# Copyright (c) 2025 FoundationVision # SPDX-License-Identifier: MIT import gc import json import math import os import os.path as osp import random import sys import time import traceback from collections import deque from contextlib import nullcontext from functools import partial from distutils.util import strtobool from typing import List, Optional, Tuple os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1' # os.environ["TORCH_LOGS"] = "+dynamo" # os.environ["TORCHDYNAMO_VERBOSE"] = '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 AutoTokenizer, T5EncoderModel, T5TokenizerFast import torch.distributed as tdist import infinity.utils.dist as dist from infinity.dataset.build import build_joint_dataset from infinity.utils.save_and_load import CKPTSaver, omnistoreCheckpoint, auto_resume, omnistore_auto_resume from infinity.models.ema import get_ema_model from infinity.utils import arg_util, misc, wandb_utils from infinity.trainer import get_trainer # from infinity.utils.mfu.mfu import mfutool def build_everything_from_args(args: arg_util.Args, saver): # set seed args.set_initial_seed(benchmark=True) # build tokenizer print(f'Loading T5 from {args.t5_path}...') if 'flan-t5' in args.t5_path: from transformers import T5EncoderModel, T5TokenizerFast text_tokenizer: T5TokenizerFast = AutoTokenizer.from_pretrained(args.t5_path, revision=None, legacy=True) # text_tokenizer.model_max_length is 512 text_tokenizer.model_max_length = args.tlen text_encoder: T5EncoderModel = T5EncoderModel.from_pretrained(args.t5_path, torch_dtype=torch.float16) text_encoder.to(args.device) text_encoder.eval() text_encoder.requires_grad_(False) args.text_tokenizer_type = 'flan_t5' args.text_tokenizer = text_tokenizer else: # umt5 raise ValueError("Only flan-t5 is supported now.") # 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) # IMPORTANT: import heavy package `InfinityTrainer` after the Dataloader object creation/iteration to avoid OOM InfinityTrainer = get_trainer(args) # build trainer trainer = InfinityTrainer( device=args.device, raw_scale_schedule=args.scale_schedule, vae_local=vae_local, gpt_wo_ddp=gpt_wo_ddp, gpt=gpt_ddp, gpt_opt=gpt_optim, label_smooth=args.label_smooth, zero=args.zero, vae_type=args.vae_type, 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': auto_resume_info, start_ep, global_it, acc_str, _, trainer_state, _ = auto_resume(args, 'global_step_*') if trainer_state is not None and len(trainer_state): trainer.load_state_dict(trainer_state, strict=False, skip_vae=True) elif args.checkpoint_type == 'omnistore': resume_path, info = omnistore_auto_resume(args, 'global_step_*') if not resume_path and args.rush_omnistore_resume: resume_path = args.rush_omnistore_resume if resume_path: print(f"omnistore resume from {resume_path}", flush=True) args_state, start_ep, start_it, global_it, acc_str, eval_milestone = saver.load(resume_path, fsdp_object=trainer.gpt, optimizer_object=trainer.gpt_opt.optimizer) dist.barrier() if args.rush_omnistore_resume == resume_path: global_it = 0 auto_resume_info, acc_str, eval_milestone, trainer_state, args_state = info, '[no acc str]', [], {}, {} 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 infinity.models.infinity import Infinity, MultipleLayers from infinity.models.init_param import init_weights from infinity.utils.amp_opt import AmpOptimizer from infinity.utils.lr_control import filter_params from infinity.utils.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') if args.tini < 0: args.tini = math.sqrt(1 / gpt_wo_ddp.C / 3) init_weights(gpt_wo_ddp, other_std=args.tini) 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=}") cpu_d = torch.load(args.rush_resume, 'cpu') 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): 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.weight' in state_dict and (state_dict['head.weight'].shape[0] != gpt_wo_ddp.head.out_features): print(f'[rush_resume] drop head.weight') del state_dict['head.weight'] if 'head.bias' in state_dict and (state_dict['head.bias'].shape[0] != gpt_wo_ddp.head.bias.shape[0]): print(f'[rush_resume] drop head.bias') del state_dict['head.bias'] 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] if 'semantic_head.weight' in state_dict: print(f'[rush_resume] replace semantic_head with semantic_head2') state_dict['semantic_head2.weight'] = state_dict['semantic_head.weight'] state_dict['semantic_head2.bias'] = state_dict['semantic_head.bias'] del state_dict['semantic_head.weight'] del state_dict['semantic_head.bias'] if 'semantic_head2.weight' in state_dict and (state_dict['semantic_head2.weight'].shape[0] != gpt_wo_ddp.semantic_head2.out_features): print(f'[rush_resume] drop semantic_head2.weight, semantic_head2.bias') del state_dict['semantic_head2.weight'] del state_dict['semantic_head2.bias'] 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) elif args.torchshard_resume: from transformers.modeling_utils import load_sharded_checkpoint load_sharded_checkpoint(gpt_wo_ddp, args.torchshard_resume, 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][#para], GPT={num_para:.2f}\n\n') gpt_uncompiled = gpt_wo_ddp 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 gpt_wo_ddp.num_block_chunks == 1: # no chunks auto_wrap_policy = ModuleWrapPolicy([type(gpt_wo_ddp.unregistered_blocks[0]), ]) else: auto_wrap_policy = ModuleWrapPolicy([MultipleLayers, ]) if args.enable_hybrid_shard: 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)) 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=args.fsdp_orig, sync_module_states=True, limit_all_gathers=True, ) 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=False, broadcast_buffers=False) torch.cuda.synchronize() # =============== build optimizer =============== nowd_keys = set() if args.disable_weight_decay: 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', } names, paras, para_groups = filter_params(gpt_ddp if args.zero else gpt_wo_ddp, ndim_dict, nowd_keys=nowd_keys) del ndim_dict if '_' in args.ada: beta0, beta1 = map(float, args.ada.split('_')) else: beta0, beta1 = float(args.ada), -1 opt_clz = { 'sgd': partial(torch.optim.SGD, momentum=beta0, nesterov=True), 'adam': partial(torch.optim.AdamW, betas=(beta0, beta1), fused=args.fused_adam), 'adamw': partial(torch.optim.AdamW, betas=(beta0, beta1), fused=args.fused_adam), }[args.opt] opt_kw = dict(lr=args.tlr, weight_decay=0) if args.adam_eps: opt_kw['eps'] = args.adam_eps print(f'[vgpt] optim={opt_clz}, opt_kw={opt_kw}\n') gpt_optim = AmpOptimizer('gpt', args.fp16, opt_clz(params=para_groups, **opt_kw), gpt_ddp if args.zero else gpt_wo_ddp, args.r_accu, args.grad_clip, args.zero) 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.data_path, args.video_data_path, 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': saver = CKPTSaver(dist.is_master(), eval_milestone=None) elif args.checkpoint_type == 'omnistore': saver = omnistoreCheckpoint(eval_milestone=None) else: raise ValueError(f'{args.checkpoint_type=}') ret = build_everything_from_args(args, saver) if ret is None: return text_tokenizer, text_encoder, trainer, start_global_it = ret gc.collect(), torch.cuda.empty_cache() seg5 = np.linspace(1, args.epoch, 5+1, dtype=int).tolist() time.sleep(3), gc.collect(), torch.cuda.empty_cache(), time.sleep(3) ep_lg = max(1, args.epoch // 10) if args.epoch <= 100 else max(1, args.epoch // 20) # ============================================= epoch loop begins ============================================= # build wandb logger if dist.is_master(): wandb_utils.wandb.init(project=args.project_name, name=args.exp_name, config={}) for ep in range(args.epoch): # build data at each epoch to ensure read meta take effects for each dataloader worker args.epoch = ep if ep == 0: train_dataset = build_dataset(args) iters_train = len(train_dataset) start_ep = start_global_it // iters_train start_it = start_global_it % iters_train print(f'[PT info] from ep{start_ep} it{start_it} {iters_train=}=======> bed: {args.bed} <=======\n') if ep < start_ep: continue if ep > start_ep: train_dataset = build_dataset(args) iters_train = len(train_dataset) # [train one epoch] train_dataloader = DataLoader(dataset=train_dataset, num_workers=args.workers, pin_memory=True, batch_size=None) stats = train_one_epoch( epoch=ep, is_first_ep=ep == start_ep, start_it=start_it if ep == start_ep else 0, start_global_it=start_global_it, me=None, saver=saver, args=args, dataloader_iter=iter(train_dataloader), iters_train=iters_train, text_tokenizer=text_tokenizer, text_encoder=text_encoder, trainer=trainer, ) del stats, train_dataset, train_dataloader return g_speed_ls = deque(maxlen=128) def train_one_epoch( epoch: int, is_first_ep: bool, start_it: int, start_global_it: int, me: misc.MetricLogger, saver: CKPTSaver, args: arg_util.Args, dataloader_iter, iters_train: int, text_tokenizer: T5TokenizerFast, text_encoder: T5EncoderModel, trainer, ): # IMPORTANT: import heavy packages after the Dataloader object creation/iteration to avoid OOM step_cnt = 0 header = f'[Ep]: [{epoch:4d}/{args.epoch}]' last_touch = time.time() g_it, max_it = epoch * iters_train, args.epoch * iters_train doing_profiling = args.prof and epoch == 0 and (args.profall or dist.is_master()) maybe_record_function = record_function if doing_profiling else nullcontext trainer.gpt_wo_ddp.maybe_record_function = maybe_record_function last_t_perf = time.time() speed_ls: deque = g_speed_ls FREQ = min(args.prof_freq, iters_train//2-1) NVIDIA_IT_PLUS_1 = set(FREQ*i for i in (1, 2, 3, 4, 6, 8)) ranges = set([2 ** i for i in range(20)]) if epoch <= 1: ranges |= {1, 2, 3, 4, 6, 8, 10, 12, 16, 20, 24, 32, 40} PRINTABLE_IT_PLUS_1 = set(FREQ*i for i in ranges) 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 ============================================= for it, data in me.log_every(start_it, iters_train, dataloader_iter, args.log_freq, args.log_every_iter, header, args): g_it = epoch * iters_train + it # mfutool.step() # mfu_val = mfutool.get_mfu() * 100 # to percent # print(f"[MFU] step={g_it}, mfu={mfu_val:.2f} %, mfu.iter_time = {mfutool.iter_time():.4f} s") if (it+1) % FREQ == 0: speed_ls.append((time.time() - last_t_perf) / FREQ) last_t_perf = time.time() if (g_it+1) % args.save_model_iters_freq == 0: if args.checkpoint_type == 'torch': saver.sav(args=args, g_it=(g_it+1), next_ep=epoch, next_it=it+1, trainer=trainer, acc_str=f'[todo]', eval_milestone=None, also_save_to=None, best_save_to=None) elif args.checkpoint_type == 'omnistore': saver.sav(args=args, global_it=(g_it+1), next_ep=epoch, next_it=it+1, fsdp_object=trainer.gpt, optimizer_object=trainer.gpt_opt.optimizer, acc_str=None, eval_milestone=None) with maybe_record_function('before_train'): # [get data] images, captions, raw_features_bcthw, feature_cache_files4images, media = data['images'], data['captions'], data['raw_features_bcthw'], data['feature_cache_files4images'], data['media'] # # [prepare text features] if args.text_tokenizer_type == 'flan_t5': tokens = text_tokenizer(text=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) else: text_features = text_encoder(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) 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] # [logging] if dist.is_local_master() and (it >= start_it + 10) and (time.time() - last_touch > 90): args.dump_log() last_touch = time.time() # [get scheduled hyperparameters] progress = g_it / (max_it - 1) clip_decay_ratio = (0.3 ** (20 * progress) + 0.2) if args.cdec else 1 stepping = (g_it + 1) % args.ac == 0 step_cnt += int(stepping) with maybe_record_function('in_training'): grad_norm_t, scale_log2_t = trainer.train_step( epoch=epoch, it=it, g_it=g_it, stepping=stepping, clip_decay_ratio=clip_decay_ratio, metric_lg=me, inp_B3HW=images, raw_features_bcthw=raw_features_bcthw, feature_cache_files4images=feature_cache_files4images, text_cond_tuple=text_cond_tuple, media=media, args=args, ) with maybe_record_function('after_train'): me.update(tlr=args.tlr) # ============================================= iteration loop ends ============================================= me.synchronize_between_processes() return {k: meter.global_avg for k, meter in me.meters.items()} def main(): args: arg_util.Args = arg_util.init_dist_and_get_args() main_train(args) print(f'final args:\n\n{str(args)}') args.dump_log() if isinstance(sys.stdout, dist.BackupStreamToFile) and isinstance(sys.stderr, dist.BackupStreamToFile): sys.stdout.close(), sys.stderr.close() dist.barrier() if __name__ == '__main__': main()