trainer/predictor_v4_rollout.py

"""Detached chunk+timestep FPPP rollout training for Self-Forcing Predictor-v4."""

from __future__ import annotations

import gc
import json
import math
import os
import random
import time
from contextlib import nullcontext
from pathlib import Path
from typing import Any

import torch
import torch.distributed as dist
import torch.nn.functional as F
from omegaconf import OmegaConf
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.predictor_v4 import SelfForcingPredictorV4
from pipeline.causal_inference import CausalInferencePipeline
from predictor_training.checkpoint import (
    atomic_torch_save,
    capture_rng_state,
    restore_rng_state,
    save_predictor_weights,
    trainable_state_dict,
)
from predictor_training.rollout_cache import (
    assert_clean_history_extent,
    build_predictor_workspace,
    reset_main_caches,
    reset_predictor_workspace,
)
from predictor_training.trajectory_dataset import (
    PredictorV4TrajectoryDataset,
    load_offline_ffff_target,
    trajectory_collate,
)
from utils.misc import set_seed


CHUNK_FRAMES = 3
NUM_CHUNKS = 7
NUM_STEPS = 4
TRAIN_STEPS = (1, 2, 3)
TOKENS_PER_FRAME = 1560
TOKENS_PER_CHUNK = CHUNK_FRAMES * TOKENS_PER_FRAME
EXPECTED_TIMESTEPS = torch.tensor(
    [1000.0, 937.5, 833.3333129882812, 625.0], dtype=torch.float32
)


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)


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


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
            ),
            lambda step: _cosine_with_linear_warmup(
                step, warmup_steps, max_steps
            ),
        ],
    )


def _optimizer_groups(
    model: torch.nn.Module,
) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter], dict[str, str]]:
    fusion: list[torch.nn.Parameter] = []
    blocks: list[torch.nn.Parameter] = []
    names: dict[str, str] = {}
    for name, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        if "predictor_blocks" in name:
            blocks.append(parameter)
            names[name] = "blocks"
        else:
            fusion.append(parameter)
            names[name] = "fusion_residual"
    if not fusion or not blocks:
        raise RuntimeError(
            f"Invalid Stage-2 optimizer groups: fusion={len(fusion)}, "
            f"blocks={len(blocks)}"
        )
    return fusion, blocks, names


def _prune_snapshots(output_dir: Path, keep: int) -> None:
    snapshots = sorted(output_dir.glob("predictor_step_*.safetensors"))
    for path in snapshots[:-keep] if keep > 0 else snapshots:
        path.unlink()


def _load_stage1_model_state(
    model: SelfForcingPredictorV4,
    path: Path,
    *,
    expected_step: int,
) -> None:
    state = torch.load(path, map_location="cpu", weights_only=False)
    actual_step = int(state.get("global_step", -1))
    if actual_step != int(expected_step):
        raise ValueError(
            f"Stage-1 initialization must be step {expected_step}, got {actual_step}"
        )
    expected_blocks = tuple(model.source_block_ids)
    saved_blocks = tuple(int(value) for value in state["source_block_ids"])
    if saved_blocks != expected_blocks:
        raise ValueError(
            f"Stage-1 source blocks {saved_blocks} != requested {expected_blocks}"
        )
    result = model.load_state_dict(state["model"], strict=False)
    trainable = {
        name for name, parameter in model.named_parameters() if parameter.requires_grad
    }
    missing_trainable = sorted(trainable.intersection(result.missing_keys))
    if result.unexpected_keys or missing_trainable:
        raise RuntimeError(
            "Stage-1 Predictor state mismatch: "
            f"unexpected={result.unexpected_keys}, "
            f"missing_trainable={missing_trainable}"
        )
    del state


def _configure_predictor_precision(
    model: SelfForcingPredictorV4,
    *,
    device: torch.device,
) -> None:
    model.to(device=device, dtype=torch.bfloat16)
    with torch.no_grad():
        for parameter in model.parameters():
            if parameter.requires_grad:
                parameter.data = parameter.data.float()


def _load_pipeline_and_predictor(
    cfg,
    *,
    device: torch.device,
) -> tuple[CausalInferencePipeline, SelfForcingPredictorV4]:
    dmd_config = OmegaConf.merge(
        OmegaConf.load(str(cfg.default_config)),
        OmegaConf.load(str(cfg.dmd_config)),
    )
    if int(dmd_config.num_frame_per_block) != CHUNK_FRAMES:
        raise ValueError("Stage-2 rollout requires three latent frames per chunk")
    if bool(dmd_config.independent_first_frame):
        raise ValueError("Stage-2 rollout requires independent_first_frame=false")
    if bool(getattr(dmd_config, "reuse_first_step_velocity", False)):
        raise ValueError("Stage-2 Teacher must use Full FFFF denoising")

    pipeline = CausalInferencePipeline(
        dmd_config,
        device=device,
        vae=torch.nn.Identity(),
    )
    checkpoint = torch.load(
        str(cfg.teacher_checkpoint), map_location="cpu", weights_only=False
    )
    checkpoint_key = str(cfg.teacher_checkpoint_key)
    if checkpoint_key not in checkpoint:
        raise KeyError(
            f"{cfg.teacher_checkpoint} lacks checkpoint key {checkpoint_key!r}"
        )
    pipeline.generator.load_state_dict(
        checkpoint[checkpoint_key],
        strict=bool(cfg.strict_teacher_load),
    )
    del checkpoint

    predictor = SelfForcingPredictorV4.from_teacher(
        pipeline.generator.model,
        source_block_ids=tuple(int(value) for value in cfg.source_block_ids),
    )

    pipeline.generator.eval().requires_grad_(False)
    pipeline.text_encoder.eval().requires_grad_(False)
    pipeline.generator.to(device=device, dtype=torch.bfloat16)
    pipeline.text_encoder.to(device=device, dtype=torch.bfloat16)
    _configure_predictor_precision(predictor, device=device)
    # Load the Stage-1 FP32 master state only after trainable parameters have
    # been promoted back to FP32; loading before .to(bfloat16) would silently
    # quantize the initialization.
    _load_stage1_model_state(
        predictor,
        Path(str(cfg.stage1_training_state)).resolve(),
        expected_step=int(cfg.stage1_expected_step),
    )
    predictor.train()
    pipeline._initialize_kv_cache(
        batch_size=1, dtype=torch.bfloat16, device=device
    )
    pipeline._initialize_crossattn_cache(
        batch_size=1, dtype=torch.bfloat16, device=device
    )
    actual = pipeline.denoising_step_list.detach().float().cpu()
    if not torch.equal(actual, EXPECTED_TIMESTEPS):
        raise ValueError(
            f"Stage-2 requires exact FP32 timesteps {EXPECTED_TIMESTEPS.tolist()}, "
            f"got {actual.tolist()}"
        )
    return pipeline, predictor


def _fixed_noise_trajectory(
    *,
    device: torch.device,
) -> tuple[torch.Tensor, list[list[torch.Tensor]]]:
    """Reproduce inference seed-0 initial and per-transition noise draws."""

    set_seed(0)
    initial = torch.randn(
        [1, NUM_CHUNKS * CHUNK_FRAMES, 16, 60, 104],
        device=device,
        dtype=torch.bfloat16,
    )
    transition_noise = [
        [
            torch.randn(
                [1, CHUNK_FRAMES, 16, 60, 104],
                device=device,
                dtype=torch.bfloat16,
            )
            for _ in range(NUM_STEPS - 1)
        ]
        for _ in range(NUM_CHUNKS)
    ]
    return initial, transition_noise


def _timestep(
    timesteps: torch.Tensor,
    step_id: int,
    *,
    device: torch.device,
) -> torch.Tensor:
    return torch.full(
        (1, CHUNK_FRAMES),
        float(timesteps[int(step_id)].item()),
        device=device,
        dtype=torch.float32,
    )


def _renoise(
    pipeline: CausalInferencePipeline,
    clean: torch.Tensor,
    noise: torch.Tensor,
    *,
    next_timestep: torch.Tensor,
) -> torch.Tensor:
    return pipeline.scheduler.add_noise(
        clean.detach().flatten(0, 1),
        noise.flatten(0, 1),
        next_timestep.flatten(0, 1),
    ).unflatten(0, clean.shape[:2]).detach()


@torch.no_grad()
def _full_step(
    pipeline: CausalInferencePipeline,
    *,
    latent: torch.Tensor,
    conditional_dict: dict[str, torch.Tensor],
    timestep: torch.Tensor,
    current_start: int,
    capture_hidden: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        flow, clean, hidden = pipeline._full_step_with_optional_hidden(
            noisy_input=latent,
            conditional_dict=conditional_dict,
            timestep=timestep,
            current_start=current_start,
            capture_hidden=capture_hidden,
        )
    return flow.detach(), clean.detach(), None if hidden is None else hidden.detach()


@torch.no_grad()
def _commit_clean_history(
    pipeline: CausalInferencePipeline,
    *,
    clean: torch.Tensor,
    conditional_dict: dict[str, torch.Tensor],
    current_start: int,
    expected_end: int,
) -> None:
    timestep = torch.full(
        (1, clean.shape[1]),
        float(pipeline.args.context_noise),
        device=clean.device,
        dtype=torch.float32,
    )
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        pipeline.generator(
            noisy_image_or_video=clean,
            conditional_dict=conditional_dict,
            timestep=timestep,
            kv_cache=pipeline.kv_cache1,
            crossattn_cache=pipeline.crossattn_cache,
            current_start=current_start,
        )
    assert_clean_history_extent(
        pipeline.kv_cache1, expected_tokens=expected_end
    )


class Trainer:
    """Trainer dispatch target for both Stage-2 supervision variants."""

    def __init__(self, config) -> None:
        if not torch.cuda.is_available():
            raise RuntimeError("Predictor-v4 Stage-2 requires CUDA")
        required = {"RANK", "WORLD_SIZE", "LOCAL_RANK"}
        if not required.issubset(os.environ):
            raise RuntimeError("Launch Predictor-v4 Stage-2 with torchrun")

        self.root_config = config
        self.cfg = config.predictor_v4_rollout
        self.rank = int(os.environ["RANK"])
        self.world_size = int(os.environ["WORLD_SIZE"])
        self.local_rank = int(os.environ["LOCAL_RANK"])
        if self.world_size != 8:
            raise ValueError(f"Formal Stage-2 requires 8 ranks, got {self.world_size}")
        torch.cuda.set_device(self.local_rank)
        self.device = torch.device("cuda", self.local_rank)
        dist.init_process_group(backend="nccl")
        self.is_main = self.rank == 0
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.set_float32_matmul_precision("high")

        seed = int(self.root_config.seed)
        random.seed(seed + self.rank)
        torch.manual_seed(seed + self.rank)
        torch.cuda.manual_seed_all(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.dataset = PredictorV4TrajectoryDataset(
            self.cfg.cases_path,
            data_root=self.cfg.data_root,
            max_cases=(
                None if self.cfg.max_cases is None else int(self.cfg.max_cases)
            ),
        )
        self.sampler = DistributedSampler(
            self.dataset,
            num_replicas=self.world_size,
            rank=self.rank,
            shuffle=True,
            seed=seed,
            drop_last=True,
        )
        self.loader = DataLoader(
            self.dataset,
            batch_size=1,
            sampler=self.sampler,
            num_workers=int(self.cfg.num_workers),
            pin_memory=bool(self.cfg.pin_memory),
            persistent_workers=(
                int(self.cfg.num_workers) > 0
                and bool(self.cfg.persistent_workers)
            ),
            collate_fn=trajectory_collate,
        )
        if not self.loader:
            raise ValueError("Stage-2 trajectory loader is empty")

        self.pipeline, self.model = _load_pipeline_and_predictor(
            self.cfg, device=self.device
        )
        fusion, blocks, group_names = _optimizer_groups(self.model)
        self.optimizer = AdamW(
            [
                {
                    "params": fusion,
                    "lr": float(self.cfg.fusion_lr),
                    "name": "fusion_residual",
                },
                {
                    "params": blocks,
                    "lr": float(self.cfg.blocks_lr),
                    "name": "blocks",
                },
            ],
            betas=(float(self.cfg.beta1), float(self.cfg.beta2)),
            weight_decay=float(self.cfg.weight_decay),
        )
        self.scheduler = _make_scheduler(
            self.optimizer,
            warmup_steps=int(self.cfg.warmup_steps),
            max_steps=int(self.cfg.max_steps),
        )
        self.global_step = 0
        self.epoch = 0
        self.batch_in_epoch = 0
        self._resume(group_names)
        self.ddp = DDP(
            self.model,
            device_ids=[self.local_rank],
            output_device=self.local_rank,
            broadcast_buffers=False,
            gradient_as_bucket_view=True,
            find_unused_parameters=False,
        )
        self.run_config = self._run_config(group_names)
        self.swanlab_run = self._initialize_swanlab()
        if self.is_main:
            _atomic_json(self.output_dir / "train_config.json", self.run_config)
            print(json.dumps(self.run_config, indent=2, default=str), flush=True)

    @property
    def supervision_mode(self) -> str:
        return str(self.cfg.supervision_mode)

    def _resume(self, group_names: dict[str, str]) -> None:
        resume = self.cfg.resume
        if resume is None or str(resume).lower() in {"", "none", "null"}:
            return
        path = Path(str(resume)).resolve()
        state = torch.load(path, map_location="cpu", weights_only=False)
        if str(state["supervision_mode"]) != self.supervision_mode:
            raise ValueError("Stage-2 resume supervision mode changed")
        if state["parameter_groups"] != group_names:
            raise ValueError("Stage-2 resume optimizer grouping changed")
        result = self.model.load_state_dict(state["model"], strict=False)
        trainable = {
            name
            for name, parameter in self.model.named_parameters()
            if parameter.requires_grad
        }
        missing_trainable = sorted(trainable.intersection(result.missing_keys))
        if result.unexpected_keys or missing_trainable:
            raise RuntimeError(
                "Stage-2 resume mismatch: "
                f"unexpected={result.unexpected_keys}, "
                f"missing_trainable={missing_trainable}"
            )
        self.optimizer.load_state_dict(state["optimizer"])
        self.scheduler.load_state_dict(state["scheduler"])
        self.global_step = int(state["global_step"])
        self.epoch = int(state["epoch"])
        self.batch_in_epoch = int(state["batch_in_epoch"])
        rng_path = path.parent / f"rng_rank_{self.rank:02d}.pt"
        if not rng_path.is_file():
            raise FileNotFoundError(rng_path)
        restore_rng_state(
            torch.load(rng_path, map_location="cpu", weights_only=False)
        )

    def _run_config(self, group_names: dict[str, str]) -> dict[str, Any]:
        return {
            **OmegaConf.to_container(self.cfg, resolve=True),
            "trainer_semantics": (
                "chunk0_ffff_chunks1-6_fppp_detached_timestep_and_chunk"
            ),
            "train_predictor_steps": [1, 2, 3],
            "persistent_cache_semantics": "final_clean_full_pass_only",
            "exact_online_timesteps": EXPECTED_TIMESTEPS.tolist(),
            "offline_target_timestep_semantics": (
                "saved_stage1_int64" if self.supervision_mode == "offline_ffff"
                else None
            ),
            "world_size": self.world_size,
            "global_batch_trajectories": self.world_size,
            "supervised_states_per_optimizer_step": (
                self.world_size * (NUM_CHUNKS - 1) * len(TRAIN_STEPS)
            ),
            "parameter_groups": group_names,
            "trainable_parameters": sum(
                parameter.numel()
                for parameter in self.model.parameters()
                if parameter.requires_grad
            ),
            "activation_dtype": "bfloat16",
            "trainable_parameter_dtype": "float32",
        }

    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
        run = 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,
        )
        self.run_config["swanlab_run_id"] = getattr(run, "id", None)
        return run

    def _save(self) -> None:
        dist.barrier()
        atomic_torch_save(
            capture_rng_state(),
            self.output_dir / f"rng_rank_{self.rank:02d}.pt",
        )
        if self.is_main:
            weights_path = (
                self.output_dir
                / f"predictor_step_{self.global_step:05d}.safetensors"
            )
            save_predictor_weights(
                self.model,
                weights_path,
                metadata={
                    "source_block_ids": list(self.model.source_block_ids),
                    "global_step": self.global_step,
                    "stage": 2,
                    "training_rollout": "fppp",
                    "supervision_mode": self.supervision_mode,
                    "predictor_config": self.model.config_dict,
                },
            )
            atomic_torch_save(
                {
                    "model": trainable_state_dict(self.model),
                    "optimizer": self.optimizer.state_dict(),
                    "scheduler": self.scheduler.state_dict(),
                    "global_step": self.global_step,
                    "epoch": self.epoch,
                    "batch_in_epoch": self.batch_in_epoch,
                    "supervision_mode": self.supervision_mode,
                    "source_block_ids": tuple(self.model.source_block_ids),
                    "parameter_groups": self.run_config["parameter_groups"],
                    "config": self.run_config,
                    "weights_path": str(weights_path),
                },
                self.output_dir / "training_latest.pt",
            )
            _prune_snapshots(self.output_dir, int(self.cfg.keep_snapshots))
        dist.barrier()

    def _target(
        self,
        *,
        case_id: int,
        chunk_id: int,
        step_id: int,
        latent: torch.Tensor,
        conditional_dict: dict[str, torch.Tensor],
        timestep: torch.Tensor,
        current_start: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if self.supervision_mode == "onpolicy":
            flow, _, hidden = _full_step(
                self.pipeline,
                latent=latent,
                conditional_dict=conditional_dict,
                timestep=timestep,
                current_start=current_start,
                capture_hidden=True,
            )
            if hidden is None:
                raise RuntimeError("On-policy Full Teacher did not return hidden")
            return flow, hidden
        if self.supervision_mode == "offline_ffff":
            target = load_offline_ffff_target(
                self.dataset.offline_step_path(case_id, chunk_id),
                step_id=step_id,
                device=self.device,
            )
            return target["flow"], target["hidden"]
        raise ValueError(f"Unknown supervision mode {self.supervision_mode!r}")

    def _run_trajectory(self, record: dict[str, Any]) -> dict[str, float]:
        started = time.perf_counter()
        case_id = int(record["case_id"])
        if int(record["seed"]) != 0:
            raise ValueError(f"case {case_id} does not use seed 0")

        reset_main_caches(
            self.pipeline.kv_cache1, self.pipeline.crossattn_cache
        )
        initial_noise, transition_noise = _fixed_noise_trajectory(
            device=self.device
        )
        with torch.no_grad(), torch.autocast(
            device_type="cuda", dtype=torch.bfloat16
        ):
            conditional_dict = self.pipeline.text_encoder(
                text_prompts=[str(record["prompt"])]
            )
        timesteps = self.pipeline.denoising_step_list.to(
            device=self.device, dtype=torch.float32
        )
        previous_hidden: dict[int, torch.Tensor] = {}
        metric_sums = torch.zeros(7, device=self.device, dtype=torch.float64)
        # total, h1, h2, h3, flow1, flow2, flow3
        num_predictor_calls = (NUM_CHUNKS - 1) * len(TRAIN_STEPS)
        predictor_call = 0
        self.optimizer.zero_grad(set_to_none=True)

        for chunk_id in range(NUM_CHUNKS):
            current_start = chunk_id * TOKENS_PER_CHUNK
            current_end = current_start + TOKENS_PER_CHUNK
            latent = initial_noise[
                :, chunk_id * CHUNK_FRAMES : (chunk_id + 1) * CHUNK_FRAMES
            ].detach()
            predictor_cache = None
            if chunk_id > 0:
                predictor_cache = build_predictor_workspace(
                    self.pipeline.kv_cache1,
                    source_block_ids=tuple(self.model.source_block_ids),
                    history_tokens=current_start,
                    current_tokens=TOKENS_PER_CHUNK,
                )

            if chunk_id == 0:
                chunk_hidden: dict[int, torch.Tensor] = {}
                final_clean = None
                for step_id in range(NUM_STEPS):
                    timestep = _timestep(
                        timesteps, step_id, device=self.device
                    )
                    _, clean, hidden = _full_step(
                        self.pipeline,
                        latent=latent,
                        conditional_dict=conditional_dict,
                        timestep=timestep,
                        current_start=current_start,
                        capture_hidden=step_id in TRAIN_STEPS,
                    )
                    if step_id in TRAIN_STEPS:
                        if hidden is None:
                            raise RuntimeError("Chunk 0 Full hidden is missing")
                        chunk_hidden[step_id] = hidden
                    final_clean = clean
                    if step_id < NUM_STEPS - 1:
                        latent = _renoise(
                            self.pipeline,
                            clean,
                            transition_noise[chunk_id][step_id],
                            next_timestep=_timestep(
                                timesteps, step_id + 1, device=self.device
                            ),
                        )
                if final_clean is None:
                    raise RuntimeError("Chunk 0 did not produce a final clean latent")
                _commit_clean_history(
                    self.pipeline,
                    clean=final_clean,
                    conditional_dict=conditional_dict,
                    current_start=current_start,
                    expected_end=current_end,
                )
                previous_hidden = {
                    step: chunk_hidden[step].detach() for step in TRAIN_STEPS
                }
                del chunk_hidden, final_clean
                continue

            timestep0 = _timestep(timesteps, 0, device=self.device)
            _, clean0, anchor = _full_step(
                self.pipeline,
                latent=latent,
                conditional_dict=conditional_dict,
                timestep=timestep0,
                current_start=current_start,
                capture_hidden=True,
            )
            if anchor is None:
                raise RuntimeError(f"Chunk {chunk_id} has no Full step-0 anchor")
            latent = _renoise(
                self.pipeline,
                clean0,
                transition_noise[chunk_id][0],
                next_timestep=_timestep(timesteps, 1, device=self.device),
            )
            current_hidden: dict[int, torch.Tensor] = {}

            for step_id in TRAIN_STEPS:
                reset_predictor_workspace(
                    predictor_cache,
                    history_tokens=current_start,
                )
                if latent.requires_grad or anchor.requires_grad:
                    raise RuntimeError(
                        f"Chunk {chunk_id} step {step_id} received a "
                        "gradient-connected detached rollout state"
                    )
                if previous_hidden[step_id].requires_grad:
                    raise RuntimeError(
                        f"Chunk {chunk_id} step {step_id} previous hidden "
                        "retained a graph across chunks"
                    )
                if any(
                    cache["k"].requires_grad or cache["v"].requires_grad
                    for cache in predictor_cache.values()
                ):
                    raise RuntimeError(
                        f"Chunk {chunk_id} step {step_id} Predictor workspace "
                        "retained an earlier timestep graph"
                    )
                timestep = _timestep(timesteps, step_id, device=self.device)
                target_flow, target_hidden = self._target(
                    case_id=case_id,
                    chunk_id=chunk_id,
                    step_id=step_id,
                    latent=latent,
                    conditional_dict=conditional_dict,
                    timestep=timestep,
                    current_start=current_start,
                )
                predictor_call += 1
                sync = predictor_call == num_predictor_calls
                sync_context = nullcontext() if sync else self.ddp.no_sync()
                with sync_context:
                    with torch.autocast(
                        device_type="cuda", dtype=torch.bfloat16
                    ):
                        output = self.ddp(
                            target_latent=latent,
                            target_timestep=timestep,
                            anchor_hidden=anchor,
                            previous_chunk_hidden=previous_hidden[step_id],
                            kv_cache=predictor_cache,
                            crossattn_cache=self.pipeline.crossattn_cache,
                            current_start=current_start,
                        )
                        hidden_loss = F.mse_loss(
                            output["pred_hidden"].float(),
                            target_hidden.float(),
                        )
                        flow_loss = F.mse_loss(
                            output["pred_flow"].float(),
                            target_flow.float(),
                        )
                        step_loss = (
                            float(self.cfg.flow_loss_weight) * flow_loss
                            + float(self.cfg.hidden_loss_weight) * hidden_loss
                        )
                    pred_hidden = output["pred_hidden"].detach()
                    pred_flow = output["pred_flow"].detach()
                    (step_loss / num_predictor_calls).backward()

                metric_sums[0] += step_loss.detach()
                metric_sums[step_id] += hidden_loss.detach()
                metric_sums[3 + step_id] += flow_loss.detach()
                current_hidden[step_id] = pred_hidden
                anchor = pred_hidden
                pred_clean = self.pipeline.generator._convert_flow_pred_to_x0(
                    flow_pred=pred_flow.flatten(0, 1),
                    xt=latent.flatten(0, 1),
                    timestep=timestep.flatten(0, 1),
                ).unflatten(0, pred_flow.shape[:2]).detach()
                if step_id < NUM_STEPS - 1:
                    latent = _renoise(
                        self.pipeline,
                        pred_clean,
                        transition_noise[chunk_id][step_id],
                        next_timestep=_timestep(
                            timesteps, step_id + 1, device=self.device
                        ),
                    )
                del (
                    output,
                    target_flow,
                    target_hidden,
                    hidden_loss,
                    flow_loss,
                    step_loss,
                    pred_flow,
                )

            # Training rollout is always FPPP: P3 supplies both final latent and
            # next-chunk step-3 hidden. Only this final clean state is committed.
            _commit_clean_history(
                self.pipeline,
                clean=pred_clean,
                conditional_dict=conditional_dict,
                current_start=current_start,
                expected_end=current_end,
            )
            previous_hidden = {
                step: current_hidden[step].detach() for step in TRAIN_STEPS
            }
            del (
                predictor_cache,
                current_hidden,
                pred_clean,
                clean0,
                anchor,
            )

        if predictor_call != num_predictor_calls:
            raise RuntimeError(
                f"Expected {num_predictor_calls} Predictor calls, got "
                f"{predictor_call}"
            )
        grad_norm = torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), float(self.cfg.grad_clip)
        )
        if not torch.isfinite(grad_norm):
            raise FloatingPointError(f"Non-finite Stage-2 grad norm {grad_norm}")
        self.optimizer.step()
        self.scheduler.step()
        self.optimizer.zero_grad(set_to_none=True)
        metric_sums[0] /= num_predictor_calls
        metric_sums[1:] /= NUM_CHUNKS - 1
        result = {
            "loss": float(metric_sums[0]),
            "hidden_mse_step1": float(metric_sums[1]),
            "hidden_mse_step2": float(metric_sums[2]),
            "hidden_mse_step3": float(metric_sums[3]),
            "flow_mse_step1": float(metric_sums[4]),
            "flow_mse_step2": float(metric_sums[5]),
            "flow_mse_step3": float(metric_sums[6]),
            "grad_norm": float(grad_norm),
            "trajectory_time_s": time.perf_counter() - started,
        }
        del (
            conditional_dict,
            initial_noise,
            transition_noise,
            previous_hidden,
            metric_sums,
        )
        return result

    def train(self) -> None:
        max_steps = int(self.cfg.max_steps)
        try:
            while self.global_step < max_steps:
                self.sampler.set_epoch(self.epoch)
                for batch_index, record in enumerate(self.loader):
                    if batch_index < self.batch_in_epoch:
                        continue
                    self.batch_in_epoch = batch_index + 1
                    metrics = self._run_trajectory(record)
                    self.global_step += 1
                    values = torch.tensor(
                        [
                            metrics["loss"],
                            metrics["hidden_mse_step1"],
                            metrics["hidden_mse_step2"],
                            metrics["hidden_mse_step3"],
                            metrics["flow_mse_step1"],
                            metrics["flow_mse_step2"],
                            metrics["flow_mse_step3"],
                            metrics["trajectory_time_s"],
                        ],
                        device=self.device,
                        dtype=torch.float64,
                    )
                    averaged = _distributed_mean(values)
                    max_grad = torch.tensor(
                        metrics["grad_norm"],
                        device=self.device,
                        dtype=torch.float64,
                    )
                    dist.all_reduce(max_grad, op=dist.ReduceOp.MAX)
                    if (
                        self.global_step == 1
                        or self.global_step % int(self.cfg.log_every) == 0
                    ):
                        record_out = {
                            "global_step": self.global_step,
                            "epoch": self.epoch,
                            "supervision_mode": self.supervision_mode,
                            "loss": float(averaged[0]),
                            "hidden_mse_step1": float(averaged[1]),
                            "hidden_mse_step2": float(averaged[2]),
                            "hidden_mse_step3": float(averaged[3]),
                            "flow_mse_step1": float(averaged[4]),
                            "flow_mse_step2": float(averaged[5]),
                            "flow_mse_step3": float(averaged[6]),
                            "trajectory_time_s": float(averaged[7]),
                            "grad_norm_max": float(max_grad),
                            "lr_fusion": self.optimizer.param_groups[0]["lr"],
                            "lr_blocks": self.optimizer.param_groups[1]["lr"],
                            "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_out, sort_keys=True) + "\n"
                                )
                            print(json.dumps(record_out, sort_keys=True), flush=True)
                            if self.swanlab_run is not None:
                                import swanlab

                                swanlab.log(record_out, step=self.global_step)
                    should_save = (
                        not bool(self.root_config.no_save)
                        and (
                            self.global_step % int(self.cfg.save_every) == 0
                            or self.global_step == max_steps
                        )
                    )
                    if should_save:
                        self._save()
                    del metrics, values, averaged, max_grad
                    if self.global_step >= max_steps:
                        break
                if self.global_step < max_steps:
                    self.epoch += 1
                    self.batch_in_epoch = 0
        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"]