Spaces:
Running on Zero
Running on Zero
| from contextlib import nullcontext | |
| import torch | |
| import os | |
| from tqdm import tqdm | |
| from einops import rearrange | |
| from models.dynamics import DynamicsModel | |
| from datasets.data_utils import visualize_reconstruction, load_data_and_data_loaders | |
| from tqdm import tqdm | |
| from einops import rearrange | |
| from utils.wandb_utils import ( | |
| init_wandb, log_training_metrics, log_learning_rate, log_system_metrics, finish_wandb, log_action_distribution | |
| ) | |
| from utils.scheduler_utils import create_cosine_scheduler | |
| from utils.utils import ( | |
| readable_timestamp, | |
| save_training_state, | |
| load_videotokenizer_from_checkpoint, | |
| load_latent_actions_from_checkpoint, | |
| load_dynamics_from_checkpoint, | |
| prepare_pipeline_run_root, | |
| prepare_stage_dirs, | |
| ) | |
| from utils.config import DynamicsConfig, load_stage_config_merged | |
| import 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(): | |
| # dynamics config merged with training_config.yaml (training takes priority), plus CLI overrides | |
| args: DynamicsConfig = load_stage_config_merged(DynamicsConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'dynamics.yaml')) | |
| # DDP setup | |
| dist_setup = init_distributed_from_env() | |
| # run save dir if it doesn't exist (running not from full train) | |
| is_main = dist_setup['is_main'] | |
| run_root = os.environ.get('NG_RUN_ROOT_DIR') | |
| if not run_root: | |
| run_root, _ = prepare_pipeline_run_root(base_cwd=os.getcwd()) | |
| stage_dir, checkpoints_dir, visualizations_dir = prepare_stage_dirs(run_root, 'dynamics') | |
| if is_main: | |
| print(f"Dynamics Training") | |
| print(f"Results will be saved in {stage_dir}") | |
| # load video tokenizer and latent action model | |
| if os.path.isdir(args.video_tokenizer_path): | |
| video_tokenizer, vq_ckpt = load_videotokenizer_from_checkpoint( | |
| checkpoint_path=args.video_tokenizer_path, | |
| device=args.device, | |
| is_distributed=dist_setup['is_distributed'], | |
| ) | |
| video_tokenizer.eval() | |
| for p in video_tokenizer.parameters(): | |
| p.requires_grad = False | |
| else: | |
| raise FileNotFoundError(f"Video tokenizer checkpoint not found at {args.video_tokenizer_path}") | |
| if os.path.isdir(args.latent_actions_path): | |
| latent_action_model, latent_action_ckpt = load_latent_actions_from_checkpoint( | |
| checkpoint_path=args.latent_actions_path, | |
| device=args.device, | |
| is_distributed=dist_setup['is_distributed'], | |
| ) | |
| unwrap_model(latent_action_model).eval() | |
| for p in unwrap_model(latent_action_model).parameters(): | |
| p.requires_grad = False | |
| else: | |
| raise FileNotFoundError(f"Latent Action Model checkpoint not found at {args.latent_actions_path}") | |
| # init dynamics model and optional ckpt load | |
| dynamics_model = DynamicsModel( | |
| 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, | |
| conditioning_dim=unwrap_model(latent_action_model).action_dim, | |
| latent_dim=args.latent_dim, | |
| num_bins=args.num_bins, | |
| use_moe=getattr(args, 'use_moe', False), | |
| num_experts=getattr(args, 'num_experts', 4), | |
| top_k_experts=getattr(args, 'top_k_experts', 2), | |
| moe_aux_loss_coeff=getattr(args, 'moe_aux_loss_coeff', 0.01), | |
| ).to(args.device) | |
| if args.checkpoint: | |
| dynamics_model, _ = load_dynamics_from_checkpoint( | |
| checkpoint_path=args.checkpoint, | |
| device=args.device, | |
| model=dynamics_model, | |
| is_distributed=dist_setup['is_distributed'], | |
| ) | |
| # optional DDP, compile, param count, tf32 | |
| print_param_count_if_main(dynamics_model, "DynamicsModel", is_main) | |
| if args.compile: | |
| video_tokenizer = torch.compile(video_tokenizer, mode="reduce-overhead", fullgraph=False, dynamic=True) | |
| latent_action_model = torch.compile(latent_action_model, mode="reduce-overhead", fullgraph=False, dynamic=True) | |
| dynamics_model = torch.compile(dynamics_model, mode="reduce-overhead", fullgraph=False, dynamic=True) | |
| print("Compiled all models for training.") | |
| dynamics_model = prepare_model_for_distributed( | |
| dynamics_model, | |
| args.distributed, | |
| model_type=dynamics_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(dynamics_model, args) | |
| # cosine scheduler for lr warmup and AMP grad scaler | |
| 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, | |
| 'dynamics_losses': [], | |
| 'loss_vals': [], | |
| } | |
| # init wandb | |
| if args.use_wandb and is_main: | |
| run_name = f"dynamics_{readable_timestamp()}" | |
| init_wandb(args.wandb_project, asdict(args), run_name) | |
| unwrap_model(dynamics_model).train() | |
| # 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_loader, _, _ = 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, | |
| ) | |
| train_iter = iter(training_loader) | |
| use_moe = getattr(args, 'use_moe', False) | |
| for i in tqdm(range(0, args.n_updates), disable=not is_main): | |
| for opt in optimizers: | |
| opt.zero_grad(set_to_none=True) | |
| if isinstance(dynamics_model, FSDPModule): | |
| dynamics_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) # [batch_size, seq_len, channels, height, width] | |
| # get video tokens for batch | |
| video_tokens = video_tokenizer.tokenize(x) # [B, T, P] | |
| video_latents = video_tokenizer.quantizer.get_latents_from_indices(video_tokens, dim=-1) # [B, T, P, L] | |
| if args.use_actions: | |
| quantized_actions = latent_action_model.encode(x) # [B, T - 1, A] | |
| else: | |
| quantized_actions = None | |
| # predict masked frame latents with dynamics model (masking in dynamics model) | |
| with train_ctx: | |
| predicted_next_logits, mask_positions, loss = dynamics_model( | |
| video_latents, training=True, conditioning=quantized_actions, targets=video_tokens | |
| ) | |
| # add MoE load-balancing auxiliary loss | |
| moe_aux_scalar = 0.0 | |
| if use_moe: | |
| aux_loss = unwrap_model(dynamics_model).transformer.moe_aux_loss() | |
| moe_aux_scalar = aux_loss.item() | |
| loss = loss + aux_loss | |
| if isinstance(dynamics_model, FSDPModule): | |
| if (micro_batch + 1) % args.gradient_accumulation_steps == 0: | |
| dynamics_model.set_requires_gradient_sync(True) | |
| loss /= args.gradient_accumulation_steps | |
| loss.backward() | |
| results['n_updates'] = i | |
| results['dynamics_losses'].append(loss.detach().cpu()) | |
| results['loss_vals'].append(loss.detach().cpu()) | |
| torch.nn.utils.clip_grad_norm_(unwrap_model(dynamics_model).parameters(), max_norm=1.0) | |
| for opt in optimizers: | |
| opt.step() | |
| for sched in schedulers: | |
| sched.step() | |
| # wandb logging | |
| if args.use_wandb and is_main: | |
| log_dict = {'train/loss': loss.item()} | |
| if use_moe: | |
| log_dict['train/moe_aux_loss'] = moe_aux_scalar | |
| # per-expert utilization across blocks | |
| expert_util = unwrap_model(dynamics_model).transformer.moe_expert_utilization() | |
| for block_idx, fracs in expert_util.items(): | |
| for expert_i, frac in enumerate(fracs.tolist()): | |
| log_dict[f'moe/block{block_idx}_expert{expert_i}_frac'] = frac | |
| wandb.log(log_dict, step=i) | |
| log_system_metrics(i) | |
| log_learning_rate(optimizers[0], i) | |
| if args.use_actions: | |
| action_indices = latent_action_model.quantizer.get_indices_from_latents(quantized_actions) | |
| log_action_distribution(action_indices, i, args.n_actions) | |
| # save model and visualize results | |
| if i % args.log_interval == 0: | |
| if args.use_wandb: | |
| predicted_next_indices = torch.argmax(predicted_next_logits, dim=-1) | |
| predicted_next_latents = video_tokenizer.quantizer.get_latents_from_indices(predicted_next_indices, dim=-1) | |
| with torch.no_grad(): | |
| predicted_frames = video_tokenizer.decoder(predicted_next_latents[:16]) # [B, T, C, H, W] | |
| # convert mask_positions to patch-level mask for visualization | |
| B, T, P = mask_positions.shape | |
| patch_size = args.patch_size | |
| H, W = args.frame_size, args.frame_size | |
| pixel_mask = torch.zeros(B, T, H, W, device=mask_positions.device) | |
| # for each pixel patch, mask if equivalent token is masked | |
| for b in range(B): | |
| for t in range(T): | |
| for p in range(P): | |
| if mask_positions[b, t, p]: | |
| patch_row = (p // (W // patch_size)) * patch_size | |
| patch_col = (p % (W // patch_size)) * patch_size | |
| pixel_mask[b, t, patch_row:patch_row+patch_size, patch_col:patch_col+patch_size] = 1 # assigning 1 to the patch in the mask of dim [1, 1, Hp, Wp] | |
| pixel_mask_expanded = rearrange(pixel_mask, 'b t h w -> b t 1 h w') | |
| masked_frames = x * (1 - pixel_mask_expanded) | |
| hyperparameters = args.__dict__ | |
| ckpt_path = save_training_state(dynamics_model, optimizers[0], schedulers[0], hyperparameters, checkpoints_dir, prefix='dynamics', step=i) | |
| # save secondary optimizer/scheduler state when using split optimizers (Muon+AdamW) | |
| if len(optimizers) > 1: | |
| import pathlib | |
| torch.save( | |
| {f'optimizer_{i}': o.state_dict() for i, o in enumerate(optimizers)}, | |
| pathlib.Path(ckpt_path) / 'all_optimizers.pt', | |
| ) | |
| torch.save( | |
| {f'scheduler_{i}': s.state_dict() for i, s in enumerate(schedulers)}, | |
| pathlib.Path(ckpt_path) / 'all_schedulers.pt', | |
| ) | |
| if is_main: | |
| save_path = os.path.join(visualizations_dir, f'dynamics_prediction_step_{i}.png') | |
| visualize_reconstruction(masked_frames[:16].cpu(), predicted_frames[:16].cpu(), 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() | |