scripts/distill/train_pcm.py

#!/usr/bin/env python3
"""
Anima PCM (Phased Consistency Model) distillation
=================================================

Reference: G-U-N/Phased-Consistency-Model
            code/text_to_image_sd3/train_pcm_lora_sd3.py

SD3 PCM は FlowMatch + v-pred + LoRA で、Anima の rectified flow と math 完全一致。
ε↔v 変換不要、scheduler/loss を直接移植可能。

Algorithm:
  - num_euler_timesteps=N (e.g. 50) で grid を作り、K phase に等分割
  - 各 step ランダムに index in [0, N) を選び、現 phase の終端まで Euler 1 step
  - student の predicted x_phase と、teacher 1 step 前進 → student EMA-target の x_phase を MSE
  - pseudo-Huber loss + 任意 adv loss (skip in v1)

memory: 1 grad-through student fwd + 3 no_grad fwd (teacher cond, teacher uncond, student-as-EMA)
        → ~60-80 GB on B200 (LoRA-only, batch=1, 768²)

データ: LADD precompute cache の emb を流用 (prompt 再 encode 不要)
        x0 は teacher rollout で online 生成するので不要 (ただし noise / cap 起点として cache を使う)
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from safetensors.torch import save_file, load_file

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from distill.anima_loader import AnimaPaths, build_anima, AnimaBundle
from distill.dmd2_trainer import attach_wide_lora
from distill.train_traj import load_warm_lora, save_lora_state
from distill.pcm_scheduler import make_pcm_solver, euler_multiphase, pseudo_huber


# ----- dataset: precomputed cache の emb と任意で x0 を使う ----------------
class PCMCacheDataset(Dataset):
    """LADD/Reflow cache の (caption, emb, x0) を読む。x0 は initial noise の source。"""
    def __init__(self, cache_dir: str | Path):
        self.cache_dir = Path(cache_dir)
        self.meta = json.loads((self.cache_dir / "metadata.json").read_text(encoding="utf-8"))

    def __len__(self):
        return len(self.meta)

    def __getitem__(self, idx):
        m = self.meta[idx]
        x0 = torch.load(m["x0_path"], map_location="cpu", weights_only=True).squeeze(0)
        emb = torch.load(m["emb_path"], map_location="cpu", weights_only=True).squeeze(0)
        return {"x0": x0, "emb": emb}


def pcm_collate(batch):
    return {
        "x0": torch.stack([b["x0"] for b in batch]),
        "emb": torch.stack([b["emb"] for b in batch]),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True, type=str,
                    help="LADD cache (emb と x0 を使用)")
    ap.add_argument("--out", required=True, type=str)
    ap.add_argument("--warm-lora", default="", type=str)
    ap.add_argument("--total-steps", type=int, default=8000)
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--grad-accum", type=int, default=4)
    ap.add_argument("--resolution", type=int, default=768)
    ap.add_argument("--num-euler-timesteps", type=int, default=50)
    ap.add_argument("--num-phases", type=int, default=4)
    ap.add_argument("--sigma-shift", type=float, default=3.0)
    ap.add_argument("--w-min", type=float, default=4.0, help="CFG-aug min")
    ap.add_argument("--w-max", type=float, default=5.0, help="CFG-aug max")
    ap.add_argument("--w-fixed", type=float, default=-1.0, help=">=0 で範囲指定無視して固定")
    ap.add_argument("--huber-c", type=float, default=1e-3)
    ap.add_argument("--lr", type=float, default=5e-6)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--neg-prompt", default="")
    ap.add_argument("--log-every", type=int, default=10)
    ap.add_argument("--sample-every", type=int, default=500)
    ap.add_argument("--num-workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda")
    dtype = torch.bfloat16
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ----- load Anima -----
    print("[load] Anima bundle")
    bundle = build_anima(device=device, dtype=dtype)

    # teacher = frozen deepcopy
    print("[setup] teacher = frozen deepcopy")
    teacher_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype).eval()
    for p in teacher_transformer.parameters():
        p.requires_grad = False

    # student = wide LoRA
    student_transformer = attach_wide_lora(bundle.transformer, rank=args.lora_rank)
    student_transformer.to(device=device, dtype=dtype)
    for n, p in student_transformer.named_parameters():
        p.requires_grad = ("lora_" in n)
    student_params = [p for p in student_transformer.parameters() if p.requires_grad]
    print(f"[setup] student trainable: {sum(p.numel() for p in student_params)/1e6:.1f}M")
    bundle.transformer = student_transformer

    if args.warm_lora:
        load_warm_lora(student_transformer, args.warm_lora)

    # solver
    solver = make_pcm_solver(
        num_euler_timesteps=args.num_euler_timesteps,
        num_phases=args.num_phases,
        sigma_shift=args.sigma_shift,
        device=device, dtype=torch.float32,
    )
    print(f"[schedule] N={solver.num_steps} K={len(solver.phase_ends)} "
          f"phase_ends={solver.phase_ends.tolist()}")

    # cond_neg
    with torch.no_grad():
        cond_neg = bundle.text_encode([args.neg_prompt or ""])

    # optimizer
    opt = torch.optim.AdamW(student_params, lr=args.lr, betas=(0.9, 0.999),
                            weight_decay=args.weight_decay, eps=1e-8)

    # dataset (emb は cache から、x0 は initial noise の source として使う —
    #  実際の noise は毎 step randn)
    print(f"[data] {args.cache_dir}")
    dataset = PCMCacheDataset(args.cache_dir)
    print(f"   {len(dataset)} captions")
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, collate_fn=pcm_collate,
        drop_last=True, pin_memory=True,
    )

    print(f"[train] steps={args.total_steps} bs={args.batch_size} accum={args.grad_accum} "
          f"N={args.num_euler_timesteps} K={args.num_phases}")
    log_path = out_dir / "pcm_log.jsonl"
    log_f = open(log_path, "a", buffering=1)
    t0 = time.time()
    data_iter = iter(loader)

    def _next():
        nonlocal data_iter
        try:
            return next(data_iter)
        except StopIteration:
            data_iter = iter(loader)
            return next(data_iter)

    for step in range(args.total_steps):
        student_transformer.train()
        opt.zero_grad()
        metrics = {}
        for _ in range(args.grad_accum):
            batch = _next()
            x0 = batch["x0"].to(device=device, dtype=dtype)
            emb = batch["emb"].to(device=device, dtype=dtype)
            B = x0.size(0)
            cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg

            # ----- index sampling -----
            index = torch.randint(0, solver.num_steps, (B,), device=device)
            sigma_cur = solver.sigmas[index].to(dtype=dtype)
            sigma_prev = solver.sigmas[index + 1].to(dtype=dtype)

            # ----- x_t (forward process) -----
            noise = torch.randn_like(x0)
            sigma_cur_ = sigma_cur.view(-1, *([1] * (x0.dim() - 1)))
            x_t = sigma_cur_ * noise + (1.0 - sigma_cur_) * x0

            # ----- student forward (grad) -----
            v_student = AnimaBundle.dit_forward(student_transformer, x_t, sigma_cur, emb)
            # phase-snap: student の "予測" を当該 phase 終端まで進める
            pred_phase, sigma_phase = euler_multiphase(x_t, v_student, index, solver)

            # ----- target side (no_grad) -----
            with torch.no_grad():
                # CFG-aug w sample
                if args.w_fixed >= 0:
                    w = torch.full((B,), args.w_fixed, device=device, dtype=dtype)
                else:
                    w = torch.empty(B, device=device, dtype=dtype).uniform_(args.w_min, args.w_max)
                v_t_cond = AnimaBundle.dit_forward(teacher_transformer, x_t, sigma_cur, emb)
                v_t_uncond = AnimaBundle.dit_forward(teacher_transformer, x_t, sigma_cur, cond_neg_b)
                w_ = w.view(-1, *([1] * (x_t.dim() - 1)))
                v_cfg = v_t_uncond + w_ * (v_t_cond - v_t_uncond)
                # 1 teacher step forward (x_t → x_prev)
                dt = (sigma_prev - sigma_cur).view(-1, *([1] * (x_t.dim() - 1)))
                x_prev = x_t + dt * v_cfg
                # student-as-EMA target prediction at x_prev
                v_tgt = AnimaBundle.dit_forward(student_transformer, x_prev, sigma_prev, emb)
                target_phase, _ = euler_multiphase(x_prev, v_tgt, (index + 1).clamp(max=solver.num_steps - 1), solver)

            # ----- loss -----
            diff = pred_phase - target_phase.detach()
            loss = pseudo_huber(diff, c=args.huber_c).mean() / args.grad_accum
            loss.backward()
            metrics = {
                "loss": float((loss * args.grad_accum).detach()),
                "pred_abs": float(pred_phase.detach().abs().mean()),
                "target_abs": float(target_phase.detach().abs().mean()),
                "w_mean": float(w.mean().detach()),
                "sigma_cur_mean": float(sigma_cur.mean().detach()),
            }
        torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip)
        opt.step()

        if step % args.log_every == 0:
            metrics["step"] = step
            metrics["elapsed"] = time.time() - t0
            log_f.write(json.dumps(metrics) + "\n")
            msg = " ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
                           for k, v in metrics.items() if k != "step")
            print(f"[step {step}/{args.total_steps}] {msg}", flush=True)

        if step > 0 and step % args.sample_every == 0:
            save_lora_state(student_transformer, out_dir, f"pcm_step{step:05d}")
            print(f"[save] pcm_step{step:05d}.safetensors", flush=True)
            try:
                import modal
                modal.Volume.from_name("anima-outputs").commit()
            except Exception:
                pass

    print("[done] saving final")
    save_lora_state(student_transformer, out_dir, "pcm_final")
    log_f.close()


if __name__ == "__main__":
    main()