def get_pp_merged_fp32_from_fp16_param_groups(optimizer,
fp32_from_fp16_groups,
param_name_groups=None):
pp_group_fp32_from_fp16_groups = smp.allgather(fp32_from_fp16_groups,
smp.PP_GROUP)
if param_name_groups is not None:
index_to_param_name_groups = []
# obtain index_to_param_name mapping across tp_group
for param_name_group in param_name_groups:
index_to_param_name = {}
for param_name, index in param_name_group.items():
index_to_param_name[index] = param_name
index_to_param_name_groups.append(index_to_param_name)
# allgather the index_to_param_name_groups across the pp_group
pp_index_to_param_name_groups = smp.allgather(
index_to_param_name_groups, smp.PP_GROUP)
else:
raise ValueError(
"Merging is not supported when param_name_groups is None")
pp_merged_fp32_from_fp16_groups = []
result_param_groups = []
# iterate through all the groups for rank 0
for group_idx in range(len(pp_group_fp32_from_fp16_groups[0])):
merged = []
start_idx = 0
result_param_group = {}
# for each group iterate through all ranks and merge the param groups across pp_ranks
for rank, group in enumerate(pp_group_fp32_from_fp16_groups):
cur_g = group[group_idx]
start_idx += len(merged)
for i, _ in enumerate(cur_g):
param_name = pp_index_to_param_name_groups[rank][group_idx][i]
if param_name in result_param_group:
raise ValueError(
"same param_name present in the param_groups of different pipeline parallel partitions"
)
result_param_group[param_name] = i + start_idx
merged.extend(cur_g)
pp_merged_fp32_from_fp16_groups.append(merged)
result_param_groups.append(result_param_group)
return pp_merged_fp32_from_fp16_groups, result_param_groups
def smp_gather(tensor):
if isinstance(tensor, (list, tuple)):
return type(tensor)(smp_gather(t) for t in tensor)
elif isinstance(tensor, dict):
return type(tensor)({k: smp_gather(v) for k, v in tensor.items()})
elif not isinstance(tensor, torch.Tensor):
raise TypeError(
f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors."
)
all_tensors = smp.allgather(tensor, smp.CommGroup.DP_GROUP)
return torch.cat([t.cpu() for t in all_tensors], dim=0)
def smp_lossgather(loss, args):
if args.use_horovod or args.use_ddp:
# Rubik: If using data parallelism, gather all losses across different model
# replicas and check if losses match.
losses = smp.allgather(loss, smp.DP_GROUP)
for l in losses:
assert math.isclose(l, losses[0])
assert loss < 0.14
else:
assert loss < 0.08
def main():
parser = get_parser()
args = parser.parse_args()
if not torch.cuda.is_available():
raise ValueError(
"The script requires CUDA support, but CUDA not available")
use_ddp = args.ddp > 0
use_horovod = args.horovod > 0
# Fix seeds in order to get the same losses across runs
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cfg = {
"microbatches": args.num_microbatches,
"placement_strategy": "spread",
"pipeline": args.pipeline,
"optimize": "speed",
"partitions": args.num_partitions,
"horovod": use_horovod,
"ddp": use_ddp,
}
smp.init(cfg)
# SM Distributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")
kwargs = {"batch_size": args.batch_size}
kwargs.update({"num_workers": 1, "pin_memory": True, "shuffle": False})
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])
# SM Distributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
if smp.local_rank() == 0:
dataset1 = datasets.MNIST("../data",
train=True,
download=True,
transform=transform)
smp.barrier()
dataset1 = datasets.MNIST("../data",
train=True,
download=False,
transform=transform)
if (use_ddp or use_horovod) and smp.dp_size() > 1:
partitions_dict = {
f"{i}": 1 / smp.dp_size()
for i in range(smp.dp_size())
}
dataset1 = SplitDataset(dataset1, partitions=partitions_dict)
dataset1.select(f"{smp.dp_rank()}")
# Download and create dataloaders for train and test dataset
dataset2 = datasets.MNIST("../data", train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
model = GroupedNet()
# SMP handles the transfer of parameters to the right device
# and the user doesn't need to call 'model.to' explicitly.
# model.to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
# SM Distributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
scaler = smp.amp.GradScaler()
optimizer = smp.DistributedOptimizer(optimizer)
if args.partial_checkpoint:
checkpoint = smp.load(args.partial_checkpoint, partial=True)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
elif args.full_checkpoint:
checkpoint = smp.load(args.full_checkpoint, partial=False)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, scaler, device, train_loader, optimizer, epoch)
test_loss = test(args, model, device, test_loader)
scheduler.step()
if args.save_partial_model:
if smp.dp_rank() == 0:
model_dict = model.local_state_dict()
opt_dict = optimizer.local_state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
f"./pt_mnist_checkpoint.pt",
partial=True,
)
if args.save_full_model:
if smp.dp_rank() == 0:
model_dict = model.state_dict()
opt_dict = optimizer.state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
"./pt_mnist_checkpoint.pt",
partial=False,
)
# Waiting the save checkpoint to be finished before run another allgather_object
smp.barrier()
if args.assert_losses:
if use_horovod or use_ddp:
# SM Distributed: If using data parallelism, gather all losses across different model
# replicas and check if losses match.
losses = smp.allgather(test_loss, smp.DP_GROUP)
for l in losses:
assert math.isclose(l, losses[0])
assert test_loss < 0.18
else:
assert test_loss < 0.08
def get_tp_merged_fp32_from_fp16_param_groups(optimizer,
cpu_fp32_from_fp16_groups):
def _merge_param_group_tp_group(group_idx, param_group):
result_fp32_from_fp16_param_group = []
param_name_group = {}
for i, param in enumerate(param_group):
# for each param, obtain param_name from param using two dicts above for tp_rank 0
param_index = param_id_to_index_tp_group[rank_0][
fp32_from_fp16_paramid_groups_tp_group[rank_0][group_idx][i]]
param_name = param_index_to_name_tp_group[rank_0][param_index]
# obtain distribution axis for the param and check if its distributed
# axis = master_distribution_axis_tp_rank_0[fp32_from_fp16_paramid_groups_tp_group[rank_0][group_idx][i]]
axis = master_distribution_axis_tp_rank_0.get(
fp32_from_fp16_paramid_groups_tp_group[rank_0][group_idx][i],
None)
if axis is not None:
tensors = []
for r in range(smp.tp_size()):
# if distributed, for each rank, obtain param id from index using above two dicts
param_index_r = param_name_to_index_tp_group[r][param_name]
param_id_r = param_index_to_id_tp_group[r][param_index_r]
# search param id in fp32_from_fp16_groups_param_ids and find the index.
group_param_idx = fp32_from_fp16_paramid_groups_tp_group[
r][group_idx].index(param_id_r)
# use the param corresponding to the index from fp32_from_fp16_groups for concatenation along axis
tensors.append(fp32_from_fp16_param_groups_tp_group[r]
[group_idx][group_param_idx])
result_fp32_from_fp16_param_group.append(
torch.cat(tensors, axis))
else:
# if not distributed set tp_rank 0 param as the param
result_fp32_from_fp16_param_group.append(param)
param_name_group[param_name] = i
return result_fp32_from_fp16_param_group, param_name_group
# get param_index_to_name all and param_name_to_index_all
param_index_to_name_tp_group = smp_state.param_index_to_name_tp_group
param_name_to_index_tp_group = smp_state.param_name_to_index_tp_group
# get mapping of param_id_to_index_all and param_index_to_id_all
param_id_to_index = optimizer._param_id_to_index()
param_id_to_index_tp_group = smp.allgather(param_id_to_index, smp.TP_GROUP)
param_index_to_id_tp_group = _get_param_index_to_id(
param_id_to_index_tp_group)
# allgather all param ids and all params for fp32_from_fp16_groups
fp32_from_fp16_paramid_groups = optimizer.fp32_from_fp16_paramid_groups
fp32_from_fp16_paramid_groups_tp_group = smp.allgather(
fp32_from_fp16_paramid_groups, smp.TP_GROUP)
fp32_from_fp16_param_groups_tp_group = smp.allgather(
cpu_fp32_from_fp16_groups, smp.TP_GROUP)
# broadcast distribution axis from tp_rank 0 to all tp_ranks
master_distribution_axis_tp_rank_0 = None
if smp.tp_rank() == 0:
master_distribution_axis_tp_rank_0 = optimizer.master_distribution_axis
smp.broadcast(master_distribution_axis_tp_rank_0, smp.TP_GROUP)
else:
master_distribution_axis_tp_rank_0 = smp.recv_from(
0, smp.RankType.TP_RANK)
result_fp32_from_fp16_param_groups = []
param_name_groups = []
rank_0 = 0
# iterate through all the params for tp_group_fp32_from_fp16_groups[rank_0]
for group_idx, param_group in enumerate(
fp32_from_fp16_param_groups_tp_group[rank_0]):
result_fp32_from_fp16_param_group, param_name_group = _merge_param_group_tp_group(
group_idx, param_group)
result_fp32_from_fp16_param_groups.append(
result_fp32_from_fp16_param_group)
param_name_groups.append(param_name_group)
return result_fp32_from_fp16_param_groups, param_name_groups
def main():
parser = get_parser()
args = parser.parse_args()
if not torch.cuda.is_available():
raise ValueError(
"The script requires CUDA support, but CUDA not available")
args.rank = -1
args.world_size = 1
if args.model_parallel:
args.deepspeed = False
cfg = {
"microbatches": args.num_microbatches,
"placement_strategy": args.placement_strategy,
"pipeline": args.pipeline,
"optimize": args.optimize,
"partitions": args.num_partitions,
"horovod": args.horovod,
"ddp": args.ddp,
}
smp.init(cfg)
torch.cuda.set_device(smp.local_rank())
args.rank = smp.dp_rank()
args.world_size = smp.size()
else:
# initialize deepspeed
print(f"args.deepspeed : {args.deepspeed}")
deepspeed_utils.init_deepspeed(args.deepspeed)
if deepspeed_utils.is_root_worker():
args.rank = 0
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed + args.rank)
np.random.seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
# args.LEARNING_RATE = args.LEARNING_RATE * float(args.world_size)
cudnn.deterministic = True
if cudnn.deterministic:
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
args.kwargs = {'num_workers': args.num_worker, 'pin_memory': True}
device = torch.device("cuda")
logger.debug(f"args.image_folder : {args.image_folder}")
logger.debug(f"args.rank : {args.rank}")
## SageMaker
try:
if os.environ.get('SM_MODEL_DIR') is not None:
args.model_dir = os.environ.get('SM_MODEL_DIR')
# args.output_dir = os.environ.get('SM_OUTPUT_DATA_DIR')
args.image_folder = os.environ.get('SM_CHANNEL_TRAINING')
except:
logger.debug("not SageMaker")
pass
IMAGE_SIZE = args.image_size
IMAGE_PATH = args.image_folder
EPOCHS = args.EPOCHS
BATCH_SIZE = args.BATCH_SIZE
LEARNING_RATE = args.LEARNING_RATE
LR_DECAY_RATE = args.LR_DECAY_RATE
NUM_TOKENS = args.NUM_TOKENS
NUM_LAYERS = args.NUM_LAYERS
NUM_RESNET_BLOCKS = args.NUM_RESNET_BLOCKS
SMOOTH_L1_LOSS = args.SMOOTH_L1_LOSS
EMB_DIM = args.EMB_DIM
HID_DIM = args.HID_DIM
KL_LOSS_WEIGHT = args.KL_LOSS_WEIGHT
STARTING_TEMP = args.STARTING_TEMP
TEMP_MIN = args.TEMP_MIN
ANNEAL_RATE = args.ANNEAL_RATE
NUM_IMAGES_SAVE = args.NUM_IMAGES_SAVE
# transform = Compose(
# [
# RandomResizedCrop(args.image_size, args.image_size),
# OneOf(
# [
# IAAAdditiveGaussianNoise(),
# GaussNoise(),
# ],
# p=0.2
# ),
# VerticalFlip(p=0.5),
# OneOf(
# [
# MotionBlur(p=.2),
# MedianBlur(blur_limit=3, p=0.1),
# Blur(blur_limit=3, p=0.1),
# ],
# p=0.2
# ),
# OneOf(
# [
# CLAHE(clip_limit=2),
# IAASharpen(),
# IAAEmboss(),
# RandomBrightnessContrast(),
# ],
# p=0.3
# ),
# HueSaturationValue(p=0.3),
# # Normalize(
# # mean=[0.485, 0.456, 0.406],
# # std=[0.229, 0.224, 0.225],
# # )
# ],
# p=1.0
# )
transform = T.Compose([
T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
T.Resize(IMAGE_SIZE),
T.CenterCrop(IMAGE_SIZE),
T.ToTensor()
])
sampler = None
dl = None
# data
logger.debug(f"IMAGE_PATH : {IMAGE_PATH}")
# ds = AlbumentationImageDataset(
# IMAGE_PATH,
# transform=transform,
# args=args
# )
ds = ImageFolder(
IMAGE_PATH,
transform=transform,
)
if args.model_parallel and (args.ddp
or args.horovod) and smp.dp_size() > 1:
partitions_dict = {
f"{i}": 1 / smp.dp_size()
for i in range(smp.dp_size())
}
ds = SplitDataset(ds, partitions=partitions_dict)
ds.select(f"{smp.dp_rank()}")
dl = DataLoader(ds,
BATCH_SIZE,
shuffle=True,
drop_last=args.model_parallel,
**args.kwargs)
vae_params = dict(image_size=IMAGE_SIZE,
num_layers=NUM_LAYERS,
num_tokens=NUM_TOKENS,
codebook_dim=EMB_DIM,
hidden_dim=HID_DIM,
num_resnet_blocks=NUM_RESNET_BLOCKS)
vae = DiscreteVAE(**vae_params,
smooth_l1_loss=SMOOTH_L1_LOSS,
kl_div_loss_weight=KL_LOSS_WEIGHT).to(device)
# optimizer
opt = Adam(vae.parameters(), lr=LEARNING_RATE)
sched = ExponentialLR(optimizer=opt, gamma=LR_DECAY_RATE)
if args.model_parallel:
import copy
dummy_codebook = copy.deepcopy(vae.codebook)
dummy_decoder = copy.deepcopy(vae.decoder)
vae = smp.DistributedModel(vae)
scaler = smp.amp.GradScaler()
opt = smp.DistributedOptimizer(opt)
if args.partial_checkpoint:
args.checkpoint = smp.load(args.partial_checkpoint, partial=True)
vae.load_state_dict(args.checkpoint["model_state_dict"])
opt.load_state_dict(args.checkpoint["optimizer_state_dict"])
elif args.full_checkpoint:
args.checkpoint = smp.load(args.full_checkpoint, partial=False)
vae.load_state_dict(args.checkpoint["model_state_dict"])
opt.load_state_dict(args.checkpoint["optimizer_state_dict"])
assert len(ds) > 0, 'folder does not contain any images'
if (not args.model_parallel) and args.rank == 0:
print(f'{len(ds)} images found for training')
# weights & biases experiment tracking
# import wandb
model_config = dict(num_tokens=NUM_TOKENS,
smooth_l1_loss=SMOOTH_L1_LOSS,
num_resnet_blocks=NUM_RESNET_BLOCKS,
kl_loss_weight=KL_LOSS_WEIGHT)
# run = wandb.init(
# project = 'dalle_train_vae',
# job_type = 'train_model',
# config = model_config
# )
def save_model(path):
if not args.rank == 0:
return
save_obj = {'hparams': vae_params, 'weights': vae.state_dict()}
torch.save(save_obj, path)
# distribute with deepspeed
if not args.model_parallel:
deepspeed_utils.check_batch_size(BATCH_SIZE)
deepspeed_config = {'train_batch_size': BATCH_SIZE}
(distr_vae, opt, dl, sched) = deepspeed_utils.maybe_distribute(
args=args,
model=vae,
optimizer=opt,
model_parameters=vae.parameters(),
training_data=ds if args.deepspeed else dl,
lr_scheduler=sched,
config_params=deepspeed_config,
)
try:
# Rubik: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(vae, images, temp):
# logger.debug(f"args.amp : {args.amp}")
with autocast(enabled=(args.amp > 0)):
loss, recons = vae(images,
return_loss=True,
return_recons=True,
temp=temp)
scaled_loss = scaler.scale(loss) if args.amp else loss
vae.backward(scaled_loss)
# torch.nn.utils.clip_grad_norm_(vae.parameters(), 5)
return loss, recons
@smp.step
def get_codes_step(vae, images, k):
images = images[:k]
logits = vae.forward(images, return_logits=True)
codebook_indices = logits.argmax(dim=1).flatten(1)
return codebook_indices
def hard_recons_step(dummy_decoder, dummy_codebook, codebook_indices):
from functools import partial
for module in dummy_codebook.modules():
method = smp_state.patch_manager.get_original_method(
"forward", type(module))
module.forward = partial(method, module)
image_embeds = dummy_codebook.forward(codebook_indices)
b, n, d = image_embeds.shape
h = w = int(sqrt(n))
image_embeds = rearrange(image_embeds,
'b (h w) d -> b d h w',
h=h,
w=w)
for module in dummy_decoder.modules():
method = smp_state.patch_manager.get_original_method(
"forward", type(module))
module.forward = partial(method, module)
hard_recons = dummy_decoder.forward(image_embeds)
return hard_recons
except:
pass
# starting temperature
global_step = 0
temp = STARTING_TEMP
for epoch in range(EPOCHS):
##
batch_time = util.AverageMeter('Time', ':6.3f')
data_time = util.AverageMeter('Data', ':6.3f')
losses = util.AverageMeter('Loss', ':.4e')
top1 = util.AverageMeter('Acc@1', ':6.2f')
top5 = util.AverageMeter('Acc@5', ':6.2f')
progress = util.ProgressMeter(
len(dl), [batch_time, data_time, losses, top1, top5],
prefix="Epoch: [{}]".format(epoch))
vae.train()
start = time.time()
for i, (images, _) in enumerate(dl):
images = images.to(device, non_blocking=True)
opt.zero_grad()
if args.model_parallel:
loss, recons = train_step(vae, images, temp)
# Rubik: Average the loss across microbatches.
loss = loss.reduce_mean()
recons = recons.reduce_mean()
else:
loss, recons = distr_vae(images,
return_loss=True,
return_recons=True,
temp=temp)
if (not args.model_parallel) and args.deepspeed:
# Gradients are automatically zeroed after the step
distr_vae.backward(loss)
distr_vae.step()
elif args.model_parallel:
if args.amp:
scaler.step(opt)
scaler.update()
else:
# some optimizers like adadelta from PT 1.8 dont like it when optimizer.step is called with no param
if len(list(vae.local_parameters())) > 0:
opt.step()
else:
loss.backward()
opt.step()
logs = {}
if i % 10 == 0:
if args.rank == 0:
# if deepspeed_utils.is_root_worker():
k = NUM_IMAGES_SAVE
with torch.no_grad():
if args.model_parallel:
model_dict = vae.state_dict()
model_dict_updated = {}
for key, val in model_dict.items():
if "decoder" in key:
key = key.replace("decoder.", "")
elif "codebook" in key:
key = key.replace("codebook.", "")
model_dict_updated[key] = val
dummy_decoder.load_state_dict(model_dict_updated,
strict=False)
dummy_codebook.load_state_dict(model_dict_updated,
strict=False)
codes = get_codes_step(vae, images, k)
codes = codes.reduce_mean().to(torch.long)
hard_recons = hard_recons_step(
dummy_decoder, dummy_codebook, codes)
else:
codes = vae.get_codebook_indices(images[:k])
hard_recons = vae.decode(codes)
images, recons = map(lambda t: t[:k], (images, recons))
images, recons, hard_recons, codes = map(
lambda t: t.detach().cpu(),
(images, recons, hard_recons, codes))
images, recons, hard_recons = map(
lambda t: make_grid(t.float(),
nrow=int(sqrt(k)),
normalize=True,
range=(-1, 1)),
(images, recons, hard_recons))
# logs = {
# **logs,
# 'sample images': wandb.Image(images, caption = 'original images'),
# 'reconstructions': wandb.Image(recons, caption = 'reconstructions'),
# 'hard reconstructions': wandb.Image(hard_recons, caption = 'hard reconstructions'),
# 'codebook_indices': wandb.Histogram(codes),
# 'temperature': temp
# }
if args.model_parallel:
filename = f'{args.model_dir}/vae.pt'
if smp.dp_rank == 0:
if args.save_full_model:
model_dict = vae.state_dict()
opt_dict = opt.state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
filename,
partial=False,
)
else:
model_dict = vae.local_state_dict()
opt_dict = opt.local_state_dict()
smp.save(
{
"model_state_dict": model_dict,
"optimizer_state_dict": opt_dict
},
filename,
partial=True,
)
smp.barrier()
else:
save_model(f'{args.model_dir}/vae.pt')
# wandb.save(f'{args.model_dir}/vae.pt')
# temperature anneal
temp = max(temp * math.exp(-ANNEAL_RATE * global_step),
TEMP_MIN)
# lr decay
sched.step()
# Collective loss, averaged
if args.model_parallel:
avg_loss = loss.detach().clone()
# print("args.world_size : {}".format(args.world_size))
avg_loss /= args.world_size
else:
avg_loss = deepspeed_utils.average_all(loss)
if args.rank == 0:
if i % 100 == 0:
lr = sched.get_last_lr()[0]
print(epoch, i, f'lr - {lr:6f}, loss - {avg_loss.item()},')
logs = {
**logs, 'epoch': epoch,
'iter': i,
'loss': avg_loss.item(),
'lr': lr
}
# wandb.log(logs)
global_step += 1
if args.rank == 0:
# Every print_freq iterations, check the loss, accuracy, and speed.
# For best performance, it doesn't make sense to print these metrics every
# iteration, since they incur an allreduce and some host<->device syncs.
# Measure accuracy
# prec1, prec5 = util.accuracy(output, target, topk=(1, 5))
# to_python_float incurs a host<->device sync
losses.update(util.to_python_float(loss), images.size(0))
# top1.update(util.to_python_float(prec1), images.size(0))
# top5.update(util.to_python_float(prec5), images.size(0))
# Waiting until finishing operations on GPU (Pytorch default: async)
torch.cuda.synchronize()
batch_time.update((time.time() - start) / args.log_interval)
end = time.time()
print(
'Epoch: [{0}][{1}/{2}] '
'Train_Time={batch_time.val:.3f}: avg-{batch_time.avg:.3f}, '
'Train_Speed={3:.3f} ({4:.3f}), '
'Train_Loss={loss.val:.10f}:({loss.avg:.4f}),'.format(
epoch,
i,
len(dl),
args.world_size * BATCH_SIZE / batch_time.val,
args.world_size * BATCH_SIZE / batch_time.avg,
batch_time=batch_time,
loss=losses))
# if deepspeed_utils.is_root_worker():
# save trained model to wandb as an artifact every epoch's end
# model_artifact = wandb.Artifact('trained-vae', type = 'model', metadata = dict(model_config))
# model_artifact.add_file(f'{args.model_dir}/vae.pt')
# run.log_artifact(model_artifact)
if args.rank == 0:
# if deepspeed_utils.is_root_worker():
# save final vae and cleanup
if args.model_parallel:
logger.debug('save model_parallel')
else:
save_model(os.path.join(args.model_dir, 'vae-final.pt'))
# wandb.save(f'{args.model_dir}/vae-final.pt')
# model_artifact = wandb.Artifact('trained-vae', type = 'model', metadata = dict(model_config))
# model_artifact.add_file(f'{args.model_dir}/vae-final.pt')
# run.log_artifact(model_artifact)
# wandb.finish()
if args.model_parallel:
if args.assert_losses:
if args.horovod or args.ddp:
# SM Distributed: If using data parallelism, gather all losses across different model
# replicas and check if losses match.
losses = smp.allgather(loss, smp.DP_GROUP)
for l in losses:
print(l)
assert math.isclose(l, losses[0])
assert loss < 0.18
else:
assert loss < 0.08
smp.barrier()
print("SMP training finished successfully")
