import os import sys import yaml import argparse import wandb import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler, Subset from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from tqdm import tqdm import time import numpy as np import signal # Global variables for signal handling _model = None _optimizer = None _step = 0 _epoch = 0 _ckpt_dir = "" _wandb_run_id = None def signal_handler(sig, frame): """Save checkpoint on SIGTERM (Slurm timeout/preemption).""" global _model, _optimizer, _step, _epoch, _ckpt_dir, _wandb_run_id if _model is not None and _ckpt_dir: rank = 0 if dist.is_initialized(): rank = dist.get_rank() if rank == 0: print(f"\n[SIGNAL {sig}] Saving emergency checkpoint at step {_step}...") ckpt_path = os.path.join(_ckpt_dir, f"checkpoint_signal_{_step}.pt") save_checkpoint(_model, _optimizer, _step, _epoch, ckpt_path, wandb_run_id=_wandb_run_id) print(f"--- Emergency Checkpoint Saved: {ckpt_path} ---") wandb.finish() sys.exit(0) # Register signal handler signal.signal(signal.SIGTERM, signal_handler) # Add project root to path sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) from wm.model.interface import get_dynamics_class from wm.dataset.dataset import RoboticsDatasetWrapper from wm.utils.visualization import visualize_layout def setup_ddp(): if 'RANK' in os.environ: dist.init_process_group("nccl") rank = int(os.environ['RANK']) local_rank = int(os.environ['LOCAL_RANK']) world_size = int(os.environ['WORLD_SIZE']) torch.cuda.set_device(local_rank) return rank, local_rank, world_size else: return 0, 0, 1 def cleanup_ddp(): if dist.is_initialized(): dist.destroy_process_group() def save_checkpoint(model, optimizer, step, epoch, path, wandb_run_id=None, save_numbered=True): checkpoint = { 'model_state_dict': model.module.state_dict() if hasattr(model, 'module') else model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'step': step, 'epoch': epoch, 'wandb_run_id': wandb_run_id } if save_numbered and path: torch.save(checkpoint, path) # Also save a 'latest.pt' for easy resuming ckpt_dir = os.path.dirname(path) if path else _ckpt_dir latest_path = os.path.join(ckpt_dir, "latest.pt") torch.save(checkpoint, latest_path) def load_checkpoint(model, optimizer, path, device): if not os.path.exists(path): return 0, 0, None print(f"--- Loading Checkpoint from {path} ---") checkpoint = torch.load(path, map_location=device, weights_only=False) # Handle DDP/FSDP vs single GPU state_dict = checkpoint['model_state_dict'] # Filter out scheduler buffers that might cause size mismatches # These are recalculated anyway, so they shouldn't be in the state_dict scheduler_buffers = [ 'scheduler.sigmas', 'scheduler.timesteps', 'scheduler.linear_timesteps_weights' ] for k in scheduler_buffers: if k in state_dict: del state_dict[k] if f"module.{k}" in state_dict: del state_dict[f"module.{k}"] if hasattr(model, 'module'): model.module.load_state_dict(state_dict, strict=False) else: # If loading a DDP state dict into a non-DDP model, strip 'module.' if any(k.startswith('module.') for k in state_dict.keys()): state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} model.load_state_dict(state_dict, strict=False) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) return checkpoint['step'], checkpoint['epoch'], checkpoint.get('wandb_run_id') def log_videos_to_wandb(model, val_loader, device, step, dataset_name, gen_mode="parallel", num_inference_steps=50): model.eval() video_logs = [] with torch.no_grad(): try: # Just take the first batch batch = next(iter(val_loader)) except StopIteration: return obs = batch['obs'].to(device) # [B, T, C, H, W] action = batch['action'].to(device) # [B, T, A] # Use generate to predict video # obs[:, 0] is the first frame # (B, H, W, C) -> permute back from (B, C, H, W) o_0 = obs[:, 0].permute(0, 2, 3, 1).contiguous() # Pass gen_mode and num_inference_steps to generate if it's DiffusionForcing_WM if hasattr(model, 'module'): curr_model = model.module else: curr_model = model try: pred_video = curr_model.generate(o_0, action, mode=gen_mode, num_inference_steps=num_inference_steps) except TypeError: # Fallback for models that don't support mode/steps pred_video = curr_model.generate(o_0, action) # pred_video: [B, T, H, W, 3] in [0, 1] for b in range(min(obs.shape[0], 8)): # Log up to 8 samples # Visualize layout on GT gt_with_layout = visualize_layout(obs[b].cpu().numpy(), action[b].cpu().numpy(), dataset_name) # Visualize layout on Pred # pred_video[b] is [T, H, W, 3], visualize_layout expects [T, C, H, W] pred_obs_b = pred_video[b].permute(0, 3, 1, 2).cpu().numpy() pred_with_layout = visualize_layout(pred_obs_b, action[b].cpu().numpy(), dataset_name) # Combine GT and Pred side by side for wandb # gt and pred are [T, H, W, 3] combined = np.concatenate([gt_with_layout, pred_with_layout], axis=2) # [T, H, 2W, 3] combined = combined.transpose(0, 3, 1, 2) # [T, 3, H, 2W] video_logs.append(wandb.Video(combined, fps=10, format="mp4", caption=f"Step {step} - Sample {b} (GT vs Pred)")) if video_logs: wandb.log({"val/videos": video_logs}, step=step) model.train() def evaluate_mse(model, val_loader, device, step, num_batches=1, gen_mode="parallel", num_inference_steps=50): model.eval() all_mse = [] with torch.no_grad(): for i, batch in enumerate(val_loader): if i >= num_batches: break obs = batch['obs'].to(device) # [B, T, C, H, W] action = batch['action'].to(device) # [B, T, A] o_0 = obs[:, 0].permute(0, 2, 3, 1).contiguous() # Use generate for rollout if hasattr(model, 'module'): curr_model = model.module else: curr_model = model try: pred_video = curr_model.generate(o_0, action, mode=gen_mode, num_inference_steps=num_inference_steps) except TypeError: pred_video = curr_model.generate(o_0, action) # pred_video: [B, T, H, W, 3], in [0, 1] gt_video = obs.permute(0, 1, 3, 4, 2).contiguous() # [B, T, H, W, 3] mse = torch.mean((pred_video - gt_video) ** 2) all_mse.append(mse.item()) avg_mse = sum(all_mse) / len(all_mse) if all_mse else 0 wandb.log({"val/mse_rollout": avg_mse}, step=step) model.train() return avg_mse def main(): global _model, _optimizer, _step, _epoch, _ckpt_dir, _wandb_run_id parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, required=True, help="Path to yaml config") parser.add_argument("--resume", action="store_true", help="Resume from latest checkpoint in wandb dir") parser.add_argument("--ckpt_path", type=str, default=None, help="Explicit path to checkpoint to resume from") args = parser.parse_args() # Load config with open(args.config, 'r') as f: config = yaml.safe_load(f) rank, local_rank, world_size = setup_ddp() device = torch.device(f"cuda:{local_rank}") if torch.cuda.is_available() else torch.device("cpu") # Setup Checkpoint Directory ckpt_dir = os.path.join("checkpoints", config['wandb']['run_name']) os.makedirs(ckpt_dir, exist_ok=True) _ckpt_dir = ckpt_dir # 1. Initialize model dynamics_class_name = config['dynamics_class'] model_name = config['model_name'] model_config = config['model_config'] if rank == 0: print(f"--- Initializing Dynamics Model: {dynamics_class_name} ({model_name}) ---") dynamics_class = get_dynamics_class(dynamics_class_name) dynamics_model = dynamics_class(model_name, model_config).to(device) # 2. Optimizer (needs to be created before loading checkpoint) optimizer = torch.optim.AdamW(dynamics_model.parameters(), lr=float(config['training']['learning_rate'])) _optimizer = optimizer # 3. Resume Logic start_step = 0 start_epoch = 0 wandb_run_id = None resume_path = args.ckpt_path if args.resume and resume_path is None: potential_latest = os.path.join(ckpt_dir, "latest.pt") if os.path.exists(potential_latest): resume_path = potential_latest if resume_path: start_step, start_epoch, wandb_run_id = load_checkpoint(dynamics_model, optimizer, resume_path, device) _wandb_run_id = wandb_run_id # Distributed wrapper if config['distributed']['use_fsdp']: model = FSDP(dynamics_model) elif world_size > 1: model = DDP(dynamics_model, device_ids=[local_rank], find_unused_parameters=True) else: model = dynamics_model _model = model if rank == 0: params = sum(p.numel() for p in dynamics_model.model.parameters() if p.requires_grad) print(f"Model Parameters: {params / 1e6:.2f}M") print(f"--- Distributed Setup Finished ---") print(f"World Size: {world_size}") print(f"Device: {device}") # 4. Initialize WandB if rank == 0: if config['wandb']['api_key'] != "YOUR_WANDB_API_KEY_HERE": os.environ["WANDB_API_KEY"] = config['wandb']['api_key'] wandb.init( project=config['wandb']['project'], name=config['wandb']['run_name'], config=config, id=wandb_run_id, resume="allow" ) wandb_run_id = wandb.run.id _wandb_run_id = wandb_run_id # Dataset and Dataloader dataset_name = config['dataset']['name'] if rank == 0: print(f"--- Loading Dataset: {dataset_name} ---") # Support separate train/eval sequence lengths train_seq_len = config['dataset'].get('train_seq_len', config['dataset'].get('seq_len')) eval_seq_len = config['dataset'].get('eval_seq_len', config['dataset'].get('seq_len')) gen_mode = config['training'].get('gen_mode', 'parallel') inference_steps = config['training'].get('inference_steps', 50) # We initialize two datasets to handle different sequence lengths train_dataset_full = RoboticsDatasetWrapper.get_dataset(dataset_name, seq_len=train_seq_len) val_dataset_full = RoboticsDatasetWrapper.get_dataset(dataset_name, seq_len=eval_seq_len) # Train/Test split (Trajectory-based to avoid data leakage) # Use trajectory indices from train_dataset_full (should be same as val_dataset_full) unique_traj_ids = sorted(list(set([idx[0] for idx in train_dataset_full.indices]))) num_total_trajs = len(unique_traj_ids) split_ratio = config['dataset'].get('train_test_split', 10) num_val_trajs = max(1, num_total_trajs // (split_ratio + 1)) # Fixed split for reproducibility import random random.seed(42) random.shuffle(unique_traj_ids) val_traj_ids = set(unique_traj_ids[:num_val_trajs]) train_indices = [] val_indices = [] # Map indices back to the respective datasets for i, (traj_idx, start_f) in enumerate(train_dataset_full.indices): if traj_idx not in val_traj_ids: train_indices.append(i) for i, (traj_idx, start_f) in enumerate(val_dataset_full.indices): if traj_idx in val_traj_ids: val_indices.append(i) if rank == 0: print(f"Split: {len(train_indices)} train windows (T={train_seq_len}), " f"{len(val_indices)} val windows (T={eval_seq_len})") train_dataset = Subset(train_dataset_full, train_indices) val_dataset = Subset(val_dataset_full, val_indices) train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank) if world_size > 1 else None train_loader = DataLoader( train_dataset, batch_size=config['training']['batch_size'], sampler=train_sampler, shuffle=(train_sampler is None), num_workers=config['training']['num_workers'], pin_memory=True ) # Validation loader with fixed shuffle for diverse but consistent visualization val_g = torch.Generator() val_g.manual_seed(42) val_loader = DataLoader( val_dataset, batch_size=config['training']['batch_size'], shuffle=True, num_workers=config['training']['num_workers'], generator=val_g ) # Training Loop num_epochs = config['training']['num_epochs'] step = start_step _step = step _epoch = start_epoch if rank == 0: print(f"--- Starting Training Loop: {num_epochs} Epochs (Starting from Epoch {start_epoch}, Step {start_step}) ---") for epoch in range(start_epoch, num_epochs): _epoch = epoch if train_sampler: train_sampler.set_epoch(epoch) model.train() pbar = tqdm(train_loader, desc=f"Epoch {epoch}", disable=(rank != 0)) last_step_end = time.time() for batch in pbar: # (1) Load data time data_time = time.time() - last_step_end obs = batch['obs'].to(device) # [B, T, C, H, W] action = batch['action'].to(device) # [B, T, A] optimizer.zero_grad() # Synchronize for accurate timing if torch.cuda.is_available(): torch.cuda.synchronize() # (2) VAE Encoding Time t_enc_start = time.time() # We call encode_obs separately to time it with torch.no_grad(): z = model.module.encode_obs(obs) if hasattr(model, 'module') else model.encode_obs(obs) if torch.cuda.is_available(): torch.cuda.synchronize() vae_time = time.time() - t_enc_start # (3) Forward and Update Time t_update_start = time.time() # training_loss handles the DiT forward pass loss = model.module.training_loss(z, action) if hasattr(model, 'module') else model.training_loss(z, action) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), config['training']['grad_clip']) optimizer.step() if torch.cuda.is_available(): torch.cuda.synchronize() update_time = time.time() - t_update_start step += 1 # Logging if rank == 0: step_time = time.time() - last_step_end pbar.set_postfix({ "loss": f"{loss.item():.4f}", "dt": f"{data_time:.2f}s", "vae": f"{vae_time:.2f}s", "up": f"{update_time:.2f}s", "st": f"{step_time:.2f}s" }) if step % config['training']['log_freq'] == 0: print(f"Step {step} (Epoch {epoch}) | Loss: {loss.item():.4f} | " f"Data: {data_time:.3f}s | VAE: {vae_time:.3f}s | " f"Update: {update_time:.3f}s | Step: {step_time:.3f}s") wandb.log({ "train/loss": loss.item(), "train/epoch": epoch, "time/data_loading": data_time, "time/vae_encoding": vae_time, "time/model_update": update_time, "time/seconds_per_step": step_time, }, step=step) # Validation MSE (Rollout) - More frequent eval_freq = config['training'].get('eval_freq', 50) if step % eval_freq == 0: print(f"\n--- Calculating Val MSE at Step {step} ---") evaluate_mse(model, val_loader, device, step, num_batches=2, gen_mode=gen_mode, num_inference_steps=inference_steps) # Video Logging - Less frequent if step % config['training']['val_freq'] == 0: print(f"\n--- Logging Validation Videos at Step {step} ---") log_videos_to_wandb(model, val_loader, device, step, dataset_name, gen_mode=gen_mode, num_inference_steps=inference_steps) print(f"--- Video Logging Finished ---") # Checkpoint # Save numbered checkpoint every 2000 steps ckpt_freq = config['training'].get('checkpoint_freq', 2000) latest_freq = config['training'].get('latest_freq', 500) if step % ckpt_freq == 0: ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{step}.pt") save_checkpoint(model, optimizer, step, epoch, ckpt_path, wandb_run_id=wandb_run_id, save_numbered=True) print(f"--- Numbered Checkpoint Saved: {ckpt_path} ---") # Save latest.pt every 500 steps (if not already saved by numbered checkpoint) elif step % latest_freq == 0: save_checkpoint(model, optimizer, step, epoch, None, wandb_run_id=wandb_run_id, save_numbered=False) print(f"--- Latest Checkpoint Updated (Step {step}) ---") _step = step last_step_end = time.time() dist.barrier() if world_size > 1 else None if rank == 0: wandb.finish() cleanup_ddp() if __name__ == "__main__": main()