import gc import logging from model import CausalDiffusion from utils.dataset import ShardingLMDBDataset, WorldMemCSVDataset, cycle from utils.pose_utils import poses_to_viewmats from utils.misc import set_seed import torch.distributed as dist from omegaconf import OmegaConf import torch import wandb import time import os from utils.distributed import barrier, fsdp_wrap, fsdp_state_dict, fsdp_optim_state_dict, load_fsdp_optim_state_dict, launch_distributed_job # from utils.distributed import EMA_FSDP class Trainer: def __init__(self, config): self.config = config self.step = 0 # Step 1: Initialize the distributed training environment (rank, seed, dtype, logging etc.) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True launch_distributed_job() global_rank = dist.get_rank() self.dtype = torch.bfloat16 if config.mixed_precision else torch.float32 self.device = torch.cuda.current_device() self.is_main_process = global_rank == 0 self.causal = config.causal self.disable_wandb = config.disable_wandb # If resume_ckpt is set, load it now so we can restore step / seed / wandb_run_id # BEFORE seeding RNG and initializing wandb. The full state stays in CPU memory # and is consumed (then freed) after the optimizer / EMA are constructed. self._resume_state = None resume_path = getattr(config, "resume_ckpt", None) if resume_path: if global_rank == 0: print(f"Loading resume checkpoint: {resume_path}") self._resume_state = torch.load(resume_path, map_location="cpu", weights_only=False) self.step = int(self._resume_state.get("step", 0)) resumed_seed = self._resume_state.get("seed", None) if resumed_seed is not None: config.seed = int(resumed_seed) if global_rank == 0: print(f"Resumed: step={self.step}, seed={config.seed}, " f"wandb_run_id={self._resume_state.get('wandb_run_id')}") # use a random seed for the training if config.seed == 0: random_seed = torch.randint(0, 10000000, (1,), device=self.device) dist.broadcast(random_seed, src=0) config.seed = random_seed.item() set_seed(config.seed + global_rank) if self.is_main_process and not self.disable_wandb: wandb.login(host=config.wandb_host, key=config.wandb_key) wandb_init_kwargs = dict( config=OmegaConf.to_container(config, resolve=True), name=config.config_name, mode="online", entity=config.wandb_entity, project=config.wandb_project, dir=config.wandb_save_dir, ) wandb_run_id = self._resume_state.get("wandb_run_id") if self._resume_state else None if wandb_run_id: # resume="must" → wandb errors out if the id is missing, instead of # silently creating a new run that diverges from the original. wandb_init_kwargs["id"] = wandb_run_id wandb_init_kwargs["resume"] = "must" wandb.init(**wandb_init_kwargs) self.output_path = config.logdir # Step 2: Initialize the model and optimizer # For WorldMem training (CSV data path), skip text encoder to save GPU memory data_path = config.data_path self.use_worldmem_data = isinstance(data_path, str) and data_path.endswith(".csv") if self.use_worldmem_data: config.use_text_encoder = False self.model = CausalDiffusion(config, device=self.device) if global_rank == 0: def _fmt(n: int) -> str: for unit in ("", "K", "M", "B"): if abs(n) < 1000: return f"{n:.2f}{unit}" n /= 1000 return f"{n:.2f}T" total = sum(p.numel() for p in self.model.generator.parameters()) trainable = sum(p.numel() for p in self.model.generator.parameters() if p.requires_grad) print(f"[generator] params total={_fmt(total)} ({total:,}) " f"trainable={_fmt(trainable)} ({trainable:,})", flush=True) if self.model.text_encoder is not None: te = sum(p.numel() for p in self.model.text_encoder.parameters()) print(f"[text_encoder] params total={_fmt(te)} ({te:,})", flush=True) vae = sum(p.numel() for p in self.model.vae.parameters()) print(f"[vae] params total={_fmt(vae)} ({vae:,})", flush=True) self.model.generator = fsdp_wrap( self.model.generator, sharding_strategy=config.sharding_strategy, mixed_precision=config.mixed_precision, wrap_strategy=config.generator_fsdp_wrap_strategy ) if self.model.text_encoder is not None: self.model.text_encoder = fsdp_wrap( self.model.text_encoder, sharding_strategy=config.sharding_strategy, mixed_precision=config.mixed_precision, wrap_strategy=config.text_encoder_fsdp_wrap_strategy ) if not config.no_visualize or config.load_raw_video: self.model.vae = self.model.vae.to( device=self.device, dtype=torch.bfloat16 if config.mixed_precision else torch.float32) self.generator_optimizer = torch.optim.AdamW( [param for param in self.model.generator.parameters() if param.requires_grad], lr=config.lr, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay ) # Step 3: Initialize the dataloader if self.use_worldmem_data: num_latent_frames = getattr(config, "num_frames", 21) sample_strategy = getattr(config, "sample_strategy", "random") dataset = WorldMemCSVDataset( data_path, num_latent_frames=num_latent_frames, sample_strategy=sample_strategy, ) else: dataset = ShardingLMDBDataset(data_path, max_pair=int(1e8)) sampler = torch.utils.data.distributed.DistributedSampler( dataset, shuffle=True, drop_last=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=config.batch_size, sampler=sampler, num_workers=8, pin_memory=True, persistent_workers=True) if dist.get_rank() == 0: world_size = dist.get_world_size() per_rank_samples = len(sampler) # honors drop_last + world_size split global_batch_size = world_size * config.batch_size print("=" * 80, flush=True) print(f"[dataset] total samples : {len(dataset)}", flush=True) print(f"[dataset] per-rank : {per_rank_samples} " f"(drop_last=True, world_size={world_size})", flush=True) print(f"[batch] per-rank : {config.batch_size}", flush=True) print(f"[batch] global : {global_batch_size} " f"(= {world_size} ranks x {config.batch_size})", flush=True) print(f"[steps/epoch] : {per_rank_samples // config.batch_size}", flush=True) print("-" * 80, flush=True) print("[config]") print(OmegaConf.to_yaml(config, resolve=True), flush=True) print("=" * 80, flush=True) self.dataloader = cycle(dataloader) ############################################################################################################## # 6. Set up EMA parameter containers rename_param = ( lambda name: name.replace("_fsdp_wrapped_module.", "") .replace("_checkpoint_wrapped_module.", "") .replace("_orig_mod.", "") ) self.name_to_trainable_params = {} for n, p in self.model.generator.named_parameters(): if not p.requires_grad: continue renamed_n = rename_param(n) self.name_to_trainable_params[renamed_n] = p # ema_weight = config.ema_weight # self.generator_ema = None # if (ema_weight is not None) and (ema_weight > 0.0): # print(f"Setting up EMA with weight {ema_weight}") # self.generator_ema = EMA_FSDP(self.model.generator, decay=ema_weight) ############################################################################################################## # 7. (If resuming) Load the model and optimizer, lr_scheduler statedicts if self._resume_state is not None: # Full resume: weights + optimizer + EMA (step already restored in __init__ top). if global_rank == 0: print("Restoring generator / optimizer from resume checkpoint.") # Use FULL_STATE_DICT context so FSDP scatters the (CPU) full # state-dict into local GPU shards instead of leaving params on CPU. from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import StateDictType with FSDP.state_dict_type(self.model.generator, StateDictType.FULL_STATE_DICT): self.model.generator.load_state_dict(self._resume_state["generator"], strict=False) if "optimizer" in self._resume_state: load_fsdp_optim_state_dict( self.model.generator, self.generator_optimizer, self._resume_state["optimizer"], ) # if "generator_ema" in self._resume_state and self.generator_ema is not None: # self.generator_ema.load_state_dict(self._resume_state["generator_ema"]) self._resume_state = None torch.cuda.empty_cache() elif getattr(config, "generator_ckpt", False): # Warm-start: weights only, step counter and optimizer stay fresh. print(f"Loading pretrained generator from {config.generator_ckpt}") state_dict = torch.load(config.generator_ckpt, map_location="cpu") if "generator" in state_dict: state_dict = state_dict["generator"] elif "model" in state_dict: state_dict = state_dict["model"] # strict=False: dense ckpts lack LTM-only params (e.g. mem_proj); keep init values. missing, unexpected = self.model.generator.load_state_dict( state_dict, strict=False ) if missing: print( f" warm-start missing keys ({len(missing)}): " f"{missing[:5]}{'...' if len(missing) > 5 else ''}" ) if unexpected: print( f" warm-start unexpected keys ({len(unexpected)}): " f"{unexpected[:5]}{'...' if len(unexpected) > 5 else ''}" ) ############################################################################################################## # # Let's delete EMA params for early steps to save some computes at training and inference # if self.step < config.ema_start_step: # self.generator_ema = None self.max_grad_norm = 10.0 self.previous_time = None def save(self): print("Start gathering distributed model states...") # Both calls are collectives: every rank must participate, only rank 0 ends up # holding the full CPU tensors (rank0_only=True inside the helpers). generator_state_dict = fsdp_state_dict(self.model.generator) optimizer_state_dict = fsdp_optim_state_dict( self.model.generator, self.generator_optimizer) wandb_run_id = ( wandb.run.id if (self.is_main_process and not self.disable_wandb and wandb.run is not None) else None ) state_dict = { "generator": generator_state_dict, "optimizer": optimizer_state_dict, "step": self.step, "seed": int(self.config.seed), "wandb_run_id": wandb_run_id, } # if self.step > self.config.ema_start_step and self.generator_ema is not None: # state_dict["generator_ema"] = self.generator_ema.state_dict() if self.is_main_process: os.makedirs(os.path.join(self.output_path, f"checkpoint_model_{self.step:06d}"), exist_ok=True) torch.save(state_dict, os.path.join(self.output_path, f"checkpoint_model_{self.step:06d}", "model.pt")) print("Model saved to", os.path.join(self.output_path, f"checkpoint_model_{self.step:06d}", "model.pt")) def train_one_step(self, batch): torch.cuda.synchronize() t1 = time.time() self.log_iters = 1 if self.step % 20 == 0: torch.cuda.empty_cache() # Step 1: Get the next batch of text prompts text_prompts = batch["prompts"] if not self.config.load_raw_video: # precomputed latent clean_latent = batch["ode_latent"][:, -1].to( device=self.device, dtype=self.dtype) else: # encode raw video to latent frames = batch["frames"].to( device=self.device, dtype=self.dtype) with torch.no_grad(): clean_latent = self.model.vae.encode_to_latent( frames).to(device=self.device, dtype=self.dtype) image_latent = clean_latent[:, 0:1, ] batch_size = len(text_prompts) image_or_video_shape = list(self.config.image_or_video_shape) image_or_video_shape[0] = batch_size # Step 2: Build conditional_dict proportion_empty_prompts = getattr(self.config, "proportion_empty_prompts", 0.0) if self.use_worldmem_data: # WorldMem: no text conditioning; only actions and viewmats conditional_dict = {} unconditional_dict = {} actions = batch["actions"].to(device=self.device, dtype=self.dtype) # (B, T, 25) poses = batch["poses"].to(device=self.device, dtype=self.dtype) # (B, T, 5) viewmats = poses_to_viewmats(poses) # (B, T, 4, 4) # CFG: randomly drop action conditioning if proportion_empty_prompts > 0.0: drop_mask = torch.rand(batch_size, device=self.device) < proportion_empty_prompts actions = actions.clone() # viewmats = viewmats.clone() actions[drop_mask] = 0.0 # eye = torch.eye(4, dtype=viewmats.dtype, device=viewmats.device) # viewmats[drop_mask] = eye.unsqueeze(0).unsqueeze(0).expand( # drop_mask.sum(), viewmats.shape[1], -1, -1) conditional_dict["actions"] = actions conditional_dict["viewmats"] = viewmats else: with torch.no_grad(): conditional_dict = self.model.text_encoder(text_prompts=text_prompts) if not getattr(self, "unconditional_dict", None): unconditional_dict = self.model.text_encoder( text_prompts=[self.config.negative_prompt] * batch_size) unconditional_dict = {k: v.detach() for k, v in unconditional_dict.items()} self.unconditional_dict = unconditional_dict else: unconditional_dict = self.unconditional_dict # Step 3: Train the generator generator_loss, log_dict = self.model.generator_loss( image_or_video_shape=image_or_video_shape, conditional_dict=conditional_dict, unconditional_dict=unconditional_dict, clean_latent=clean_latent, initial_latent=image_latent ) self.generator_optimizer.zero_grad() generator_loss.backward() generator_grad_norm = self.model.generator.clip_grad_norm_( self.max_grad_norm) self.generator_optimizer.step() # Increment the step since we finished gradient update self.step += 1 loss_unweighted_val = log_dict["loss_unweighted"].item() wandb_loss_dict = { "generator_loss": generator_loss.item(), # weighted, used for backprop "generator_loss_unweighted": loss_unweighted_val, # raw MSE, comparable across runs "generator_grad_norm": generator_grad_norm.item(), } # Per-frame unweighted loss bucketed by frame position (10 frames per bucket) bucket_str = "" loss_per_frame = log_dict.get("loss_per_frame", None) if loss_per_frame is not None: per_frame = loss_per_frame.float().mean(dim=0) # (F,) num_frame = per_frame.shape[0] bucket_size = 10 bucket_means = [] for start in range(0, num_frame, bucket_size): end = min(start + bucket_size, num_frame) v = per_frame[start:end].mean().item() bucket_means.append((start, end, v)) wandb_loss_dict[f"loss_frames_{start:02d}-{end - 1:02d}"] = v bucket_str = " " + " ".join( f"f{start:02d}-{end - 1:02d}={v:.4f}" for start, end, v in bucket_means ) # Step 4: Logging torch.cuda.synchronize() t2 = time.time() if self.is_main_process: print( f"step={self.step} loss={generator_loss.item():.4f} " f"loss_unweighted={loss_unweighted_val:.4f} " f"grad_norm={generator_grad_norm.item():.4f}", f"time={t2 - t1:.2f} seconds", flush=True, ) if not self.disable_wandb: wandb.log(wandb_loss_dict, step=self.step) if self.step % self.config.gc_interval == 0: if dist.get_rank() == 0: logging.info("DistGarbageCollector: Running GC.") gc.collect() def generate_video(self, pipeline, prompts, image=None): batch_size = len(prompts) sampled_noise = torch.randn( [batch_size, 21, 16, 60, 104], device="cuda", dtype=self.dtype ) video, _ = pipeline.inference( noise=sampled_noise, text_prompts=prompts, return_latents=True ) current_video = video.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0 return current_video def train(self): while True: batch = next(self.dataloader) self.train_one_step(batch) if (not self.config.no_save) and self.step % self.config.log_iters == 0: torch.cuda.empty_cache() self.save() torch.cuda.empty_cache() barrier() if self.is_main_process: current_time = time.time() if self.previous_time is None: self.previous_time = current_time else: if not self.disable_wandb: wandb.log({"per iteration time": current_time - self.previous_time}, step=self.step) self.previous_time = current_time