import time
from contextlib import suppress
import numpy as np
import torch
from tqdm import tqdm
import datetime
import os
import gc
from torch.distributed.fsdp import (
FullyShardedDataParallel as FSDP,
MixedPrecision,
BackwardPrefetch,
ShardingStrategy,
FullStateDictConfig,
StateDictType,
)
from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
from torch.distributed.fsdp.wrap import (
transformer_auto_wrap_policy,
enable_wrap,
wrap,
)
from torch.utils.tensorboard import SummaryWriter
import logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s %(message)s',
datefmt='%m/%d %I:%M:%S',
)
def get_cast_dtype(precision: str):
cast_dtype = None
if precision == "bf16":
cast_dtype = torch.bfloat16
elif precision == "fp16":
cast_dtype = torch.float16
return cast_dtype
def get_autocast(precision):
if precision == "amp_fp16":
return lambda: torch.cuda.amp.autocast(dtype=torch.float16)
elif precision == "amp_bfloat16" or precision == "amp_bf16":
# amp_bfloat16 is more stable than amp float16 for clip training
return lambda: torch.cuda.amp.autocast(dtype=torch.bfloat16)
else:
return suppress
def get_sync(model, flag):
if flag:
return suppress
else:
return lambda: model.no_sync()
def train_one_epoch(
args,
model,
laion_loader,
pile_loader,
tokenizer,
optimizer,
lr_scheduler,
device_id,
writer: SummaryWriter,
optim_groups,
scaler,
total_laion_token: int,
total_pile_token: int,
total_laion_sample: int,
total_step: int,
):
world_size = torch.distributed.get_world_size()
autocast = get_autocast(args.precision)
cast_dtype = get_cast_dtype(args.precision)
media_token_id = tokenizer("<|#image#|>", add_special_tokens=False)["input_ids"][-1]
endofmedia_token_id = tokenizer("<|#endofimage#|>", add_special_tokens=False)["input_ids"][-1]
visual_token_id = tokenizer("<|#visual#|>", add_special_tokens=False)["input_ids"][-1]
if args.add_box:
box_token_id = tokenizer("<|#box#|>", add_special_tokens=False)["input_ids"][-1]
endofobject_token_id = tokenizer("<|#endofobject#|>", add_special_tokens=False)["input_ids"][-1]
endofattr_token_id = tokenizer("<|#endofattr#|>", add_special_tokens=False)["input_ids"][-1]
if args.use_format_v2:
prebox_token_id = tokenizer("<|#prebox#|>", add_special_tokens=False)["input_ids"][-1]
previsual_token_id = tokenizer("<|#previsual#|>", add_special_tokens=False)["input_ids"][-1]
if args.rank == 0:
logging.info(f"train from: {total_step} step")
model.train()
# loop through dataloader
last_logging_step = total_step
last_save_step = total_step
for num_steps, (batch_laion, batch_pile) in tqdm(
enumerate(zip(laion_loader, pile_loader)),
disable=args.rank != 0 or "SLURM_PROCID" in os.environ,
total=args.num_steps * args.gradient_accumulation_steps,
initial=total_step * args.gradient_accumulation_steps,
):
#### LAION FORWARD PASS ####
images = (
batch_laion[0]
.to(device_id, dtype=cast_dtype, non_blocking=True)
.unsqueeze(1)
.unsqueeze(1)
)
image_nums = batch_laion[1]
image_start_index_list = batch_laion[2]
# TODO: OPT model: input_ids is not started with </s> while input_ids2 is?
input_ids = batch_laion[3].to(device_id, non_blocking=True).long()
attention_mask = batch_laion[4].to(device_id, dtype=cast_dtype, non_blocking=True)
added_bbox_list = [x.to(device_id) for x in batch_laion[5]] # list object
total_laion_token += int(attention_mask.sum().long()) * world_size
total_laion_sample += sum(image_nums) * world_size
labels = input_ids.clone()
if args.add_box:
labels[input_ids == visual_token_id] = -100
labels[input_ids == box_token_id] = -100
labels[input_ids == endofattr_token_id] = -100
if args.use_format_v2:
labels[input_ids == previsual_token_id] = -100
labels[input_ids == prebox_token_id] = -100
labels[torch.roll(input_ids == prebox_token_id, 1)] = -100
labels[torch.roll(input_ids == box_token_id, 1)] = -100
labels[:, 0] = -100
labels[input_ids == tokenizer.pad_token_id] = -100
labels[input_ids == media_token_id] = -100
labels[input_ids == endofmedia_token_id] = -100
labels.to(device_id)
current_laion_num = input_ids.shape[0]
#### PILE FORWARD PASS ####
if batch_pile is not None and batch_pile[0] is not None and batch_pile[1] is not None:
input_ids2 = batch_pile[0].to(device_id, non_blocking=True).long()
attention_mask2 = batch_pile[1].to(device_id, dtype=cast_dtype, non_blocking=True)
input_length = input_ids.shape[-1]
input_ids2 = torch.cat([input_ids2, torch.ones((input_ids2.shape[0], input_length - input_ids2.shape[1]), device=input_ids2.device, dtype=input_ids2.dtype) * tokenizer.pad_token_id], dim=-1)
attention_mask2 = torch.cat([attention_mask2, torch.zeros((attention_mask2.shape[0], input_length - attention_mask2.shape[1]), device=attention_mask2.device, dtype=attention_mask2.dtype)], dim=-1)
labels2 = input_ids2.clone()
labels2[labels2 == tokenizer.pad_token_id] = -100
labels2[:, 0] = -100
labels2.to(device_id)
if (num_steps != 0 and num_steps % args.pile_freq == 0) or args.pile_freq == 1:
image_nums = image_nums + [0] * len(input_ids2)
image_start_index_list = image_start_index_list + [[]] * len(input_ids2)
input_ids = torch.cat([input_ids, input_ids2], dim=0)
attention_mask = torch.cat([attention_mask, attention_mask2], dim=0)
labels = torch.cat([labels, labels2], dim=0)
total_pile_token += int(attention_mask2.sum().long()) * world_size
else:
del input_ids2
del attention_mask2
del labels2
if args.instruct:
answer_token_id = tokenizer(" Answer").input_ids[0]
answer_token_loc = (input_ids == answer_token_id).nonzero()
for batch_idx, idx in answer_token_loc:
labels[batch_idx][:idx+2] = -100
if args.relation and not args.instruct:
relations = batch_laion[6]
else:
relations = None
if len(added_bbox_list) == 0:
added_bbox_list = None
update_flag = (num_steps != 0 and num_steps % args.gradient_accumulation_steps == 0) or args.gradient_accumulation_steps == 1
# do_sync = get_sync(model, update_flag)
with autocast():
# modify:
# /gpfs/u/home/LMCG/LMCGljnn/scratch/miniconda3-ppc64le/envs/unified/lib/python3.9/site-packages/transformers/models/codegen/modeling_codegen.py
# /gpfs/u/home/LMCG/LMCGljnn/scratch/miniconda3-ppc64le/envs/unified/lib/python3.9/site-packages/transformers/models/opt/modeling_opt.py
# CrossEntropyLoss(reduction="none")
outputs = model(
vision_x=images,
lang_x=input_ids,
attention_mask=attention_mask,
labels=labels,
image_nums=image_nums,
image_start_index_list=image_start_index_list,
added_bbox_list=added_bbox_list,
add_box=args.add_box,
relations=relations,
)
loss_total = outputs.loss.reshape(labels.shape[0], -1)
loss_sample = loss_total.sum(-1) / (loss_total != 0).sum(-1)
loss_sample_for_laion = loss_sample[:current_laion_num]
nan_mask = torch.isnan(loss_sample_for_laion)
if nan_mask.sum() > 0:
logging.warning(f"caption NaN: {nan_mask}")
if nan_mask.sum() == len(loss_sample_for_laion) or not model.valid:
logging.info("WARNING: skip this caption loss due to some error")
loss_laion = torch.tensor(0.0).cuda()
else:
loss_laion = loss_sample_for_laion[~nan_mask].mean()
loss_caption = loss_laion
divided_loss_laion = loss_laion / args.gradient_accumulation_steps
if current_laion_num != loss_sample.shape[0]:
loss_pile = loss_sample[current_laion_num:].mean()
else:
loss_pile = torch.tensor(0.0).cuda()
divided_loss_pile = loss_pile / args.gradient_accumulation_steps
if "detection_losses" in outputs:
loss_det = outputs["detection_losses"]["loss"]
loss_iou = outputs["detection_losses"]["loss_iou"]
loss_obj = outputs["detection_losses"]["loss_obj"]
loss_cls = outputs["detection_losses"]["loss_cls"]
else:
loss_det = torch.tensor(0.0).cuda()
loss_iou = torch.tensor(0.0).cuda()
loss_obj = torch.tensor(0.0).cuda()
loss_cls = torch.tensor(0.0).cuda()
if "loss_dict" in outputs:
visual_loss_iou = outputs["loss_dict"][0]["loss_iou"]
previsual_loss_iou = outputs["loss_dict"][1]["loss_iou"]
visual_loss_obj = outputs["loss_dict"][0]["loss_obj"]
previsual_loss_obj = outputs["loss_dict"][1]["loss_obj"]
else:
visual_loss_iou = torch.tensor(0.0).cuda()
previsual_loss_iou = torch.tensor(0.0).cuda()
visual_loss_obj = torch.tensor(0.0).cuda()
previsual_loss_obj = torch.tensor(0.0).cuda()
divided_loss_det = loss_det / args.gradient_accumulation_steps
loss_rel = outputs.get("rel_loss", torch.tensor(0.0).cuda())
divided_loss_rel = loss_rel / args.gradient_accumulation_steps
loss = (
divided_loss_laion * args.loss_multiplier_laion +
divided_loss_pile * args.loss_multiplier_pile +
divided_loss_det * args.loss_multiplier_det +
divided_loss_rel * args.loss_multiplier_rel
)
scaler.scale(loss).backward()
# for logging only
loss = (
loss_laion * args.loss_multiplier_laion
+ loss_pile * args.loss_multiplier_pile
+ loss_det * args.loss_multiplier_det
+ loss_rel * args.loss_multiplier_rel
).detach()
# step optimizer and log
if update_flag:
#### MASK GRADIENTS FOR EMBEDDINGS ####
# Note (anas): Do not apply weight decay to embeddings as it will break this function.
# ! not an important point
# if args.ddp:
# def mask_embedding(m):
# if isinstance(m, torch.nn.Embedding) and m.weight.requires_grad:
# zero_mask = torch.zeros_like(m.weight.grad)
# zero_mask[media_token_id] = torch.ones_like(zero_mask[media_token_id])
# zero_mask[endofmedia_token_id] = torch.ones_like(zero_mask[endofmedia_token_id])
# m.weight.grad = m.weight.grad * zero_mask
# model.apply(mask_embedding)
total_step += 1
scaler.unscale_(optimizer)
if args.ddp:
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
else:
model.clip_grad_norm_(1.0)
scaler.step(optimizer)
scaler.update()
lr_scheduler.step()
optimizer.zero_grad()
# https://github.com/facebookresearch/fairscale/issues/627
model.zero_grad(set_to_none=True)
if args.rank == 0 and total_step % args.logging_steps == 0 and total_step != last_logging_step:
last_logging_step = total_step
global_step = total_step
lr = optimizer.param_groups[0]["lr"]
writer.add_scalar("lr", lr, global_step)
writer.add_scalar("scale", scaler.get_scale(), global_step)
writer.add_scalar("loss_groundcaption", loss_laion.item(), global_step)
writer.add_scalar("loss_laion", loss_caption.item(), global_step)
writer.add_scalar("loss_pile", loss_pile.item(), global_step)
writer.add_scalar("loss", loss.item(), global_step)
writer.add_scalar("loss_det", loss_det.item(), global_step)
writer.add_scalar("loss_iou", loss_iou.item(), global_step)
writer.add_scalar("loss_obj", loss_obj.item(), global_step)
writer.add_scalar("loss_cls", loss_cls.item(), global_step)
if loss_rel.item() != 0:
writer.add_scalar("loss_rel", loss_rel.item(), global_step)
if args.use_format_v2:
writer.add_scalar("loss_iou_visual", visual_loss_iou.item(), global_step)
writer.add_scalar("loss_obj_visual", visual_loss_obj.item(), global_step)
writer.add_scalar("loss_iou_previsual", previsual_loss_iou.item(), global_step)
writer.add_scalar("loss_obj_previsual", previsual_loss_obj.item(), global_step)
global_sample_num = total_laion_sample
writer.add_scalar("loss_groundcaption_vs_sample_num", loss_laion.item(), global_sample_num)
writer.add_scalar("loss_laion_vs_sample_num", loss_caption.item(), global_sample_num)
writer.add_scalar("loss_pile_vs_sample_num", loss_pile.item(), global_sample_num)
writer.add_scalar("loss_vs_sample_num", loss.item(), global_sample_num)
writer.add_scalar("loss_det_vs_sample_num", loss_det.item(), global_sample_num)
writer.add_scalar("loss_iou_vs_sample_num", loss_iou.item(), global_sample_num)
writer.add_scalar("loss_obj_vs_sample_num", loss_obj.item(), global_sample_num)
if loss_rel.item() != 0:
writer.add_scalar("loss_rel_vs_sample_num", loss_rel.item(), global_sample_num)
writer.add_scalar("lr_vs_sample_num", optimizer.param_groups[0]["lr"], global_sample_num)
writer.add_scalar("loss_groundcaption_vs_token", loss_laion.item(), total_laion_token)
writer.add_scalar("loss_laion_vs_token", loss_caption.item(), total_laion_token)
writer.add_scalar("loss_pile_vs_token", loss_pile.item(), total_pile_token)
writer.add_scalar("loss_det_vs_token", loss_det.item(), total_laion_token)
writer.add_scalar("loss_iou_vs_token", loss_iou.item(), total_laion_token)
writer.add_scalar("loss_obj_vs_token", loss_obj.item(), total_laion_token)
writer.add_scalar("loss_cls_vs_token", loss_cls.item(), total_laion_token)
if loss_rel.item() != 0:
writer.add_scalar("loss_rel_vs_token", loss_rel.item(), total_laion_token)
total_token = total_laion_token + total_pile_token
writer.add_scalar("sample_num", global_sample_num, global_step)
writer.add_scalar("total_laion_token", total_laion_token, global_step)
writer.add_scalar("total_pile_token", total_pile_token, global_step)
writer.add_scalar("total_token", total_token, global_step)
logging.info(
f"[{global_step}][{total_laion_sample}][{total_token}]. total: {loss.item():.3f} // laion: {loss_caption.item():.3f} // pile: {loss_pile.item():.3f} // iou: {loss_iou.item():.4f} // obj: {loss_obj.item():.4f} // previsual_obj: {previsual_loss_obj.item():.4f} // visual_obj: {visual_loss_obj.item():.4f} // previsual_iou: {previsual_loss_iou.item():.4f} // visual_iou: {visual_loss_iou.item():.4f} // lr: {lr:.2e} // scale: {scaler.get_scale()}"
)
if total_step % args.save_interval == 0 and total_step != last_save_step:
last_save_step = total_step
torch.distributed.barrier()
if args.ddp:
cpu_state = model.state_dict()
# if args.rank == 0:
# optimizer_state = optimizer.state_dict()
else:
save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(
model, StateDictType.FULL_STATE_DICT, save_policy
):
cpu_state = model.state_dict()
torch.distributed.barrier()
# https://pytorch.org/docs/1.12/fsdp.html
# need to pass optim_groups as optim_input
# optimizer_state = FSDP.full_optim_state_dict(model, optimizer, optim_input=optim_groups)
if args.rank == 0:
checkpoint_dict = {
"model_state_dict": cpu_state,
# "optimizer_state_dict": optimizer_state,
"lr_scheduler_state_dict": lr_scheduler.state_dict(),
"scaler_state_dict": scaler.state_dict(),
"total_pile_token": total_pile_token,
"total_laion_token": total_laion_token,
"total_laion_sample": total_laion_sample,
"total_step": total_step,
}
logging.info(f"Saving checkpoint to {args.run_name}/checkpoint_{total_step}.pt")
torch.save(checkpoint_dict, f"{args.run_name}/checkpoint_{total_step}.pt")
del checkpoint_dict
if args.delete_previous_checkpoint and total_step-args.save_interval > 0 and (total_step-args.save_interval) % args.skip_delete_pattern != 0:
try:
os.remove(f"{args.run_name}/checkpoint_{total_step-args.save_interval}.pt")
except:
pass
torch.distributed.barrier()
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count