from contextlib import nullcontext import torch import os from models.video_tokenizer import VideoTokenizer from datasets.data_utils import visualize_reconstruction, load_data_and_data_loaders from utils.scheduler_utils import create_cosine_scheduler from tqdm import tqdm import wandb from utils.utils import readable_timestamp, save_training_state, prepare_stage_dirs, prepare_pipeline_run_root from utils.config import VideoTokenizerConfig, load_stage_config_merged from utils.utils import save_training_state, load_videotokenizer_from_checkpoint from utils.wandb_utils import init_wandb, log_training_metrics, log_system_metrics, log_learning_rate, finish_wandb from dataclasses import asdict from utils.distributed import init_distributed_from_env, prepare_model_for_distributed, unwrap_model, print_param_count_if_main, cleanup_distributed from torch.distributed.fsdp import FSDPModule def main(): # vidtokenizer config merged with training_config.yaml (training takes priority), plus CLI overrides args: VideoTokenizerConfig = load_stage_config_merged(VideoTokenizerConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'video_tokenizer.yaml')) # DDP setup dist_setup = init_distributed_from_env() # run save dir if it doesn't exist (running not from full train) timestamp = readable_timestamp() run_root = os.environ.get('NG_RUN_ROOT_DIR') if not run_root: run_root, _ = prepare_pipeline_run_root(base_cwd=os.getcwd()) is_main = dist_setup['is_main'] stage_dir, checkpoints_dir, visualizations_dir = prepare_stage_dirs(run_root, 'video_tokenizer') if is_main: print(f"Video Tokenizer Training") print(f'Results will be saved in {stage_dir}') # dataloader data_overrides = {} if hasattr(args, 'fps') and args.fps is not None: data_overrides['fps'] = args.fps if hasattr(args, 'preload_ratio') and args.preload_ratio is not None: data_overrides['preload_ratio'] = args.preload_ratio training_data, validation_data, training_loader, validation_loader, x_train_var = load_data_and_data_loaders( dataset=args.dataset, batch_size=args.batch_size_per_gpu, num_frames=args.context_length, distributed=dist_setup['is_distributed'], rank=dist_setup['device_mesh'].get_rank() if dist_setup['device_mesh'] is not None else 0, world_size=dist_setup['world_size'], **data_overrides, ) # print("Length of training data:", len(training_data)) # print("Length of validation data:", len(validation_data)) # init model and optional ckpt load model = VideoTokenizer( frame_size=(args.frame_size, args.frame_size), patch_size=args.patch_size, embed_dim=args.embed_dim, num_heads=args.num_heads, hidden_dim=args.hidden_dim, num_blocks=args.num_blocks, latent_dim=args.latent_dim, num_bins=args.num_bins, ).to(args.device) if args.checkpoint: model, _ = load_videotokenizer_from_checkpoint( args.checkpoint, args.device, model, dist_setup['is_distributed'], ) # optional DDP, compile, param count, tf32 print_param_count_if_main(model, "VideoTokenizer", is_main) if args.compile: model = torch.compile(model, mode="reduce-overhead", fullgraph=False, dynamic=True) model = prepare_model_for_distributed( model, args.distributed, model_type=model.model_type, device_mesh=dist_setup['device_mesh'], ) if args.tf32: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # create optimizer(s) — AdamW or Muon+AdamW split from utils.optimizer_utils import create_optimizer optimizers = create_optimizer(model, args) # cosine scheduler for lr warmup schedulers = [create_cosine_scheduler(opt, args.n_updates) for opt in optimizers] train_ctx = torch.amp.autocast(args.device, enabled=True, dtype=torch.bfloat16) if args.amp and not args.distributed.use_fsdp else nullcontext() results = { 'n_updates': 0, 'loss_vals': [], } # init wandb if args.use_wandb and is_main: cfg = asdict(args) cfg.update({'timestamp': timestamp}) run_name = f"video_tokenizer_{timestamp}" init_wandb(args.wandb_project, cfg, run_name) unwrap_model(model).train() train_iter = iter(training_loader) for i in tqdm(range(args.n_updates), disable=not is_main): for opt in optimizers: opt.zero_grad(set_to_none=True) if isinstance(model, FSDPModule): model.set_requires_gradient_sync(False) if args.compile: torch.compiler.cudagraph_mark_step_begin() for micro_batch in range(args.gradient_accumulation_steps): try: (x, _) = next(train_iter) except StopIteration: train_iter = iter(training_loader) # reset iterator when epoch ends (x, _) = next(train_iter) x = x.to(args.device, non_blocking=True) with train_ctx: loss, x_hat = model(x) loss /= args.gradient_accumulation_steps if isinstance(model, FSDPModule): if (micro_batch + 1) % args.gradient_accumulation_steps == 0: model.set_requires_gradient_sync(True) loss.backward() torch.nn.utils.clip_grad_norm_(unwrap_model(model).parameters(), max_norm=1.0) for opt in optimizers: opt.step() for sched in schedulers: sched.step() results["loss_vals"].append(loss.cpu().detach()) results["n_updates"] = i # wandb logging if args.use_wandb and is_main: metrics = { 'loss': loss.item(), 'learning_rate': schedulers[0].get_last_lr()[0], } log_training_metrics(i, metrics, prefix='train') log_system_metrics(i) log_learning_rate(optimizers[0], i) # save model and visualize results if i % args.log_interval == 0: if args.use_wandb: with torch.no_grad(): indices = unwrap_model(model).tokenize(x) unique_codes = torch.unique(indices).numel() codebook_usage = unique_codes / unwrap_model(model).codebook_size if is_main: wandb.log({'train/codebook_usage': codebook_usage}, step=i) hyperparameters = args.__dict__ save_training_state(model, optimizers[0], schedulers[0], hyperparameters, checkpoints_dir, prefix='video_tokenizer', step=i) if is_main: x_hat_vis = x_hat.detach().cpu() x_vis = x.detach().cpu() save_path = os.path.join(visualizations_dir, f'video_tokenizer_recon_step_{i}.png') visualize_reconstruction(x_vis[:16], x_hat_vis[:16], save_path) print('\n Step', i, 'Loss:', torch.mean(torch.stack(results["loss_vals"][-args.log_interval:])).item()) # finish wandb if args.use_wandb and is_main: finish_wandb() cleanup_distributed(dist_setup['is_distributed']) if __name__ == "__main__": main()