import gc import json import math import os import os.path as osp import random import sys import time from functools import partial from typing import List, Optional, Tuple os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1' import numpy as np import torch torch._dynamo.config.cache_size_limit = 64 from torch.nn import functional as F from torch.profiler import record_function from torch.utils.data import DataLoader from transformers import T5EncoderModel, T5TokenizerFast import torch.distributed as tdist import grn.utils_t2iv.dist as dist from grn.dataset.build import build_joint_dataset from grn.models.ema import get_ema_model from grn.utils_t2iv import arg_util, misc from grn.utils import wandb_utils from grn.trainer import get_trainer def build_everything_from_args(args: arg_util.Args, saver): args.set_initial_seed(benchmark=True) print(f'Loading T5 from {args.t5_path}...') from grn.models.umt5.t5 import T5EncoderModel text_encoder = T5EncoderModel( text_len=args.tlen, # 512 dtype=torch.bfloat16, # torch.bfloat16 device=args.device, checkpoint_path=osp.join(args.t5_path, 'models_t5_umt5-xxl-enc-bf16.pth'), tokenizer_path=osp.join(args.t5_path, 'umt5-xxl'), enable_fsdp=True) # False # text_encoder.model.to(args.device) text_tokenizer = text_encoder.tokenizer args.text_tokenizer_type = 'umt5' args.text_tokenizer = text_tokenizer # build models. Note that here gpt is the causal VAR transformer which performs next scale prediciton with text guidance vae_local, gpt_uncompiled, gpt_wo_ddp, gpt_ddp, gpt_wo_ddp_ema, gpt_ddp_ema, gpt_optim = build_model_optimizer(args) Trainer = get_trainer(args) # build trainer trainer = Trainer( is_visualizer=dist.is_visualizer(), device=args.device, vae_local=vae_local, gpt_wo_ddp=gpt_wo_ddp, gpt=gpt_ddp, zero=args.zero, vae_latent_dim=args.vae_latent_dim, gpt_opt=gpt_optim, reweight_loss_by_scale=args.reweight_loss_by_scale, gpt_wo_ddp_ema=gpt_wo_ddp_ema, gpt_ema=gpt_ddp_ema, use_fsdp_model_ema=args.use_fsdp_model_ema, other_args=args, ) # auto resume from broken experiment global_it = 0 if args.checkpoint_type == 'torch': from grn.utils_t2iv.save_and_load import auto_resume auto_resume_info, start_ep, start_it, acc_str, eval_milestone, trainer_state, args_state = auto_resume(args, 'ar-ckpt*.pth') print(f'initial args:\n{str(args)}') if start_ep == args.ep: print(f'[vgpt] AR finished ({acc_str}), skipping ...\n\n') return None if trainer_state is not None and len(trainer_state): trainer.load_state_dict(trainer_state, strict=False, skip_vae=True) # don't load vae again del vae_local, gpt_uncompiled, gpt_wo_ddp, gpt_ddp, gpt_wo_ddp_ema, gpt_ddp_ema, gpt_optim dist.barrier() return text_tokenizer, text_encoder, trainer, global_it def build_model_optimizer(args): from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from grn.models.grn import MultipleLayers from grn.models.init_param import init_weights from grn.utils_t2iv.lr_control import filter_params from grn.utils_t2iv.load import build_vae_gpt # disable builtin initialization for speed setattr(torch.nn.Linear, 'reset_parameters', lambda self: None) setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None) vae_local, gpt_wo_ddp = build_vae_gpt(args, device=args.model_init_device) count_p = lambda m: sum(p.numel() for p in m.parameters()) / 1e6 num_para = count_p(gpt_wo_ddp) if num_para/1000 < 20: # < 20B gpt_wo_ddp = gpt_wo_ddp.to('cuda') init_weights(gpt_wo_ddp) gpt_wo_ddp.special_init() if args.use_fsdp_model_ema: gpt_wo_ddp_ema = get_ema_model(gpt_wo_ddp) else: gpt_wo_ddp_ema = None if args.rush_resume: print(f"{args.rush_resume=}") if '.pth' in args.rush_resume: cpu_d = torch.load(args.rush_resume, 'cpu') else: from grn.utils_t2iv.save_and_load import merge_ckpt cpu_d = merge_ckpt(args.rush_resume, osp.join(args.rush_resume, 'ouput'), save=False, use_ema_model=False, fsdp_save_flatten_model=args.fsdp_save_flatten_model) if 'trainer' in cpu_d: state_dict = cpu_d['trainer']['gpt_fsdp'] ema_state_dict = cpu_d['trainer'].get('gpt_ema_fsdp', state_dict) else: state_dict = cpu_d ema_state_dict = state_dict def drop_unfit_weights(state_dict): try: if 'word_embed.weight' in state_dict and (state_dict['word_embed.weight'].shape[1] != gpt_wo_ddp.word_embed.in_features): print(f'[rush_resume] drop word_embed.weight') del state_dict['word_embed.weight'] if 'head.proj.weight' in state_dict and (state_dict['head.proj.weight'].shape[0] != gpt_wo_ddp.head.proj.out_features): print(f'[rush_resume] drop head') del state_dict['head.proj.weight'] del state_dict['head.proj.bias'] except Exception as e: print(e) for key in ['word_embed.weight', 'head.proj.weight', 'head.proj.bias']: if key in state_dict: del state_dict[key] print(f'[rush_resume] drop {key}') if 'text_proj_for_sos.ca.mat_kv.weight' in state_dict and \ (state_dict['text_proj_for_sos.ca.mat_kv.weight'].shape != gpt_wo_ddp.text_proj_for_sos.ca.mat_kv.weight.shape): print(f'[rush_resume] drop cfg_uncond') del state_dict['cfg_uncond'] for key in list(state_dict.keys()): if 'text' in key: del state_dict[key] return state_dict print(gpt_wo_ddp.load_state_dict(drop_unfit_weights(state_dict), strict=False)) if args.use_fsdp_model_ema: gpt_wo_ddp_ema.load_state_dict(drop_unfit_weights(ema_state_dict), strict=False) ndim_dict = {name: para.ndim for name, para in gpt_wo_ddp.named_parameters() if para.requires_grad} print(f'[PT] GPT model = {gpt_wo_ddp}\n\n') print(f'[PT] GPT model details:') for name, param in gpt_wo_ddp.named_parameters(): print(f"Name: {name}, Shape: {param.shape}") print(f'[PT][#para], GPT={num_para:.2f}M parameters\n\n') gpt_uncompiled = gpt_wo_ddp gpt_wo_ddp = args.compile_model(gpt_wo_ddp, args.tfast) gpt_ddp_ema = None if args.zero: from torch.distributed.fsdp import ShardingStrategy from torch.distributed.fsdp.wrap import ModuleWrapPolicy from torch.distributed.device_mesh import init_device_mesh # use mix prec: https://github.com/pytorch/pytorch/issues/76607 if args.fsdp_warp_mode == 'full': print(f'warp all modules for fsdp') def my_policy( module: torch.nn.Module, recurse: bool, **kwargs, ) -> bool: return True auto_wrap_policy = my_policy else: print(f'warp transformer blocks for fsdp') auto_wrap_policy = ModuleWrapPolicy([MultipleLayers, ]) if args.enable_hybrid_shard == 1: sharding_strategy = ShardingStrategy.HYBRID_SHARD if args.zero == 3 else ShardingStrategy._HYBRID_SHARD_ZERO2 world_size = dist.get_world_size() assert world_size % args.inner_shard_degree == 0 assert args.inner_shard_degree > 1 and args.inner_shard_degree <= world_size device_mesh = init_device_mesh('cuda', (world_size // args.inner_shard_degree, args.inner_shard_degree)) elif args.enable_hybrid_shard == -1: # no shard sharding_strategy = ShardingStrategy.NO_SHARD device_mesh = None else: sharding_strategy = ShardingStrategy.FULL_SHARD if args.zero == 3 else ShardingStrategy.SHARD_GRAD_OP device_mesh = None print(f'{">" * 45 + " " * 5} FSDP INIT with {args.zero=} {sharding_strategy=} {auto_wrap_policy=} {" " * 5 + "<" * 45}', flush=True) if args.fsdp_init_device == 'cpu': gpt_wo_ddp = gpt_wo_ddp.cpu() gpt_ddp: FSDP = FSDP( gpt_wo_ddp, device_id=dist.get_local_rank(), sharding_strategy=sharding_strategy, mixed_precision=None, auto_wrap_policy=auto_wrap_policy, use_orig_params=True, sync_module_states=True, limit_all_gathers=True, device_mesh=device_mesh, ).to(args.device) if args.use_fsdp_model_ema: gpt_wo_ddp_ema = gpt_wo_ddp_ema.to(args.device) gpt_ddp_ema: FSDP = FSDP( gpt_wo_ddp_ema, device_id=dist.get_local_rank(), sharding_strategy=sharding_strategy, mixed_precision=None, auto_wrap_policy=auto_wrap_policy, use_orig_params=True, sync_module_states=True, limit_all_gathers=True, device_mesh=device_mesh, ) else: ddp_class = DDP if dist.initialized() else misc.NullDDP gpt_ddp: DDP = ddp_class(gpt_wo_ddp, device_ids=[dist.get_local_rank()], find_unused_parameters=args.dbg, broadcast_buffers=False) torch.cuda.synchronize() # =============== build optimizer =============== nowd_keys = set() nowd_keys |= { 'cls_token', 'start_token', 'task_token', 'cfg_uncond', 'pos_embed', 'pos_1LC', 'pos_start', 'start_pos', 'lvl_embed', 'gamma', 'beta', 'ada_gss', 'moe_bias', 'scale_mul', 'text_proj_for_sos.ca.mat_q', 'scale_tokens', 'class_tokens' } names, paras, para_groups = filter_params(gpt_ddp if args.zero else gpt_wo_ddp, ndim_dict, nowd_keys=nowd_keys) del ndim_dict opt_clz = partial(torch.optim.AdamW, betas=(0.9, 0.999), fused=True) opt_kw = dict(lr=args.tlr, weight_decay=args.twd) print(f'[vgpt] optim={opt_clz}, opt_kw={opt_kw}\n') gpt_optim = opt_clz(params=para_groups, **opt_kw) del names, paras, para_groups return vae_local, gpt_uncompiled, gpt_wo_ddp, gpt_ddp, gpt_wo_ddp_ema, gpt_ddp_ema, gpt_optim def build_dataset(args): train_dataset = build_joint_dataset( args, args.meta_folders, args.meta_folder_repeats, max_caption_len=args.tlen, short_prob=args.short_cap_prob, load_vae_instead_of_image=False ) return train_dataset def main_train(args: arg_util.Args): if args.checkpoint_type == 'torch': from grn.utils_t2iv.save_and_load import CKPTSaver, auto_resume saver = CKPTSaver(dist.is_master(), eval_milestone=None) else: raise ValueError(f'{args.checkpoint_type=}') text_tokenizer, text_encoder, trainer, start_global_it = build_everything_from_args(args, saver) gc.collect(), torch.cuda.empty_cache() logging_params_milestone: List[int] = np.linspace(1, args.ep, 10+1, dtype=int).tolist() time.sleep(3), gc.collect(), torch.cuda.empty_cache(), time.sleep(3) # ============================================= epoch loop begins ============================================= # build wandb logger if dist.is_master(): wandb_utils.wandb.init(project=args.project_name, name=args.exp_name, config={}) dataloader_generator = torch.Generator() dataloader_generator.manual_seed(args.seed + tdist.get_rank()) for ep in range(args.ep): # build data at each epoch to ensure read meta take effects for each dataloader worker args.epoch = ep train_dataset = build_dataset(args) iters_train = len(train_dataset) print(f'[PT info] from {start_global_it=} {iters_train=}=======> bed: {args.bed} <=======\n') # build dataloader train_dataloader = DataLoader(dataset=train_dataset, num_workers=args.workers, pin_memory=True, batch_size=None, shuffle=True, generator=dataloader_generator) train_dataloader_iter_obj = iter(train_dataloader) # [train one epoch] stats, (sec, remain_time, finish_time) = train_one_ep( ep=ep, start_global_it=start_global_it, me=None, saver=saver, args=args, ld_or_itrt=train_dataloader_iter_obj, iters_train=iters_train, text_tokenizer=text_tokenizer, text_encoder=text_encoder, trainer=trainer, logging_params_milestone=logging_params_milestone, ) start_global_it += iters_train del stats, train_dataset, train_dataloader time.sleep(10), gc.collect(), time.sleep(10) # torch.cuda.empty_cache() return def train_one_ep( ep: int, start_global_it: int, me: misc.MetricLogger, saver, args: arg_util.Args, ld_or_itrt, iters_train: int, text_tokenizer: T5TokenizerFast, text_encoder: T5EncoderModel, trainer, logging_params_milestone, ): # IMPORTANT: import heavy packages after the Dataloader object creation/iteration to avoid OOM step_cnt = 0 header = f'[Ep]: [{ep:4d}/{args.ep}]' g_it, max_it = start_global_it, args.ep * iters_train me = misc.MetricLogger() [me.add_meter(x, misc.SmoothedValue(window_size=1, fmt='{value:.2g}')) for x in ['tlr']] [me.add_meter(x, misc.SmoothedValue(window_size=1, fmt='{median:.2f} ({global_avg:.2f})')) for x in ['tnm']] [me.add_meter(x, misc.SmoothedValue(window_size=1, fmt='{median:.3f} ({global_avg:.3f})')) for x in ['L', 'L_i', 'L_v']] [me.add_meter(x, misc.SmoothedValue(window_size=1, fmt='{median:.2f} ({global_avg:.2f})')) for x in ['Acc', 'Acc_i', 'Acc_v']] [me.add_meter(x, misc.SmoothedValue(window_size=1, fmt='{median:.2f} ({global_avg:.2f})')) for x in ['seq_usage']] # ============================================= iteration loop begins ============================================= start_it = 0 for it, data in me.log_every(start_it, iters_train, ld_or_itrt, args.log_freq, args.log_every_iter, header): # for dpo training, we will save the first iter model for comparison g_it += 1 if (g_it > 0 and g_it % args.save_model_iters_freq == 0) or (args.save_start_model and g_it == 1): if args.checkpoint_type == 'torch': saver.sav(args=args, g_it=g_it, next_ep=ep, next_it=it+1, trainer=trainer, acc_str=f'[todo]', eval_milestone=None, also_save_to=None, best_save_to=None) # [get data] images, captions, raw_features_bcthw, feature_cache_files4images, media, meta_list = data['images'], data['captions'], data['raw_features_bcthw'], data['feature_cache_files4images'], data['media'], data['meta_list'] # # [prepare text features] if args.add_class_token > 0: # c2i task text_cond_tuple = [[] for _ in range(5)] else: caption_nums = [len(item) for item in captions] flatten_captions = [] for item in captions: flatten_captions.extend(item) if args.text_tokenizer_type == 'flan_t5': tokens = text_tokenizer(text=flatten_captions, max_length=text_tokenizer.model_max_length, padding='max_length', truncation=True, return_tensors='pt') # todo: put this into dataset input_ids = tokens.input_ids.cuda(non_blocking=True) mask = tokens.attention_mask.cuda(non_blocking=True) text_features = text_encoder(input_ids=input_ids, attention_mask=mask)['last_hidden_state'].float() lens: List[int] = mask.sum(dim=-1).tolist() cu_seqlens_k = F.pad(mask.sum(dim=-1).to(dtype=torch.int32).cumsum_(0), (1, 0)) Ltext = max(lens) kv_compact = [] for text_ind, (len_i, feat_i) in enumerate(zip(lens, text_features.unbind(0))): kv_compact.append(feat_i[:len_i]) kv_compact = torch.cat(kv_compact, dim=0) text_cond_tuple: Tuple[torch.FloatTensor, List[int], torch.LongTensor, int] = (kv_compact, lens, cu_seqlens_k, Ltext, caption_nums) else: text_features = text_encoder(flatten_captions, args.device) lens = [len(item) for item in text_features] cu_seqlens_k = [0] for len_i in lens: cu_seqlens_k.append(cu_seqlens_k[-1] + len_i) cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32) Ltext = max(lens) kv_compact = torch.cat(text_features, dim=0).float() text_cond_tuple = (kv_compact, lens, cu_seqlens_k, Ltext, caption_nums) if len(images): images = [item.to(args.device, non_blocking=True) for item in images] if len(raw_features_bcthw): raw_features_bcthw = [item.to(args.device, non_blocking=True) for item in raw_features_bcthw] # [schedule learning rate and weight decay] if ep == 0 and (g_it-start_global_it) < args.wp_it: cur_lr_ratio = args.wp0 + (1-args.wp0) * (g_it-start_global_it) / args.wp_it else: cur_lr_ratio = 1 cur_lr = args.tlr * cur_lr_ratio if cur_lr_ratio < 1: cur_wd = args.twd for param_group in trainer.gpt_opt.param_groups: param_group['lr'] = cur_lr * param_group.get('lr_sc', 1) # 'lr_sc' could be assigned param_group['weight_decay'] = cur_wd * param_group.get('wd_sc', 1) # [get scheduled hyperparameters] stepping = (g_it + 1) % args.gradient_accumulation == 0 step_cnt += int(stepping) trainer.train_step( ep=ep, it=it, g_it=g_it, stepping=stepping, clip_decay_ratio=1, metric_lg=me, logging_params=stepping and step_cnt == 1 and (ep < 4 or ep in logging_params_milestone), inp_B3HW=images, raw_features_bcthw=raw_features_bcthw, feature_cache_files4images=feature_cache_files4images, text_cond_tuple=text_cond_tuple, media=media, meta_list=meta_list, args=args, ) me.update(tlr=cur_lr) # ============================================= iteration loop ends ============================================= me.synchronize_between_processes() return {k: meter.global_avg for k, meter in me.meters.items()}, me.iter_time.time_preds(max_it - (g_it + 1) + (args.ep - ep) * 15) # +15: other cost def main(): args: arg_util.Args = arg_util.init_dist_and_get_args() main_train(args) print(f'final args:\n\n{str(args)}') if isinstance(sys.stdout, dist.BackupStreamToFile) and isinstance(sys.stderr, dist.BackupStreamToFile): sys.stdout.close(), sys.stderr.close() dist.barrier() time.sleep(120) if __name__ == '__main__': main()