scripts/distill/train_dmdx.py

#!/usr/bin/env python3
"""
Anima DMDX (ADM-only) 蒸留 — arxiv 2507.18569v1 移植
=====================================================

Pipeline:
  - student: wide LoRA on Anima base (DMD2 と同じ wide target、~72M trainable)
  - discriminator: LADD-style (teacher MiniTrainDIT frozen backbone + spectral norm heads)
  - real: teacher_x0_cache の x0 を流用
  - ADM loss: hinge GAN at t-Δt after teacher 1-step evolution
  - time schedule: cubic high-noise bias (paper default)
  - alternation: N_critic × disc → 1 × generator (DMD2 風)

DMD2 との違い:
  - DMD2: dual adapter (student + fake_score)、reverse-KL grad trick
  - DMDX: single student LoRA + 別 D heads、hinge GAN、TVD 最小化

オプション:
  - --recon-weight > 0 で LADD 流 Smooth-L1 anchor 追加 (mean collapse 防止)
  - --misaligned-pairs-d : LADD 流の text-alignment trick

Modal CLI 例:
  modal run modal_app.py::train_dmdx_distill \\
      --warm-lora /models/loras/anima_turbo.safetensors \\
      --total-outer-steps 5000
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import torch
from torch.utils.data import DataLoader

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 save_lora_state
from distill.train_ladd import PrecomputedCacheDataset, ladd_collate
from distill.anima_ladd_disc import AnimaLADDDiscriminator
from distill.dmdx_loss import adm_generator_loss, adm_discriminator_loss


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True, type=str,
                    help="teacher_x0_cache (real x0 source)")
    ap.add_argument("--out", required=True, type=str)
    ap.add_argument("--warm-lora", default="", type=str)
    ap.add_argument("--total-outer-steps", type=int, default=5000,
                    help="1 outer = N_critic + 1 generator")
    ap.add_argument("--n-critic-per-gen", type=int, default=2,
                    help="DMDX paper: hinge GAN は 1:1〜2:1 が一般的 (DMD2 の 5:1 より少ない)")
    ap.add_argument("--n-student-steps", type=int, default=4,
                    help="few-step rollout step 数")
    ap.add_argument("--batch-size", type=int, default=1)
    ap.add_argument("--resolution", type=int, default=768)
    ap.add_argument("--teacher-cfg", type=float, default=4.5,
                    help="Anima 公式推奨 CFG")
    ap.add_argument("--student-cfg", type=float, default=1.0)
    ap.add_argument("--dt-ratio", type=float, default=1.0 / 64,
                    help="paper default Δt = T/64")
    ap.add_argument("--recon-weight", type=float, default=0.0,
                    help="LADD 流 Smooth-L1 anchor (>0 で有効化、stability boost)")
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--lr-gen", type=float, default=5e-6)
    ap.add_argument("--lr-disc", type=float, default=1e-5)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--block-ids", type=str, default="2,8,14,20,26",
                    help="teacher block indices for D hooks")
    ap.add_argument("--head-hidden", type=int, default=512)
    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("--neg-prompt", default="")
    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 base -----
    print("[load] Anima bundle")
    bundle = build_anima(device=device, dtype=dtype)

    # ----- teacher = deepcopy frozen -----
    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 on bundle.transformer -----
    print("[setup] 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

    # warm-start
    if args.warm_lora:
        from distill.train_traj import load_warm_lora
        load_warm_lora(student_transformer, args.warm_lora)

    # ----- discriminator (LADD-style) -----
    print("[setup] DMDX discriminator (teacher backbone + spectral heads)")
    block_ids = [int(x) for x in args.block_ids.split(",")]
    disc = AnimaLADDDiscriminator(
        teacher_transformer=teacher_transformer,
        block_ids=block_ids, head_hidden=args.head_hidden,
    )
    H_lat = args.resolution // 8
    W_lat = args.resolution // 8
    dummy_x = torch.randn(1, 16, 1, H_lat, W_lat, device=device, dtype=dtype)
    dummy_t = torch.tensor([0.5], device=device, dtype=dtype)
    with torch.no_grad():
        dummy_cond = bundle.text_encode([""])
    disc.lazy_init_heads(dummy_x, dummy_t, dummy_cond)
    disc.to(device=device, dtype=torch.float32)
    disc_params = disc.trainable_parameters()
    print(f"[setup] D heads trainable: {sum(p.numel() for p in disc_params)/1e6:.1f}M")

    # ----- optimizers -----
    opt_gen = torch.optim.AdamW(student_params, lr=args.lr_gen, betas=(0.0, 0.999),
                                weight_decay=args.weight_decay, eps=1e-8)
    opt_disc = torch.optim.AdamW(disc_params, lr=args.lr_disc, betas=(0.0, 0.999),
                                 weight_decay=args.weight_decay, eps=1e-8)

    # ----- velocity functions -----
    def student_v(x, t, cond):
        return AnimaBundle.dit_forward(student_transformer, x, t, cond)

    def teacher_v(x, t, cond):
        return AnimaBundle.dit_forward(teacher_transformer, x, t, cond)

    # ----- conditioning (neg prompt cached once) -----
    with torch.no_grad():
        cond_neg = bundle.text_encode([args.neg_prompt or ""])

    # ----- dataset (precompute cache provides real x0 + caption emb) -----
    print(f"[data] loading cache {args.cache_dir}")
    dataset = PrecomputedCacheDataset(args.cache_dir)
    print(f"   {len(dataset)} cached samples")
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, collate_fn=ladd_collate, drop_last=True,
        pin_memory=True,
    )

    # ----- training loop -----
    print(f"[train] outer={args.total_outer_steps} "
          f"cycle=({args.n_critic_per_gen} disc + 1 gen) "
          f"n_student_steps={args.n_student_steps} dt={args.dt_ratio:.4f}")
    log_path = out_dir / "dmdx_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)

    def _prep_batch():
        batch = _next()
        x0_teacher = batch["x0"].to(device=device, dtype=dtype)
        cond_pos = batch["emb"].to(device=device, dtype=dtype)
        B = x0_teacher.size(0)
        cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg
        noise = torch.randn_like(x0_teacher)
        return noise, x0_teacher, cond_pos, cond_neg_b

    for outer in range(args.total_outer_steps):
        # ---- N disc updates ----
        disc_metrics = {}
        for _ in range(args.n_critic_per_gen):
            noise, x0_teacher, cond_pos, cond_neg_b = _prep_batch()
            disc.train()
            student_transformer.eval()
            opt_disc.zero_grad()
            loss_d, m_d = adm_discriminator_loss(
                student_v, teacher_v, disc,
                noise, x0_teacher, cond_pos, cond_neg_b,
                teacher_cfg=args.teacher_cfg, student_cfg=args.student_cfg,
                n_student_steps=args.n_student_steps, dt_ratio=args.dt_ratio,
            )
            loss_d.backward()
            torch.nn.utils.clip_grad_norm_(disc_params, args.grad_clip)
            opt_disc.step()
            disc_metrics = {k: float(v) for k, v in m_d.items()}

        # ---- 1 generator update ----
        noise, x0_teacher, cond_pos, cond_neg_b = _prep_batch()
        student_transformer.train()
        disc.eval()
        opt_gen.zero_grad()
        loss_g, m_g = adm_generator_loss(
            student_v, teacher_v, disc,
            noise, cond_pos, cond_neg_b,
            real_x0=x0_teacher,
            teacher_cfg=args.teacher_cfg, student_cfg=args.student_cfg,
            n_student_steps=args.n_student_steps, dt_ratio=args.dt_ratio,
            recon_weight=args.recon_weight,
        )
        loss_g.backward()
        torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip)
        opt_gen.step()
        gen_metrics = {k: float(v) for k, v in m_g.items()}

        # ---- log ----
        if outer % args.log_every == 0:
            metrics = {"outer": outer, "elapsed": time.time() - t0,
                       **disc_metrics, **gen_metrics}
            log_f.write(json.dumps(metrics) + "\n")
            msg = " ".join(f"{k}={v:.4f}" for k, v in metrics.items() if k != "outer")
            print(f"[outer {outer}/{args.total_outer_steps}] {msg}", flush=True)

        # ---- ckpt (student LoRA only) ----
        if outer > 0 and outer % args.sample_every == 0:
            save_lora_state(student_transformer, out_dir, f"dmdx_student_step{outer:05d}")
            print(f"[save] dmdx_student_step{outer:05d}.safetensors", flush=True)
            try:
                import modal
                modal.Volume.from_name("anima-outputs").commit()
            except Exception as e:
                print(f"[save] volume commit failed: {e}", flush=True)

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


if __name__ == "__main__":
    main()