trainer/predictor_v4_dmd.py

"""Wan14B DMD training for Predictor-v4, with optional Full-Generator tuning."""

from __future__ import annotations

import gc
import json
import math
import os
import random
import shutil
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Iterator

import torch
import torch.distributed as dist
from omegaconf import OmegaConf
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

from model.dmd import DMD
from model.predictor_v4 import SelfForcingPredictorV4
from pipeline.predictor_v4_dmd_training import PredictorV4DMDTrainingPipeline
from predictor_training.checkpoint import (
    atomic_torch_save,
    save_predictor_weights,
    unwrap_model,
)
from trainer.predictor_v4_rollout import (
    EXPECTED_TIMESTEPS,
    _configure_predictor_precision,
    _load_stage1_model_state,
    _optimizer_groups,
)
from utils.dataset import TextDataset
from utils.distributed import fsdp_state_dict, fsdp_wrap, launch_distributed_job
from utils.misc import set_seed


def _cosine_with_linear_warmup(step: int, warmup: int, total: int) -> float:
    if step < warmup:
        return max(1e-8, float(step + 1) / max(1, warmup))
    progress = min(1.0, float(step - warmup) / max(1, total - warmup))
    return 0.5 * (1.0 + math.cos(math.pi * progress))


def _make_scheduler(
    optimizer: torch.optim.Optimizer,
    *,
    warmup_steps: int,
    max_steps: int,
) -> LambdaLR:
    return LambdaLR(
        optimizer,
        lambda step: _cosine_with_linear_warmup(
            step, warmup_steps, max_steps
        ),
    )


class _LocalEMA:
    """EMA over the rank-local parameter views, including FSDP shards."""

    def __init__(self, module: torch.nn.Module, decay: float) -> None:
        self.decay = float(decay)
        self.parameters = [
            parameter for parameter in module.parameters()
            if parameter.requires_grad
        ]
        self.shadow = [
            parameter.detach().float().clone()
            for parameter in self.parameters
        ]

    @torch.no_grad()
    def update(self) -> None:
        for shadow, parameter in zip(self.shadow, self.parameters):
            shadow.mul_(self.decay).add_(
                parameter.detach().float(), alpha=1.0 - self.decay
            )

    @contextmanager
    def apply(self):
        with torch.no_grad():
            backups = [
                parameter.detach().clone() for parameter in self.parameters
            ]
            try:
                for parameter, shadow in zip(self.parameters, self.shadow):
                    parameter.copy_(shadow.to(dtype=parameter.dtype))
                yield
            finally:
                for parameter, backup in zip(self.parameters, backups):
                    parameter.copy_(backup)


def _distributed_mean(values: torch.Tensor) -> torch.Tensor:
    dist.all_reduce(values, op=dist.ReduceOp.SUM)
    return values / dist.get_world_size()


def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
    temporary.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True, default=str)
        + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


class Trainer:
    """Distributed Predictor-only or joint Full+Predictor DMD trainer."""

    def __init__(self, config) -> None:
        self.root_config = config
        self.cfg = config.predictor_v4_dmd
        self.mode = str(self.cfg.training_mode).lower()
        if self.mode not in {"predictor_only", "joint"}:
            raise ValueError(f"Unknown training_mode {self.mode!r}")

        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("high")
        launch_distributed_job()
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        self.local_rank = int(os.environ["LOCAL_RANK"])
        self.device = torch.device("cuda", self.local_rank)
        self.is_main = self.rank == 0
        if self.world_size != int(self.cfg.expected_world_size):
            raise ValueError(
                f"Expected {self.cfg.expected_world_size} ranks, got {self.world_size}"
            )

        seed = int(config.seed)
        set_seed(seed + self.rank)
        random.seed(seed + self.rank)
        self.output_dir = Path(str(self.cfg.output_dir)).resolve()
        if self.is_main:
            self.output_dir.mkdir(parents=True, exist_ok=True)
        dist.barrier()
        self.log_path = self.output_dir / "train_log.jsonl"

        self.model = DMD(config, device=self.device)
        checkpoint = torch.load(
            str(self.cfg.full_checkpoint),
            map_location="cpu",
            mmap=True,
            weights_only=False,
        )
        checkpoint_key = str(self.cfg.full_checkpoint_key)
        if checkpoint_key not in checkpoint:
            raise KeyError(
                f"{self.cfg.full_checkpoint} lacks {checkpoint_key!r}"
            )
        self.model.generator.load_state_dict(
            checkpoint[checkpoint_key],
            strict=bool(self.cfg.strict_full_load),
        )
        del checkpoint

        predictor = SelfForcingPredictorV4.from_teacher(
            self.model.generator.model,
            source_block_ids=tuple(
                int(value) for value in self.cfg.source_block_ids
            ),
        )
        if self.mode == "predictor_only":
            self.model.generator.requires_grad_(False)
        else:
            self.model.generator.requires_grad_(True)

        self.model.generator = fsdp_wrap(
            self.model.generator,
            sharding_strategy=str(config.sharding_strategy),
            mixed_precision=bool(config.mixed_precision),
            wrap_strategy=str(config.generator_fsdp_wrap_strategy),
        )
        self.model.real_score = fsdp_wrap(
            self.model.real_score,
            sharding_strategy=str(config.sharding_strategy),
            mixed_precision=bool(config.mixed_precision),
            wrap_strategy=str(config.real_score_fsdp_wrap_strategy),
        )
        self.model.fake_score = fsdp_wrap(
            self.model.fake_score,
            sharding_strategy=str(config.sharding_strategy),
            mixed_precision=bool(config.mixed_precision),
            wrap_strategy=str(config.fake_score_fsdp_wrap_strategy),
        )
        self.model.text_encoder = fsdp_wrap(
            self.model.text_encoder,
            sharding_strategy=str(config.sharding_strategy),
            mixed_precision=bool(config.mixed_precision),
            wrap_strategy=str(config.text_encoder_fsdp_wrap_strategy),
            cpu_offload=bool(getattr(config, "text_encoder_cpu_offload", False)),
        )

        _configure_predictor_precision(predictor, device=self.device)
        _load_stage1_model_state(
            predictor,
            Path(str(self.cfg.stage1_training_state)).resolve(),
            expected_step=int(self.cfg.stage1_expected_step),
        )
        predictor.train()
        self.predictor = DDP(
            predictor,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            broadcast_buffers=False,
            gradient_as_bucket_view=True,
            find_unused_parameters=False,
        )

        actual_timesteps = self.model.denoising_step_list.detach().float().cpu()
        if not torch.equal(actual_timesteps, EXPECTED_TIMESTEPS):
            raise ValueError(
                "Predictor DMD requires exact warped timesteps "
                f"{EXPECTED_TIMESTEPS.tolist()}, got {actual_timesteps.tolist()}"
            )
        forced_exit = getattr(self.cfg, "forced_exit_step", None)
        self.rollout = PredictorV4DMDTrainingPipeline(
            denoising_step_list=self.model.denoising_step_list,
            scheduler=self.model.scheduler,
            generator=self.model.generator,
            predictor=self.predictor,
            training_mode=self.mode,
            context_noise=int(config.context_noise),
            forced_exit_step=(
                None if forced_exit is None else int(forced_exit)
            ),
        )
        self.model.inference_pipeline = self.rollout

        fusion, blocks, group_names = _optimizer_groups(
            unwrap_model(self.predictor)
        )
        self.predictor_optimizer = AdamW(
            [
                {
                    "params": fusion,
                    "lr": float(self.cfg.predictor_fusion_lr),
                    "name": "fusion_residual",
                },
                {
                    "params": blocks,
                    "lr": float(self.cfg.predictor_blocks_lr),
                    "name": "blocks",
                },
            ],
            betas=(
                float(self.cfg.student_beta1),
                float(self.cfg.student_beta2),
            ),
            weight_decay=float(self.cfg.weight_decay),
        )
        self.predictor_scheduler = _make_scheduler(
            self.predictor_optimizer,
            warmup_steps=int(self.cfg.warmup_steps),
            max_steps=int(self.cfg.target_predictor_updates),
        )
        self.full_optimizer = None
        self.full_scheduler = None
        if self.mode == "joint":
            self.full_optimizer = AdamW(
                [
                    parameter
                    for parameter in self.model.generator.parameters()
                    if parameter.requires_grad
                ],
                lr=float(self.cfg.full_lr),
                betas=(
                    float(self.cfg.student_beta1),
                    float(self.cfg.student_beta2),
                ),
                weight_decay=float(self.cfg.weight_decay),
            )
            self.full_scheduler = _make_scheduler(
                self.full_optimizer,
                warmup_steps=int(self.cfg.warmup_steps),
                max_steps=int(self.cfg.max_student_steps),
            )
        self.critic_optimizer = AdamW(
            [
                parameter
                for parameter in self.model.fake_score.parameters()
                if parameter.requires_grad
            ],
            lr=float(self.cfg.fake_score_lr),
            betas=(
                float(self.cfg.critic_beta1),
                float(self.cfg.critic_beta2),
            ),
            weight_decay=float(self.cfg.weight_decay),
        )

        dataset = TextDataset(prompt_path=str(self.cfg.data_path))
        self.sampler = DistributedSampler(
            dataset,
            num_replicas=self.world_size,
            rank=self.rank,
            shuffle=True,
            seed=seed,
            drop_last=True,
        )
        self.loader = DataLoader(
            dataset,
            batch_size=int(config.batch_size),
            sampler=self.sampler,
            num_workers=int(self.cfg.num_workers),
            pin_memory=bool(self.cfg.pin_memory),
            drop_last=True,
        )
        self.data_iterator: Iterator[dict[str, Any]] | None = None
        self.data_epoch = 0

        self.student_step = 0
        self.predictor_step = 0
        self.critic_step = 0
        self.predictor_ema: _LocalEMA | None = None
        self.full_ema: _LocalEMA | None = None
        self.unconditional_dict: dict[str, torch.Tensor] | None = None
        self.run_config = {
            **OmegaConf.to_container(self.cfg, resolve=True),
            "world_size": self.world_size,
            "effective_global_batch": int(config.batch_size) * self.world_size,
            "parameter_groups": group_names,
            "rollout": (
                "random_exit_P1_P2_P3"
                if self.mode == "predictor_only"
                else "random_exit_F0_P1_P2_P3"
            ),
            "dmd_frames": 21,
            "clean_context_grad": False,
        }
        if self.is_main:
            _atomic_json(self.output_dir / "train_config.json", self.run_config)
        self.swanlab_run = self._initialize_swanlab()

    def _initialize_swanlab(self):
        if not self.is_main or not bool(self.cfg.use_swanlab):
            return None
        import swanlab

        mode = str(self.cfg.swanlab_mode)
        if mode == "cloud":
            api_key = os.environ.get("SWANLAB_API_KEY")
            if api_key:
                swanlab.login(api_key=api_key, save=False)
            else:
                swanlab.login()
        workspace = self.cfg.swanlab_workspace
        return swanlab.init(
            project=str(self.cfg.swanlab_project),
            workspace=None if workspace is None else str(workspace),
            experiment_name=str(self.cfg.swanlab_experiment),
            description=str(self.cfg.swanlab_description),
            tags=list(self.cfg.swanlab_tags),
            config=self.run_config,
            logdir=str(self.output_dir / "swanlab"),
            mode=mode,
        )

    def _next_batch(self) -> dict[str, Any]:
        if self.data_iterator is None:
            self.sampler.set_epoch(self.data_epoch)
            self.data_iterator = iter(self.loader)
        try:
            return next(self.data_iterator)
        except StopIteration:
            self.data_epoch += 1
            self.sampler.set_epoch(self.data_epoch)
            self.data_iterator = iter(self.loader)
            return next(self.data_iterator)

    @torch.no_grad()
    def _conditional_dicts(
        self, prompts: list[str]
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        conditional = self.model.text_encoder(text_prompts=prompts)
        if self.unconditional_dict is None:
            unconditional = self.model.text_encoder(
                text_prompts=[str(self.root_config.negative_prompt)]
                * len(prompts)
            )
            self.unconditional_dict = {
                key: value.detach() for key, value in unconditional.items()
            }
        return conditional, self.unconditional_dict

    def _image_shape(self, batch_size: int) -> list[int]:
        shape = list(self.root_config.image_or_video_shape)
        shape[0] = int(batch_size)
        return shape

    def _critic_update(self) -> dict[str, float]:
        batch = self._next_batch()
        prompts = list(batch["prompts"])
        conditional, unconditional = self._conditional_dicts(prompts)
        self.critic_optimizer.zero_grad(set_to_none=True)
        loss, logs = self.model.critic_loss(
            image_or_video_shape=self._image_shape(len(prompts)),
            conditional_dict=conditional,
            unconditional_dict=unconditional,
            clean_latent=None,
            initial_latent=None,
        )
        loss.backward()
        grad_norm = self.model.fake_score.clip_grad_norm_(
            float(self.cfg.critic_grad_clip)
        )
        if not torch.isfinite(grad_norm):
            raise FloatingPointError(f"Non-finite Fake Score grad norm {grad_norm}")
        self.critic_optimizer.step()
        self.critic_step += 1
        result = {
            "critic_loss": float(loss.detach()),
            "critic_grad_norm": float(grad_norm),
        }
        del batch, conditional, loss, logs
        return result

    def _maybe_initialize_ema(self) -> None:
        if (
            self.predictor_ema is None
            and self.predictor_step >= int(self.cfg.ema_start_step)
        ):
            self.predictor_ema = _LocalEMA(
                unwrap_model(self.predictor), float(self.cfg.ema_decay)
            )
        if (
            self.mode == "joint"
            and self.full_ema is None
            and self.student_step >= int(self.cfg.ema_start_step)
        ):
            self.full_ema = _LocalEMA(
                self.model.generator, float(self.cfg.ema_decay)
            )

    @torch.no_grad()
    def _sync_predictor_frozen_from_full(self) -> None:
        if self.mode != "joint":
            return
        with FSDP.summon_full_params(
            self.model.generator,
            recurse=True,
            writeback=False,
            rank0_only=False,
        ):
            unwrap_model(self.predictor).sync_frozen_from_teacher(
                self.model.generator.module.model
            )

    def _student_update(self) -> dict[str, float]:
        batch = self._next_batch()
        prompts = list(batch["prompts"])
        conditional, unconditional = self._conditional_dicts(prompts)
        self.predictor_optimizer.zero_grad(set_to_none=True)
        if self.full_optimizer is not None:
            self.full_optimizer.zero_grad(set_to_none=True)

        loss, logs = self.model.generator_loss(
            image_or_video_shape=self._image_shape(len(prompts)),
            conditional_dict=conditional,
            unconditional_dict=unconditional,
            clean_latent=None,
            initial_latent=None,
        )
        exit_step = int(self.rollout.last_exit_step)
        loss.backward()

        predictor_updated = exit_step > 0
        predictor_grad_norm = torch.zeros((), device=self.device)
        if predictor_updated:
            predictor_grad_norm = torch.nn.utils.clip_grad_norm_(
                unwrap_model(self.predictor).parameters(),
                float(self.cfg.predictor_grad_clip),
            )
            if not torch.isfinite(predictor_grad_norm):
                raise FloatingPointError(
                    f"Non-finite Predictor grad norm {predictor_grad_norm}"
                )
            self.predictor_optimizer.step()
            self.predictor_scheduler.step()
            self.predictor_step += 1

        full_grad_norm = torch.zeros((), device=self.device)
        if self.full_optimizer is not None:
            full_grad_norm = self.model.generator.clip_grad_norm_(
                float(self.cfg.full_grad_clip)
            )
            if not torch.isfinite(full_grad_norm):
                raise FloatingPointError(
                    f"Non-finite Full grad norm {full_grad_norm}"
                )
            self.full_optimizer.step()
            self.full_scheduler.step()
            self._sync_predictor_frozen_from_full()

        self.student_step += 1
        self._maybe_initialize_ema()
        if predictor_updated and self.predictor_ema is not None:
            self.predictor_ema.update()
        if self.full_optimizer is not None and self.full_ema is not None:
            self.full_ema.update()

        result = {
            "dmd_loss": float(loss.detach()),
            "dmd_gradient_norm": float(logs["dmdtrain_gradient_norm"]),
            "dmd_score_timestep": float(logs["timestep"].float().mean()),
            "exit_step": float(exit_step),
            "predictor_updated": float(predictor_updated),
            "predictor_grad_norm": float(predictor_grad_norm),
            "full_grad_norm": float(full_grad_norm),
        }
        del batch, conditional, loss, logs
        return result

    def _checkpoint_metadata(self) -> dict[str, Any]:
        model = unwrap_model(self.predictor)
        return {
            "source_block_ids": list(model.source_block_ids),
            "student_step": self.student_step,
            "predictor_step": self.predictor_step,
            "training_mode": self.mode,
            "training_rollout": (
                "P1_P2_P3" if self.mode == "predictor_only" else "F0_P1_P2_P3"
            ),
            "teacher_checkpoint_key": "generator_ema",
            "predictor_config": model.config_dict,
        }

    def _save(self, *, final: bool = False) -> None:
        dist.barrier()
        suffix = "final" if final else f"step_{self.student_step:05d}"
        checkpoint_dir = self.output_dir / f"checkpoint_{suffix}"
        if self.is_main:
            checkpoint_dir.mkdir(parents=True, exist_ok=True)
            save_predictor_weights(
                self.predictor,
                checkpoint_dir / "predictor.safetensors",
                metadata=self._checkpoint_metadata(),
            )
            if self.predictor_ema is not None:
                with self.predictor_ema.apply():
                    save_predictor_weights(
                        self.predictor,
                        checkpoint_dir / "predictor_ema.safetensors",
                        metadata={
                            **self._checkpoint_metadata(),
                            "ema_decay": float(self.cfg.ema_decay),
                        },
                    )
        dist.barrier()

        critic_state = fsdp_state_dict(self.model.fake_score)
        generator_state = None
        generator_ema_state = None
        if self.mode == "joint":
            generator_state = fsdp_state_dict(self.model.generator)
            if self.full_ema is not None:
                with self.full_ema.apply():
                    generator_ema_state = fsdp_state_dict(self.model.generator)
            else:
                generator_ema_state = generator_state
        if self.is_main:
            payload: dict[str, Any] = {
                "critic": critic_state,
                "student_step": self.student_step,
                "predictor_step": self.predictor_step,
                "training_mode": self.mode,
            }
            if generator_state is not None:
                payload["generator"] = generator_state
                payload["generator_ema"] = generator_ema_state
            atomic_torch_save(payload, checkpoint_dir / "model.pt")
            _atomic_json(
                checkpoint_dir / "state.json",
                {
                    "student_step": self.student_step,
                    "predictor_step": self.predictor_step,
                    "critic_step": self.critic_step,
                    "training_mode": self.mode,
                    "final": final,
                },
            )
            latest = self.output_dir / "latest"
            temporary = self.output_dir / f".latest.{os.getpid()}"
            if temporary.exists() or temporary.is_symlink():
                temporary.unlink()
            temporary.symlink_to(checkpoint_dir.name)
            os.replace(temporary, latest)
            keep = int(self.cfg.keep_checkpoints)
            snapshots = sorted(
                path for path in self.output_dir.glob("checkpoint_step_*")
                if path.is_dir()
            )
            for old in snapshots[:-keep] if keep > 0 else snapshots:
                shutil.rmtree(old)
        dist.barrier()

    def train(self) -> None:
        max_student_steps = int(self.cfg.max_student_steps)
        critic_updates = int(self.cfg.critic_updates_per_student)
        exit_counts = torch.zeros(4, device=self.device, dtype=torch.float64)
        try:
            while self.student_step < max_student_steps:
                started = time.perf_counter()
                critic_loss_sum = 0.0
                critic_grad_sum = 0.0
                for _ in range(critic_updates):
                    metrics = self._critic_update()
                    critic_loss_sum += metrics["critic_loss"]
                    critic_grad_sum += metrics["critic_grad_norm"]
                student = self._student_update()
                exit_counts[int(student["exit_step"])] += 1
                values = torch.tensor(
                    [
                        student["dmd_loss"],
                        student["dmd_gradient_norm"],
                        student["dmd_score_timestep"],
                        student["predictor_grad_norm"],
                        student["full_grad_norm"],
                        critic_loss_sum / critic_updates,
                        critic_grad_sum / critic_updates,
                        time.perf_counter() - started,
                    ],
                    device=self.device,
                    dtype=torch.float64,
                )
                averaged = _distributed_mean(values)
                global_exit_counts = exit_counts.clone()
                dist.all_reduce(global_exit_counts, op=dist.ReduceOp.SUM)
                global_exit_counts /= self.world_size

                if (
                    self.student_step == 1
                    or self.student_step % int(self.cfg.log_every) == 0
                ):
                    record = {
                        "student_step": self.student_step,
                        "predictor_step": self.predictor_step,
                        "critic_step": self.critic_step,
                        "training_mode": self.mode,
                        "exit_step": int(student["exit_step"]),
                        "dmd_loss": float(averaged[0]),
                        "dmd_gradient_norm": float(averaged[1]),
                        "dmd_score_timestep": float(averaged[2]),
                        "predictor_grad_norm": float(averaged[3]),
                        "full_grad_norm": float(averaged[4]),
                        "critic_loss": float(averaged[5]),
                        "critic_grad_norm": float(averaged[6]),
                        "step_time_s": float(averaged[7]),
                        "lr_predictor_fusion": self.predictor_optimizer.param_groups[0]["lr"],
                        "lr_predictor_blocks": self.predictor_optimizer.param_groups[1]["lr"],
                        "lr_full": (
                            0.0
                            if self.full_optimizer is None
                            else self.full_optimizer.param_groups[0]["lr"]
                        ),
                        "lr_fake_score": self.critic_optimizer.param_groups[0]["lr"],
                        "exit_count_f0": int(global_exit_counts[0]),
                        "exit_count_p1": int(global_exit_counts[1]),
                        "exit_count_p2": int(global_exit_counts[2]),
                        "exit_count_p3": int(global_exit_counts[3]),
                        "peak_memory_gib": (
                            torch.cuda.max_memory_allocated(self.device) / 2**30
                        ),
                    }
                    if self.is_main:
                        with self.log_path.open("a", encoding="utf-8") as handle:
                            handle.write(json.dumps(record, sort_keys=True) + "\n")
                        print(json.dumps(record, sort_keys=True), flush=True)
                        if self.swanlab_run is not None:
                            import swanlab

                            swanlab.log(
                                {
                                    key: value
                                    for key, value in record.items()
                                    if not isinstance(value, str)
                                },
                                step=self.student_step,
                            )

                should_save = (
                    not bool(self.root_config.no_save)
                    and (
                        self.student_step % int(self.cfg.save_every) == 0
                        or self.student_step == max_student_steps
                    )
                )
                if should_save:
                    self._save(final=self.student_step == max_student_steps)
                del values, averaged, global_exit_counts, student
                if self.student_step % int(self.root_config.gc_interval) == 0:
                    gc.collect()
                    torch.cuda.empty_cache()
        except BaseException as error:
            if self.swanlab_run is not None:
                try:
                    import swanlab

                    swanlab.finish(error=str(error))
                except Exception:
                    pass
            raise
        else:
            if self.swanlab_run is not None:
                import swanlab

                swanlab.finish()
        finally:
            gc.collect()
            torch.cuda.empty_cache()
            dist.destroy_process_group()


__all__ = ["Trainer"]