from json import load import os import argparse import math import glob import time import logging from distutils.util import strtobool from copy import deepcopy import gc gc.disable() import os.path as osp import torch import torch.nn.functional as F import torch.optim as optim import torch.distributed as dist from torch.profiler import record_function as torch_record_function from contextlib import nullcontext from torch.nn.parallel import DistributedDataParallel as DDP from safetensors.torch import load_file from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import StateDictType, FullStateDictConfig from videovae.utils.misc import data_prefix_manager, COLOR_BLUE, COLOR_RESET, is_torch_optim_sch from videovae.utils.distributed import init_distributed_mode, reduce_losses, average_losses, _FSDP from videovae.utils.ema import update_ema, requires_grad from videovae.models.discriminator import ImageDiscriminator, VideoDiscriminator from videovae.data import VideoData from videovae.modules import build_lpips_model from videovae.modules.loss import get_disc_loss, adopt_weight from videovae.utils.misc import get_last_ckpt, seed_everything, print_gpu_usage, print_model_summary, version_checker from videovae.utils.init_models import init_vae_only, init_vit_from_image, resume_from_ckpt, init_cnn_from_image, load_cnn from videovae.utils.nan_detector import NanDetector from videovae.utils.arguments import MainArgs, add_model_specific_args, init_args, format_args from videovae.utils.scheduler import get_lambda from videovae.utils.mfu import register_mfu_hook, get_mfu, get_tflops, get_tflops_dict from videovae.utils.context_parallel import ContextParallelUtils as cp def save_model(fsdp_model, rank, model_path, global_step): # FSDP推荐用 state_dict_type=FULL_STATE_DICT 来保存 with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)): state_dict = fsdp_model.state_dict() os.makedirs(os.path.dirname(model_path), exist_ok=True) torch.save({'vae': state_dict, 'step': global_step}, model_path) print(f"模型已保存到 {model_path}") # enable_timeline_sdk = strtobool(os.getenv("GenAI_USE_TIMELINE_SDK", "0")) enable_timeline_sdk = False def split_to_ranks(x): bs = x.shape[0] cp_size = cp.get_cp_size() if cp_size > 1 and bs % cp_size == 0: cp_rank = cp.get_cp_rank() return x.chunk(cp_size, dim=0)[cp_rank] else: return x if enable_timeline_sdk: try: import bytedance.ndtimeline as ndtimeline except ImportError: print(f"import vescale.ndtimeline failed, skipped") enable_timeline_sdk = False def init_data_scheduler(video_ranks_ratio: float = -1.0, cp_size: int = 1): if video_ranks_ratio < 0: return None,None cp_size = max(1, cp_size) rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() video_ranks = list(range(int((world_size * video_ranks_ratio) // cp_size) * cp_size)) # align to cp_size for video image_ranks = list(range(len(video_ranks), world_size)) print(f"[info] video_ranks: {video_ranks}, image_ranks: {image_ranks}") if rank in image_ranks: group = torch.distributed.new_group(image_ranks) dataset_type_on_this_rank = "image" else: group = torch.distributed.new_group(video_ranks) dataset_type_on_this_rank = "video" return group, dataset_type_on_this_rank def main(): parser = argparse.ArgumentParser() parser = MainArgs.add_main_args(parser) parser = VideoData.add_data_specific_args(parser) args, unknown = parser.parse_known_args() args, parser, vae_model = add_model_specific_args(args, parser) args = parser.parse_args() args = init_args(args) # post process args # init data_prefix_manager data_prefix_manager.set_data_root(args.data_root, username=args.username) args.default_root_dir = data_prefix_manager(args.default_root_dir) # Setup DDP: init_distributed_mode(args) rank = dist.get_rank() world_size = dist.get_world_size() device = rank % torch.cuda.device_count() seed_everything(args.seed) torch.cuda.set_device(device) # init context parallel cp_cfg = {"cp_size": args.context_parallel_size} cp.initialize_context_parallel(cp_cfg) ds_group, ds_type = init_data_scheduler(args.video_ranks_ratio, cp_size = args.context_parallel_size) # Setup an experiment folder: checkpoint_dir = f"{args.default_root_dir}/checkpoints" # Stores saved model checkpoints os.makedirs(checkpoint_dir, exist_ok=True) if rank == 0: script_str = format_args(args) with open(os.path.join(args.default_root_dir, "script.sh"), "w") as f: f.write(script_str) print(f"{COLOR_BLUE}Experiment directory created at {args.default_root_dir}{COLOR_RESET}") import wandb wandb_project = "HBQ_Tokenizer" wandb.init( project=wandb_project, name=os.path.basename(os.path.normpath(args.default_root_dir)), dir=args.default_root_dir, config=args, mode="offline" if args.debug else "online" ) # init model vae = vae_model(args).to(device) if rank == 0: model_arch_save_path = os.path.join(args.default_root_dir, "model_arch.txt") print(f"{COLOR_BLUE}Logging model architecture at {model_arch_save_path}{COLOR_RESET}") with open(model_arch_save_path, "w") as f: f.write(str(vae)) image_disc = ImageDiscriminator(args).to(device) video_disc = VideoDiscriminator(args).to(device) # init optimizers and schedulers if args.optim_type == "Adam": vae_optim = torch.optim.Adam elif args.optim_type == "AdamW": vae_optim = torch.optim.AdamW if args.disc_optim_type is None: disc_optim = vae_optim elif args.disc_optim_type == "rmsprop": disc_optim = torch.optim.RMSprop def get_param_groups(model): decay = [] no_decay = [] for name, param in model.named_parameters(): if param.requires_grad: if len(param.shape) == 1 or name.endswith(".bias") or ('scale_learnable_parameters' in name): no_decay.append(param) print(f'disable weight deacy for {name}') else: decay.append(param) optimizer_grouped_parameters = [ {'params': decay, 'weight_decay': 0.01}, {'params': no_decay, 'weight_decay': 0.0} ] return optimizer_grouped_parameters opt_vae = vae_optim(get_param_groups(vae), lr=args.lr, betas=(args.beta1, args.beta2)) if disc_optim == torch.optim.RMSprop: opt_image_disc = disc_optim(image_disc.parameters(), lr=args.lr * args.dis_lr_multiplier) opt_video_disc = disc_optim(video_disc.parameters(), lr=args.lr * args.dis_lr_multiplier) else: opt_image_disc = disc_optim(image_disc.parameters(), lr=args.lr * args.dis_lr_multiplier, betas=(args.beta1, args.beta2)) opt_video_disc = disc_optim(video_disc.parameters(), lr=args.lr * args.dis_lr_multiplier, betas=(args.beta1, args.beta2)) if args.scheduler == "no": sch_vae, sch_image_disc, sch_video_disc = None, None, None else: lr_lambda = get_lambda(args) sch_vae = optim.lr_scheduler.LambdaLR(opt_vae, lr_lambda) sch_image_disc = optim.lr_scheduler.LambdaLR(opt_image_disc, lr_lambda) sch_video_disc = optim.lr_scheduler.LambdaLR(opt_video_disc, lr_lambda) ### ema ema = None if args.ema == "yes": ema = deepcopy(vae).to(device) # Create an EMA of the model for use after training requires_grad(ema, False) print(f"EMA Parameters: {sum(p.numel() for p in ema.parameters()):,}") update_ema(ema, vae, decay=0) # Ensure EMA is initialized with synced weights ema.eval() # EMA model should always be in eval mode model_optims = { "vae" : vae, "image_disc" : image_disc, "video_disc" : video_disc, "opt_vae" : opt_vae, "opt_image_disc" : opt_image_disc, "opt_video_disc" : opt_video_disc, "sch_vae" : sch_vae, "sch_image_disc" : sch_image_disc, "sch_video_disc" : sch_video_disc, "ema": ema, } ### Resume from checkpoint in default_root_dir or load pretrained weights if specified ckpt_path = None assert not args.default_root_dir is None # required argument ckpt_path = get_last_ckpt(args.default_root_dir) init_step = 0 if ckpt_path: print(f"Resuming from {ckpt_path}") state_dict = torch.load(ckpt_path, map_location="cpu") model_optims, init_step = resume_from_ckpt(state_dict, model_optims, load_optims=args.zero<=0, remove_disc=args.remove_disc, ckpt_path=ckpt_path, args=args) elif args.pretrained is not None: args.pretrained = data_prefix_manager(args.pretrained) # read weight state_dict = torch.load(args.pretrained, map_location="cpu", weights_only=True) if args.pretrained_ema == "yes": state_dict["vae"] = state_dict["ema"] # replace vae weights with ema weight # load model if args.pretrained_mode == "weights": model_optims, _ = resume_from_ckpt(state_dict, model_optims, load_optims=False, remove_disc=args.remove_disc, remove_enlarge_factors=args.remove_enlarge_factors, args=args) # load all models and optims del state_dict else: raise NotImplementedError print(f"Successfully loaded ckpt {args.pretrained}, pretrained_mode {args.pretrained_mode}") # init dataloader data = VideoData(args, ds_group = ds_group, ds_type = ds_type) dataloaders = data.train_dataloader() dataloader_iters = [iter(loader) for loader in dataloaders] ### init epoch in resuming dataloader_init_epoch = ( init_step if init_step > 0 # in case of resuming else args.dataloader_init_epoch if args.dataloader_init_epoch > 0 # in case of fintuning else 0 ) data_epochs = [dataloader_init_epoch for _ in dataloaders] for idx in range(len(dataloaders)): print(f"Reset the {idx}th dataloader as epoch {data_epochs[idx]}") if hasattr(dataloaders[idx], "sampler"): dataloaders[idx].sampler.set_epoch(data_epochs[idx]) else: raise NotImplementedError ### torch.compile after loading all weights print_model_summary([vae, image_disc, video_disc]) if args.zero > 0: from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, ShardingStrategy, MixedPrecision, ) def my_policy( module: torch.nn.Module, recurse: bool, **kwargs, ) -> bool: return True auto_wrap_policy = my_policy vae = FSDP( vae, device_id=device, sharding_strategy=ShardingStrategy.FULL_SHARD, mixed_precision=None, auto_wrap_policy=auto_wrap_policy, use_orig_params=True, sync_module_states=True, limit_all_gathers=True, device_mesh=None, ).to(device) # vae = _FSDP(vae, device, args.zero) else: vae = DDP(vae.to(device), device_ids=[args.gpu], bucket_cap_mb=args.bucket_cap_mb, find_unused_parameters=True) image_disc = DDP(image_disc.to(device), device_ids=[args.gpu], bucket_cap_mb=args.bucket_cap_mb) video_disc = DDP(video_disc.to(device), device_ids=[args.gpu], bucket_cap_mb=args.bucket_cap_mb) image_perceptual_model, video_perceptual_model = build_lpips_model(args) image_perceptual_model = image_perceptual_model.to(device) video_perceptual_model = video_perceptual_model.to(device) if args.compile == "yes": if args.vf_weight > 0 and args.vf_weight_approx < 0: torch._functorch.config.donated_buffer = False # This backward function was compiled with non-empty donated buffers which requires create_graph=False and retain_graph=False torch._dynamo.config.cache_size_limit = 256 torch._dynamo.config.accumulated_cache_size_limit = 4096 torch._dynamo.config.automatic_dynamic_shapes = True torch._dynamo.config.suppress_errors = False torch._dynamo.config.optimize_ddp = False if args.use_checkpoint else True for k in model_optims: if k != "ema" and model_optims[k] and not is_torch_optim_sch(model_optims[k]): print(f"compiling model {k}") if k == "vae": model_optims[k].encoder.compile()#options={'fx_graph_cache':True}) model_optims[k].decoder.compile()#options={'fx_graph_cache':True}) else: model_optims[k].compile()#options={'fx_graph_cache':True}) print(f"Successfully compiled all models") disc_loss = get_disc_loss(args.disc_loss_type) if enable_timeline_sdk: version_checker("2.0.0", "3.0.0") bmq_cluster = os.getenv('CUDA_TIMER_STREAM_KAFKA_CLUSTER', 'bmq_bigbang_3rd') bmq_topic = os.getenv('CUDA_TIMER_STREAM_KAFKA_TOPIC', 'megatron_cuda_timer_tracing_original') ndtimeline.init_ndtimers( mode="fsdp", mesh_shape=(world_size,), world_size=world_size, enable_streamer=True, post_handlers=[ndtimeline.handlers.MQNDHandler(mq_sinks=[ndtimeline.handlers.format_mq_sink(bmq_cluster, bmq_topic)])], ) ndtimeline.set_global_step(init_step) print(f"init timeline successfully") # init profiler def trace_handler(p): p.export_chrome_trace(os.path.join(args.default_root_dir, f"trace_step_{p.step_num}_rank_{rank}.json.gz")) if args.turn_on_profiler: print(f"start to init profiler") tp = torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], schedule=torch.profiler.schedule( wait=args.profiler_scheduler_wait_steps, warmup=3, active=2, repeat=1, ), with_stack=True, record_shapes=True, profile_memory=True, on_trace_ready=trace_handler ) tp.start() record_function = torch_record_function print(f"finish to init profiler") else: record_function = nullcontext start_time = time.time() cnt = 0 debug_root_dir = data_prefix_manager(f"debug/{cnt}") while os.path.exists(debug_root_dir): cnt += 1 debug_root_dir = data_prefix_manager(f"debug/{cnt}") os.makedirs(debug_root_dir, exist_ok=True) if cp.get_cp_rank() == 0: cp_group = cp.get_cp_group_rank() debug_f = open(os.path.join(debug_root_dir, f"rank_{rank}_cp_group_{cp_group}.txt"), "w") else: debug_f = None for global_step in range(init_step, args.max_steps): # if args.turn_on_profiler and tp: # tp.step() loss_dicts = [] if global_step == args.discriminator_iter_start - args.disc_pretrain_iter: logging.info(f"discriminator begins pretraining ") if global_step == args.discriminator_iter_start: log_str = "add GAN loss into training" if args.disc_pretrain_iter > 0: log_str += ", discriminator ends pretraining" logging.info(log_str) for idx in range(len(dataloader_iters)): try: _batch = next(dataloader_iters[idx]) except StopIteration: data_epochs[idx] += 1 print(f"Reset the {idx}th dataloader as epoch {data_epochs[idx]}") dataloaders[idx].sampler.set_epoch(data_epochs[idx]) dataloader_iters[idx] = iter(dataloaders[idx]) # update dataloader iter _batch = next(dataloader_iters[idx]) except Exception as e: raise e x = _batch["video"] _type = _batch["type"][0] if _type == "image" and ds_type is None: x = split_to_ranks(x) disc_factor = 1. with NanDetector(vae) if args.enable_nan_detector else nullcontext(): with record_function("vae"): if _type == "image": x, x_recon, flat_frames, flat_frames_recon, vae_loss_dict, vae_log_dict = vae(x, disc_factor, image_disc=image_disc, image_perceptual_model=image_perceptual_model) elif _type == "video": if debug_f is not None: debug_f.write(f'step {idx}, {_batch["path"]}\n') x, x_recon, flat_frames, flat_frames_recon, vae_loss_dict, vae_log_dict = vae( x, disc_factor, image_disc=image_disc, video_disc=video_disc, image_perceptual_model=image_perceptual_model, video_perceptual_model=video_perceptual_model, ) g_loss = sum(vae_loss_dict.values()) opt_vae.zero_grad() g_loss.backward() # print_gpu_usage("vae") if args.max_grad_norm > 0: torch.nn.utils.clip_grad_norm_(vae.parameters(), args.max_grad_norm) # from pnp.utils import detect_anomalous_params # detect_anomalous_params(g_loss, vae) opt_vae.step() opt_vae.zero_grad() # free memory if args.ema == "yes": update_ema(ema, vae.module) with record_function("disc"): disc_loss_dict = {} # args.discriminator_iter_start=-1 args.disc_pretrain_iter=0 discloss = d_image_loss = d_video_loss = torch.tensor(0.).to(x.device) ### enable pool warmup for disc_step in range(args.disc_optim_steps): require_optim = False if _type == "image": if args.image_disc_weight > 0: require_optim = True logits_image_real = image_disc(x, pool_name="real") logits_image_fake = image_disc(x_recon.detach(), pool_name="fake") d_image_loss = disc_loss(logits_image_real, logits_image_fake) discloss = d_image_loss * args.image_disc_weight disc_loss_dict["train/logits_image_real"] = logits_image_real.mean().detach() disc_loss_dict["train/logits_image_fake"] = logits_image_fake.mean().detach() disc_loss_dict["train/d_image_loss"] = discloss.detach() opt_discs, sch_discs = [opt_image_disc], [sch_image_disc] elif _type == "video": if args.image_disc_weight > 0 and args.gan_image4video == "yes": require_optim = True logits_image_real = image_disc(flat_frames.detach(), pool_name="real") logits_image_fake = image_disc(flat_frames_recon.detach(), pool_name="fake") d_image_loss = disc_loss(logits_image_real, logits_image_fake) disc_loss_dict["train/logits_image_real"] = logits_image_real.mean().detach() disc_loss_dict["train/logits_image_fake"] = logits_image_fake.mean().detach() disc_loss_dict["train/d_image_loss"] = (d_image_loss * args.image_disc_weight).detach() if args.video_disc_weight > 0: require_optim = True logits_video_real = video_disc(x.detach(), pool_name="real") logits_video_fake = video_disc(x_recon.detach(), pool_name="fake") d_video_loss = disc_loss(logits_video_real, logits_video_fake) disc_loss_dict["train/logits_video_real"] = logits_video_real.mean().detach() disc_loss_dict["train/logits_video_fake"] = logits_video_fake.mean().detach() disc_loss_dict["train/d_video_loss"] = (d_video_loss * args.video_disc_weight).detach() discloss = d_image_loss * args.image_disc_weight + d_video_loss * args.video_disc_weight opt_discs, sch_discs = [opt_image_disc, opt_video_disc], [sch_image_disc, sch_video_disc] discloss = disc_factor * discloss if require_optim: for opt_disc in opt_discs: opt_disc.zero_grad() discloss.backward() # print_gpu_usage("disc") if args.max_grad_norm_disc > 0: torch.nn.utils.clip_grad_norm_(image_disc.parameters(), args.max_grad_norm_disc) torch.nn.utils.clip_grad_norm_(video_disc.parameters(), args.max_grad_norm_disc) for opt_disc in opt_discs: opt_disc.step() for opt_disc in opt_discs: opt_disc.zero_grad() # free memory with record_function("loss"): loss_dict = {**vae_loss_dict, **disc_loss_dict, **vae_log_dict} if (global_step+1) % args.log_every == 0: reduced_loss_dict = reduce_losses(loss_dict) else: reduced_loss_dict = {} loss_dicts.append(reduced_loss_dict) # update scheduler if not sch_vae is None: sch_vae.step() for sch_disc in sch_discs: if not sch_disc is None: sch_disc.step() # if enable_timeline_sdk: # ndtimeline.inc_step() if (global_step+1) % args.log_every == 0: avg_loss_dict = average_losses(loss_dicts) torch.cuda.synchronize() end_time = time.time() iter_speed = (end_time - start_time) / args.log_every if args.mfu_logging == "yes": tflops = get_tflops() / args.log_every tflops_log_str = f"tflops={tflops:.1f}, " tflops_dict = get_tflops_dict(args.log_every) tflops_dict_log_str = f"tflops_Dict={tflops_dict}, " mfu = get_mfu(iter_speed) / args.log_every mfu_log_str = f"mfu={mfu:.3f}, " else: tflops_log_str = "" tflops_dict_log_str = "" mfu_log_str = "" if rank == 0: avg_loss_dict["lr"] = opt_vae.param_groups[0]['lr'] for key, value in avg_loss_dict.items(): wandb.log({key: value}, step=global_step) recons_loss_sum, video_perceptual_loss_sum = 0., 0. for key in avg_loss_dict: if 'recon_loss' in key: recons_loss_sum += avg_loss_dict[key] if 'video_perceptual_loss' in key: video_perceptual_loss_sum += avg_loss_dict[key] print(f'global_step={global_step}, recon_loss={recons_loss_sum:.4f}, ' \ f'video_perceptual_loss={video_perceptual_loss_sum:.4f}, ' \ f'iter_speed={iter_speed:.2f}s, ' \ f'{mfu_log_str}' \ f'{tflops_log_str}' \ f'{tflops_dict_log_str}' \ ) start_time = time.time() if enable_timeline_sdk: ndtimeline.flush() if (global_step+1) % args.ckpt_every == 0 and global_step != init_step: checkpoint_path = os.path.join(checkpoint_dir, f'model_step_{global_step}.ckpt') if args.zero > 0: save_model(vae, rank, checkpoint_path, global_step) else: if rank == 0: save_dict = {} for k in model_optims: model = model_optims[k] save_dict[k] = None if model is None \ else model.module.state_dict() if hasattr(model, "module") \ else model.state_dict() torch.save({ 'step': global_step, **save_dict, }, checkpoint_path) print(f'Checkpoint saved at step {global_step}') if (global_step+1) % args.manual_gc_interval == 0: gc.collect() if __name__ == '__main__': main()