# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. import argparse import os from pathlib import Path import wandb import torch import torch.distributed as dist import torch.nn.functional as F from torch.utils.data import DataLoader, DistributedSampler from tqdm import tqdm from torch.distributed.checkpoint.state_dict import ( get_model_state_dict, get_optimizer_state_dict, set_optimizer_state_dict, StateDictOptions, ) from safetensors.torch import save_file, load_file import json from .configs import VA_CONFIGS from .distributed.fsdp import shard_model, apply_ac from .distributed.util import ( _configure_model, init_distributed, dist_mean, dist_max ) from einops import rearrange from .modules.utils import ( load_transformer, ) from .utils import ( init_logger, logger, get_mesh_id, sample_timestep_id, data_seq_to_patch, warmup_constant_lambda, FlowMatchScheduler ) from .dataset import MultiLatentLeRobotDataset, dataset_indexes_ready from .mcp import shift_latents_for_mcp, validate_mcp_settings import gc class Trainer: def __init__(self, config): if config.enable_wandb and config.rank == 0: wandb.login(host=os.environ['WANDB_BASE_URL'], key=os.environ['WANDB_API_KEY']) self.wandb = wandb self.wandb.init( entity=os.environ["WANDB_TEAM_NAME"], project=os.getenv("WANDB_PROJECT", "va_robotwin"), # dir=log_dir, config=config, mode="online", name='test_lln' # name=os.path.basename(os.path.normpath(job_config.job.dump_folder)) ) logger.info("WandB logging enabled") self.step = 0 self.config = config self.device = torch.device(f"cuda:{config.local_rank}") self.dtype = config.param_dtype self.patch_size = config.patch_size self.enable_mcp = getattr(config, 'enable_mcp', True) if self.enable_mcp: validate_mcp_settings( num_mcp_depths=config.num_mcp_depths, mcp_blocks_per_depth=config.mcp_blocks_per_depth, mcp_hidden_collect_layers=config.mcp_hidden_collect_layers, mcp_loss_weights=config.mcp_loss_weights, ) # Load models logger.info("Loading models...") # Load and shard transformer with FSDP logger.info("Loading transformer...") if hasattr(config, 'resume_from') and config.resume_from: transformer_path = os.path.join(config.resume_from, 'transformer') if config.rank == 0: logger.info(f"Resuming from checkpoint: {transformer_path}") else: transformer_path = os.path.join(config.wan22_pretrained_model_name_or_path, 'transformer') self.transformer = load_transformer( transformer_path, torch_dtype=torch.float32, torch_device='cpu', attn_mode="flex", disable_mcp=not self.enable_mcp, ) if self.enable_mcp: validate_mcp_settings( num_mcp_depths=config.num_mcp_depths, mcp_blocks_per_depth=config.mcp_blocks_per_depth, mcp_hidden_collect_layers=config.mcp_hidden_collect_layers, mcp_loss_weights=config.mcp_loss_weights, num_layers=len(self.transformer.blocks), ) initialized = self.transformer.enable_mcp_training( num_mcp_depths=config.num_mcp_depths, mcp_blocks_per_depth=config.mcp_blocks_per_depth, mcp_hidden_collect_layers=config.mcp_hidden_collect_layers, init_from_backbone=config.mcp_init_from_backbone, ) if initialized: total_mcp_blocks = ( config.num_mcp_depths * config.mcp_blocks_per_depth) if config.mcp_init_from_backbone: num_backbone_blocks = len(self.transformer.blocks) source_start = ( num_backbone_blocks - config.mcp_blocks_per_depth) logger.info( f"Initializing {config.num_mcp_depths} MCP depths " f"({config.mcp_blocks_per_depth} blocks per depth) " f"from backbone blocks[{source_start}:" f"{num_backbone_blocks}]" ) for depth in range(config.num_mcp_depths): logger.info( f" MCP depth {depth + 1}: initialized from " f"backbone blocks[{source_start}:" f"{num_backbone_blocks}]" ) logger.info( f"Initialized {config.num_mcp_depths} x " f"{config.mcp_blocks_per_depth} = " f"{total_mcp_blocks} MCP blocks" ) else: logger.info( f"Initialized {total_mcp_blocks} MCP blocks from " "scratch" ) logger.info( "MCP hidden-state collection uses backbone indices " f"{list(config.mcp_hidden_collect_layers)}" ) logger.info( "Initialized MCP hidden fuser and input projections " "from scratch" ) else: logger.info( f"Loaded MCP modules from checkpoint: " f"{config.num_mcp_depths} depths x " f"{config.mcp_blocks_per_depth} blocks per depth" ) logger.info("Setting up activation checkpointing ...") apply_ac(self.transformer) logger.info("Setting up FSDP...") shard_fn = shard_model self.transformer = _configure_model( model=self.transformer, shard_fn=shard_fn, param_dtype=self.dtype, device=self.device, eval_mode=False, ) self.transformer.train() self.transformer.requires_grad_(True) # Optimizer self.optimizer = torch.optim.AdamW( [p for p in self.transformer.parameters() if p.requires_grad], lr=config.learning_rate, betas=(config.beta1, config.beta2), eps=1e-8, weight_decay=config.weight_decay, fused=True, foreach=False, ) self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lambda step: warmup_constant_lambda(step, warmup_steps=config.warmup_steps)) # Setup dataloaders logger.info("Setting up datasets...") cache_ready = False if ( config.world_size > 1 and getattr(config, 'enable_dataset_index_cache', True) ): cache_ready_flag = torch.zeros(1, device=self.device, dtype=torch.int) if config.rank == 0: cache_ready_flag.fill_(int(dataset_indexes_ready(config))) dist.broadcast(cache_ready_flag, src=0) cache_ready = bool(cache_ready_flag.item()) use_rank_zero_indexing = ( config.world_size > 1 and getattr(config, 'enable_dataset_index_cache', True) and not cache_ready ) if config.rank == 0 and cache_ready: logger.info( "Dataset index and Arrow caches are complete; " "loading all ranks concurrently" ) if use_rank_zero_indexing and config.rank != 0: dist.barrier() # Rank 0 has completed any requested rebuild at this point. config.rebuild_dataset_index_cache = False train_dataset = MultiLatentLeRobotDataset(config=config) if use_rank_zero_indexing: if config.rank == 0: dist.barrier() dist.barrier() if config.rank == 0: logger.info( "Dataset ready: %d samples from %d datasets " "(%d index cache hits, %d direct Arrow loads, %d rebuilt)", len(train_dataset), len(train_dataset._datasets), train_dataset.index_cache_hits, train_dataset.hf_cache_hits, train_dataset.index_cache_misses, ) train_sampler = DistributedSampler( train_dataset, num_replicas=config.world_size, rank=config.rank, shuffle=True, seed=42 ) if config.world_size > 1 else None self.train_loader = DataLoader( train_dataset, batch_size=config.batch_size, shuffle=(train_sampler is None), num_workers=config.load_worker, sampler=train_sampler, ) self.train_scheduler_latent = FlowMatchScheduler(shift=self.config.snr_shift, sigma_min=0.0, extra_one_step=True) self.train_scheduler_latent.set_timesteps(1000, training=True) self.train_scheduler_action = FlowMatchScheduler(shift=self.config.action_snr_shift, sigma_min=0.0, extra_one_step=True) self.train_scheduler_action.set_timesteps(1000, training=True) self.train_scheduler_mcp = None if self.enable_mcp: self.train_scheduler_mcp = FlowMatchScheduler( shift=self.config.mcp_snr_shift, sigma_min=0.0, extra_one_step=True, ) self.train_scheduler_mcp.set_timesteps(1000, training=True) self.save_dir = Path(config.save_root) / "checkpoints" self.save_dir.mkdir(parents=True, exist_ok=True) self.gradient_accumulation_steps = getattr(config, 'gradient_accumulation_steps', 1) self.train_loader_iter = None # if hasattr(config, 'resume_from') and config.resume_from: # self._load_training_state(config.resume_from) def _get_next_batch(self): """Get next batch from iterator, reset if epoch is finished.""" if self.train_loader_iter is None: self.train_loader_iter = iter(self.train_loader) try: batch = next(self.train_loader_iter) except StopIteration: # Reset sampler and iterator when epoch finishes if hasattr(self.train_loader.sampler, 'set_epoch'): self.train_loader.sampler.set_epoch(self.train_loader.sampler.epoch + 1) self.train_loader_iter = iter(self.train_loader) batch = next(self.train_loader_iter) return batch @torch.no_grad() def _add_noise(self, latent, train_scheduler, action_mask=False, action_mode=False, noisy_cond_prob=0., frame_shift=0): B, C, F, H, W = latent.shape timestep_ids = sample_timestep_id(batch_size=F, num_train_timesteps=train_scheduler.num_train_timesteps) noise = torch.zeros_like(latent).normal_() timesteps = train_scheduler.timesteps[timestep_ids].to(device=self.device) noisy_latents =train_scheduler.add_noise(latent, noise, timesteps, t_dim=2) targets =train_scheduler.training_target(latent, noise, timesteps) patch_f, patch_h, patch_w = self.patch_size if action_mode: patch_f = patch_h = patch_w = 1 latent_grid_id = get_mesh_id( latent.shape[-3] // patch_f, # F latent.shape[-2] // patch_h, # H latent.shape[-1] // patch_w, # W t=1 if action_mode else 0, # 1 for action mode (0 for latent), not used f_w=1, f_shift=frame_shift, action=action_mode ).to(self.device) # shape: [4, seq_len] latent_grid_id = latent_grid_id[None].repeat(B, 1, 1) if torch.rand(1).item() < noisy_cond_prob: cond_timestep_ids = sample_timestep_id( batch_size=F, min_timestep_bd=0.5, max_timestep_bd=1.0, num_train_timesteps=train_scheduler.num_train_timesteps, ) noise = torch.zeros_like(latent).normal_() cond_timesteps = train_scheduler.timesteps[cond_timestep_ids].to(device=self.device) latent = train_scheduler.add_noise(latent, noise, cond_timesteps, t_dim=2) else: cond_timesteps = torch.zeros_like(timesteps) if action_mask is not None: noisy_latents *= action_mask.float() targets *= action_mask.float() latent *= action_mask.float() return dict( timesteps=timesteps[None].repeat(B, 1), noisy_latents=noisy_latents, targets=targets, latent=latent, cond_timesteps=cond_timesteps[None].repeat(B, 1), grid_id=latent_grid_id, ) @torch.no_grad() def _prepare_input_dict(self, batch_dict): """Prepare input dict following infer code pattern from wan_va_server.py.""" chunk_size = torch.randint(1, 5, (1,)).item() # Generate grid_id following infer code (no batch dimension yet) # For action mode: get_mesh_id(shape[-3], shape[-2], shape[-1], t=1, f_w=1, f_shift, action=True) latent_dict = self._add_noise( latent=batch_dict['latents'], train_scheduler=self.train_scheduler_latent, action_mask=None, action_mode=False, noisy_cond_prob=0.5) action_dict = self._add_noise( latent=batch_dict['actions'], train_scheduler=self.train_scheduler_action, action_mask=batch_dict['actions_mask'], action_mode=True, noisy_cond_prob=0.0) latent_dict['text_emb'] = batch_dict['text_emb'] action_dict['text_emb'] = batch_dict['text_emb'] action_dict['actions_mask'] = batch_dict['actions_mask'] input_dict = { 'latent_dict': latent_dict, 'action_dict': action_dict, 'chunk_size': chunk_size, 'window_size': torch.randint(4, 65, (1,)).item(), } if self.enable_mcp: mcp_latent_dicts = [] for depth in range(self.config.num_mcp_depths): frame_shift = (depth + 1) * chunk_size shifted_latents, valid_mask = shift_latents_for_mcp( batch_dict['latents'], frame_shift) mcp_latent_dict = self._add_noise( latent=shifted_latents, train_scheduler=self.train_scheduler_mcp, action_mask=None, action_mode=False, noisy_cond_prob=0.0, frame_shift=frame_shift, ) mcp_latent_dict.pop('latent') mcp_latent_dict.pop('cond_timesteps') mcp_latent_dict['valid_mask'] = valid_mask mcp_latent_dicts.append(mcp_latent_dict) input_dict['mcp_latent_dicts'] = mcp_latent_dicts return input_dict def convert_input_format(self, input_dict): """Convert input dict to match transformer input format if needed.""" for key, value in input_dict.items(): input_dict[key] = value.to(self.device)#.to(self.dtype) return input_dict def compute_loss(self, input_dict, pred ): if self.enable_mcp: latent_pred, action_pred, mcp_pred_list = pred else: latent_pred, action_pred = pred mcp_pred_list = [] if self.enable_mcp and len(mcp_pred_list) != self.config.num_mcp_depths: raise RuntimeError( "MCP output depth count must match num_mcp_depths") action_pred = rearrange(action_pred, 'b (f n) c -> b c f n 1', f=input_dict['action_dict']['targets'].shape[-3]) latent_pred = data_seq_to_patch( self.patch_size, latent_pred, input_dict['latent_dict']['targets'].shape[-3], input_dict['latent_dict']['targets'].shape[-2], input_dict['latent_dict']['targets'].shape[-1], batch_size=latent_pred.shape[0]) Bn, Fn = input_dict['latent_dict']['timesteps'].shape latent_loss_weight = self.train_scheduler_latent.training_weight(input_dict['latent_dict']['timesteps'].flatten()).reshape(Bn, Fn) action_loss_weight = self.train_scheduler_action.training_weight(input_dict['action_dict']['timesteps'].flatten()).reshape(Bn, Fn) # Frame-wise video loss calculation latent_loss = F.mse_loss(latent_pred.float(), input_dict['latent_dict']['targets'].float().detach(), reduction='none') latent_loss = latent_loss * latent_loss_weight[:, None, :, None, None] # Permute to (B, F, H, W, C) and flatten to (B*F, H*W*C) latent_loss = latent_loss.permute(0, 2, 3, 4, 1) # (B, C, F, H, W) -> (B, F, H, W, C) latent_loss = latent_loss.flatten(0, 1).flatten(1) # (B, F, H, W, C) -> (B*F, H*W*C) # Sum per frame and compute mask per frame latent_loss_per_frame = latent_loss.sum(dim=1) # (B*F,) latent_mask_per_frame = torch.ones_like(latent_loss).sum(dim=1) # (B*F,) latent_loss = (latent_loss_per_frame / (latent_mask_per_frame + 1e-6)).mean() # Frame-wise action loss calculation action_loss = F.mse_loss(action_pred.float(), input_dict['action_dict']['targets'].float().detach(), reduction='none') action_loss = action_loss * action_loss_weight[:, None, :, None, None] action_loss = action_loss * input_dict['action_dict']['actions_mask'].float() # Permute to (B, F, H, W, C) and flatten to (B*F, H*W*C) action_loss = action_loss.permute(0, 2, 3, 4, 1) # (B, C, F, H, W) -> (B, F, H, W, C) action_mask = input_dict['action_dict']['actions_mask'].float().permute(0, 2, 3, 4, 1) # (B, C, F, H, W) -> (B, F, H, W, C) action_loss = action_loss.flatten(0, 1).flatten(1) # (B, F, H, W, C) -> (B*F, H*W*C) action_mask = action_mask.flatten(0, 1).flatten(1) # (B, F, H, W, C) -> (B*F, H*W*C) # Sum per frame and normalize by mask per frame action_loss_per_frame = action_loss.sum(dim=1) # (B*F,) action_mask_per_frame = action_mask.sum(dim=1) # (B*F,) action_loss = (action_loss_per_frame / (action_mask_per_frame + 1e-6)).mean() mcp_losses = [] for mcp_pred, mcp_latent_dict in zip( mcp_pred_list, input_dict.get('mcp_latent_dicts', [])): mcp_pred = data_seq_to_patch( self.patch_size, mcp_pred, mcp_latent_dict['targets'].shape[-3], mcp_latent_dict['targets'].shape[-2], mcp_latent_dict['targets'].shape[-1], batch_size=mcp_pred.shape[0], ) mcp_batch_size, mcp_num_frames = mcp_latent_dict[ 'timesteps'].shape mcp_loss_weight = self.train_scheduler_mcp.training_weight( mcp_latent_dict['timesteps'].flatten()).reshape( mcp_batch_size, mcp_num_frames) mcp_loss = F.mse_loss( mcp_pred.float(), mcp_latent_dict['targets'].float().detach(), reduction='none', ) mcp_loss = mcp_loss * mcp_loss_weight[:, None, :, None, None] valid_mask = mcp_latent_dict['valid_mask'].to( device=mcp_loss.device, dtype=mcp_loss.dtype) valid_count = valid_mask.expand_as(mcp_loss).sum() mcp_loss = (mcp_loss * valid_mask).sum() / valid_count.clamp_min(1.) mcp_losses.append(mcp_loss / self.gradient_accumulation_steps) return ( latent_loss / self.gradient_accumulation_steps, action_loss / self.gradient_accumulation_steps, mcp_losses, ) def _train_step(self, batch, batch_idx): """Train a single batch, returns losses for logging.""" batch = self.convert_input_format(batch) input_dict = self._prepare_input_dict(batch) should_sync = (batch_idx + 1) % self.gradient_accumulation_steps == 0 if not should_sync: self.transformer.set_requires_gradient_sync(False) else: self.transformer.set_requires_gradient_sync(True) output = self.transformer(input_dict, train_mode=True) latent_loss, action_loss, mcp_losses = self.compute_loss( input_dict, output) mcp_loss = sum( weight * depth_loss for weight, depth_loss in zip( self.config.mcp_loss_weights, mcp_losses) ) if mcp_losses else latent_loss.new_zeros(()) loss = latent_loss + action_loss + mcp_loss loss.backward() losses = { 'latent_loss': latent_loss.detach(), 'action_loss': action_loss.detach(), 'mcp_losses': [depth_loss.detach() for depth_loss in mcp_losses], 'mcp_loss': mcp_loss.detach(), } # Only update weights after accumulating gradients if should_sync: total_norm = torch.nn.utils.clip_grad_norm_(self.transformer.parameters(), 2.0) self.optimizer.step() self.lr_scheduler.step() self.optimizer.zero_grad() losses['total_norm'] = total_norm losses['should_log'] = True else: losses['should_log'] = False return losses def save_checkpoint(self,): """Save model checkpoint in the same format as pretrained model.""" try: state_dict = get_model_state_dict( self.transformer, options=StateDictOptions(full_state_dict=True, cpu_offload=True), ) state_dict_bf16 = {k: v.to(torch.bfloat16) for k, v in state_dict.items()} # optim_state = get_optimizer_state_dict( # self.transformer, self.optimizer, # options=StateDictOptions(full_state_dict=True, cpu_offload=True), # ) # Only rank 0 saves the checkpoint if self.config.rank == 0: checkpoint_dir = self.save_dir / f"checkpoint_step_{self.step}" checkpoint_dir.mkdir(parents=True, exist_ok=True) # Save transformer in the same format as pretrained model transformer_dir = checkpoint_dir / "transformer" transformer_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Saving transformer to {transformer_dir}") # Manually save in diffusers format (outside FSDP context to avoid deadlock) # Save model weights model_file = transformer_dir / "diffusion_pytorch_model.safetensors" save_file(state_dict_bf16, model_file) # Save config (copy from original transformer config and update _name_or_path) config_file = transformer_dir / "config.json" config_dict = dict(self.transformer.config) config_dict.pop('_name_or_path', None) with open(config_file, 'w') as f: json.dump(config_dict, f, indent=2) # # Save optimizer state and training metadata in PyTorch format # training_state_path = checkpoint_dir / "training_state.pt" # logger.info(f"Saving training state to {training_state_path}") # torch.save({ # 'step': self.step, # 'optimizer_state_dict': optim_state, # 'config': vars(self.config), # }, training_state_path) logger.info(f"Checkpoint saved successfully at step {self.step}") # Synchronize all processes after saving if dist.is_initialized(): dist.barrier() except Exception as e: if self.config.rank == 0: logger.error(f"Failed to save checkpoint: {e}") import traceback logger.error(traceback.format_exc()) # Ensure all processes stay synchronized even on error if dist.is_initialized(): dist.barrier() def _load_training_state(self, checkpoint_path): """Load training state (optimizer + step) after FSDP and optimizer creation.""" checkpoint_dir = Path(checkpoint_path) training_state_path = checkpoint_dir / "training_state.pt" if not training_state_path.exists(): if self.config.rank == 0: logger.warning(f"Training state not found: {training_state_path}, starting from step 0") return if self.config.rank == 0: logger.info(f"Loading training state from {training_state_path}") # All ranks load the training state directly training_state = torch.load(training_state_path, map_location='cpu', weights_only=False) # All ranks load optimizer state (required for FSDP) set_optimizer_state_dict( self.transformer, self.optimizer, optim_state_dict=training_state['optimizer_state_dict'], options=StateDictOptions(full_state_dict=True, strict=False) ) self.step = training_state.get('step', 0) if self.config.rank == 0: logger.info(f"Training state loaded, resuming from step {self.step}") # Synchronize all ranks if dist.is_initialized(): dist.barrier() def train(self): """Main training loop - train by steps instead of epochs.""" logger.info(f"Starting training for {self.config.num_steps} steps...") self.transformer.train() progress_bar = tqdm( total=self.config.num_steps, desc="Training", disable=(self.config.rank != 0), leave=True, dynamic_ncols=True, initial=self.step ) self.optimizer.zero_grad() accumulated_latent_losses = [] accumulated_action_losses = [] accumulated_mcp_losses = [ [] for _ in range(self.config.num_mcp_depths) ] if self.enable_mcp else [] accumulated_mcp_total_losses = [] step_in_accumulation = 0 while self.step < self.config.num_steps: # Get next batch (handles epoch reset automatically) batch = self._get_next_batch() losses = self._train_step(batch, step_in_accumulation) # Accumulate losses for logging accumulated_latent_losses.append(losses['latent_loss']) accumulated_action_losses.append(losses['action_loss']) for depth, mcp_loss in enumerate(losses['mcp_losses']): accumulated_mcp_losses[depth].append(mcp_loss) accumulated_mcp_total_losses.append(losses['mcp_loss']) step_in_accumulation += 1 # Log and checkpoint when optimizer steps if losses['should_log']: lr = self.lr_scheduler.get_last_lr()[0] # Average accumulated losses latent_loss_show = dist_mean(torch.stack(accumulated_latent_losses).sum()).detach().cpu().item() action_loss_show = dist_mean(torch.stack(accumulated_action_losses).sum()).detach().cpu().item() max_latent_loss_show = dist_max(torch.stack(accumulated_latent_losses).sum()).detach().cpu().item() max_action_loss_show = dist_max(torch.stack(accumulated_action_losses).sum()).detach().cpu().item() mcp_loss_shows = [ dist_mean(torch.stack(depth_losses).sum()).detach().cpu().item() for depth_losses in accumulated_mcp_losses ] mcp_total_loss_show = dist_mean( torch.stack(accumulated_mcp_total_losses).sum() ).detach().cpu().item() # Clear accumulated losses accumulated_latent_losses = [] accumulated_action_losses = [] accumulated_mcp_losses = [ [] for _ in range(self.config.num_mcp_depths) ] if self.enable_mcp else [] accumulated_mcp_total_losses = [] step_in_accumulation = 0 torch.cuda.synchronize() if self.step % self.config.gc_interval == 0: torch.cuda.empty_cache() gc.collect() if self.config.rank == 0: total_norm = losses['total_norm'] progress_bar.n += 1 postfix = { 'latent_loss': f'{latent_loss_show:.4f}', 'action_loss': f'{action_loss_show:.4f}', 'step': self.step, 'grad_norm': f'{total_norm.item():.2f}', 'lr': f'{lr:.2e}' } if self.enable_mcp: postfix['mcp_loss'] = f'{mcp_total_loss_show:.4f}' progress_bar.set_postfix(postfix) if self.config.enable_wandb: log_values = { 'loss_metrics/global_avg_video_loss': latent_loss_show, 'loss_metrics/global_avg_action_loss': action_loss_show, 'loss_metrics/global_max_video_loss': max_latent_loss_show, 'loss_metrics/global_max_action_loss': max_action_loss_show, 'grad_norm': total_norm.item(), 'lr': lr, } if self.enable_mcp: log_values['loss_metrics/mcp_weighted_total'] = ( mcp_total_loss_show) for depth, mcp_loss_show in enumerate( mcp_loss_shows): log_values[ f'loss_metrics/mcp_depth_{depth + 1}'] = ( mcp_loss_show) self.wandb.log(log_values, step=self.step) self.step += 1 if self.step % self.config.save_interval == 0: if self.config.rank == 0: logger.info(f"Starting save model at step {self.step}") self.save_checkpoint() if dist.is_initialized(): dist.barrier() progress_bar.close() logger.info("Training completed!") def run(args): """Main entry point.""" config = VA_CONFIGS[args.config_name] overrides = { 'wan22_pretrained_model_name_or_path': args.pretrained_model_path, 'dataset_path': args.dataset_path, 'empty_emb_path': args.empty_emb_path, 'learning_rate': args.learning_rate, 'cfg_prob': args.cfg_prob, 'init_worker': args.init_worker, 'load_worker': args.load_worker, 'batch_size': args.batch_size, 'gradient_accumulation_steps': args.gradient_accumulation_steps, 'num_steps': args.num_steps, 'save_interval': args.save_interval, 'save_root': args.save_root, } for key, value in overrides.items(): if value is not None: config[key] = value if args.dataset_path is not None and args.empty_emb_path is None: config.empty_emb_path = os.path.join(args.dataset_path, 'empty_emb.pt') if args.disable_wandb: config.enable_wandb = False rank = int(os.getenv("RANK", 0)) local_rank = int(os.environ.get('LOCAL_RANK', 0)) world_size = int(os.environ.get("WORLD_SIZE", 1)) init_distributed(world_size, local_rank, rank) config.rank = rank config.local_rank = local_rank config.world_size = world_size if rank == 0: logger.info(f"Using config: {args.config_name}") logger.info(f"World size: {world_size}, Local rank: {local_rank}") trainer = Trainer(config) trainer.train() def main(): """Parse arguments and run training.""" parser = argparse.ArgumentParser(description="Train WAN model for robotics") parser.add_argument( "--config-name", type=str, default='robotwin_train', help="Config name", ) parser.add_argument( "--save-root", type=str, default=None, help="Root directory for saving checkpoints", ) parser.add_argument( "--pretrained-model-path", type=str, default=None, help="Pretrained model root containing the transformer directory", ) parser.add_argument( "--dataset-path", type=str, default=None, help="Root directory containing open-format LeRobot datasets", ) parser.add_argument( "--empty-emb-path", type=str, default=None, help="Path to empty_emb.pt (defaults to DATASET_PATH/empty_emb.pt)", ) parser.add_argument( "--disable-wandb", action="store_true", help="Disable Weights & Biases logging", ) parser.add_argument("--learning-rate", type=float, default=None) parser.add_argument("--cfg-prob", type=float, default=None) parser.add_argument("--init-worker", type=int, default=None) parser.add_argument("--load-worker", type=int, default=None) parser.add_argument("--batch-size", type=int, default=None) parser.add_argument( "--gradient-accumulation-steps", type=int, default=None) parser.add_argument("--num-steps", type=int, default=None) parser.add_argument("--save-interval", type=int, default=None) args = parser.parse_args() run(args) if __name__ == "__main__": init_logger() main()