DecMem/trainer/diffusion.py

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