scripts/distill/train_shortcut.py

#!/usr/bin/env python3
"""
Anima Shortcut Models distillation (Frans et al., 2024)

特徴:
  - 単一 LoRA、d を入力で取って 1/2/4/8/128-step 自在に切替可能
  - Flow-matching half (d=0) + Bootstrap half (d>0) を 1 step 内で混在
  - PCM と違い phase 固定なし、d を連続値で扱える

データ:
  Reflow cache (--save-noise 付き) を流用。(noise, x0, emb) triplet。
"""
from __future__ import annotations
import argparse
import copy
import json
import math
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from distill.anima_loader import build_anima, AnimaBundle
from distill.dmd2_trainer import attach_wide_lora
from distill.train_traj import load_warm_lora, save_lora_state
from distill.train_reflow import ReflowPairDataset, reflow_collate
from distill.shortcut_module import attach_shortcut_d_head, set_shortcut_d, shortcut_d_head_params


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True, type=str,
                    help="Reflow cache (--save-noise 付き)")
    ap.add_argument("--out", required=True, type=str)
    ap.add_argument("--warm-lora", default="", type=str)
    ap.add_argument("--total-steps", type=int, default=2000)
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--grad-accum", type=int, default=2)
    ap.add_argument("--denoise-timesteps", type=int, default=128,
                    help="discrete grid for sampling t and d (2^k granularity)")
    ap.add_argument("--bootstrap-every", type=int, default=8,
                    help="batch 内 bootstrap 比率: 1/N (paper default 8)")
    ap.add_argument("--resolution", type=int, default=768)
    ap.add_argument("--lr", type=float, default=2e-5)
    ap.add_argument("--lr-d-head", type=float, default=5e-4,
                    help="d_head は zero-init なので高めの lr で立ち上げ")
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--clip-x-bootstrap", type=float, default=4.0,
                    help="bootstrap 中の x_t clip range")
    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)

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

    # student = wide LoRA + d-head
    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)
    attach_shortcut_d_head(student_transformer)  # adds d_head + hook + _current_d attr
    # d_head の params は trainable に
    d_head_params = shortcut_d_head_params(student_transformer)
    for p in d_head_params:
        p.requires_grad = True
    student_lora_params = [p for n, p in student_transformer.named_parameters()
                           if p.requires_grad and "lora_" in n]
    print(f"[setup] student LoRA: {sum(p.numel() for p in student_lora_params)/1e6:.1f}M")
    print(f"[setup] d_head:       {sum(p.numel() for p in d_head_params)/1e6:.1f}M")
    bundle.transformer = student_transformer

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

    # optimizers (LoRA と d_head で別 lr)
    opt_lora = torch.optim.AdamW(student_lora_params, lr=args.lr,
                                 betas=(0.9, 0.999), weight_decay=0.01)
    opt_d_head = torch.optim.AdamW(d_head_params, lr=args.lr_d_head,
                                   betas=(0.9, 0.999), weight_decay=0.0)

    # dataset
    print(f"[data] {args.cache_dir}")
    dataset = ReflowPairDataset(args.cache_dir)
    print(f"   {len(dataset)} triplets")
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, collate_fn=reflow_collate,
        drop_last=True, pin_memory=True,
    )

    T = args.denoise_timesteps
    log2_sections = int(math.log2(T))  # 7 for T=128

    def student_v_with_d(x, t, cond, d):
        set_shortcut_d(student_transformer, d)
        try:
            return AnimaBundle.dit_forward(student_transformer, x, t, cond)
        finally:
            set_shortcut_d(student_transformer, None)

    print(f"[train] steps={args.total_steps} bs={args.batch_size} accum={args.grad_accum} "
          f"T={T} bootstrap_every={args.bootstrap_every}")
    log_path = out_dir / "shortcut_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_lora.zero_grad()
        opt_d_head.zero_grad()
        metrics = {}
        for _ in range(args.grad_accum):
            batch = _next()
            x0 = batch["x0"].to(device=device, dtype=dtype)
            noise = batch["noise"].to(device=device, dtype=dtype)
            emb = batch["emb"].to(device=device, dtype=dtype)
            B = x0.size(0)
            B_boot = max(1, B // args.bootstrap_every)
            B_flow = B - B_boot

            # ----- Flow-matching half (d=0, infinitesimal) -----
            t_fm_idx = torch.randint(0, T, (B_flow,), device=device)
            t_fm = t_fm_idx.float() / T
            t_fm_b = t_fm.view(-1, *([1] * (x0.dim() - 1)))
            x0_fm = x0[:B_flow]; noise_fm = noise[:B_flow]; emb_fm = emb[:B_flow]
            x_t_fm = (1 - t_fm_b) * x0_fm + t_fm_b * noise_fm
            v_tgt_fm = noise_fm - x0_fm
            d_fm = torch.zeros(B_flow, device=device, dtype=dtype)

            # ----- Bootstrap half (d > 0, self-consistency) -----
            if B_boot > 0:
                k = torch.randint(0, log2_sections, (B_boot,), device=device)
                d_b = (1.0 / (2.0 ** k.float())).to(dtype=dtype)  # d ∈ {1/2^k}
                # t aligned to d grid: t_b = (random int) / 2^k * 1
                t_b_max = (2 ** k).float()
                t_b_idx = (torch.rand(B_boot, device=device) * t_b_max).floor()
                t_b = t_b_idx / t_b_max.clamp(min=1.0)
                t_b_b = t_b.view(-1, *([1] * (x0.dim() - 1)))
                x0_b = x0[B_flow:]; noise_b = noise[B_flow:]; emb_b = emb[B_flow:]
                x_t_b = (1 - t_b_b) * x0_b + t_b_b * noise_b

                # 2 sub-step bootstrap target (no_grad)
                with torch.no_grad():
                    d_half = d_b * 0.5
                    v1 = student_v_with_d(x_t_b, t_b.to(dtype=dtype), emb_b, d_half)
                    dt_half = d_half.view(-1, *([1] * (x_t_b.dim() - 1)))
                    x_t2 = (x_t_b + dt_half * v1).clamp(-args.clip_x_bootstrap, args.clip_x_bootstrap)
                    t_b_half = (t_b + d_half.float()).clamp(0.0, 1.0)
                    v2 = student_v_with_d(x_t2, t_b_half.to(dtype=dtype), emb_b, d_half)
                    v_tgt_b = 0.5 * (v1 + v2)

                # concat batch
                x_cat = torch.cat([x_t_fm, x_t_b], dim=0)
                t_cat = torch.cat([t_fm.to(dtype=dtype), t_b.to(dtype=dtype)], dim=0)
                d_cat = torch.cat([d_fm, d_b], dim=0)
                emb_cat = torch.cat([emb_fm, emb_b], dim=0)
                v_tgt_cat = torch.cat([v_tgt_fm, v_tgt_b], dim=0)
            else:
                x_cat, t_cat, d_cat, emb_cat, v_tgt_cat = x_t_fm, t_fm.to(dtype=dtype), d_fm, emb_fm, v_tgt_fm

            # 1 trainable forward
            v_pred = student_v_with_d(x_cat, t_cat, emb_cat, d_cat)
            loss = F.mse_loss(v_pred.float(), v_tgt_cat.detach().float()) / args.grad_accum
            loss.backward()
            metrics = {
                "loss": float((loss * args.grad_accum).detach()),
                "B_flow": B_flow, "B_boot": B_boot,
                "v_pred_abs": float(v_pred.detach().abs().mean()),
                "v_tgt_abs": float(v_tgt_cat.detach().abs().mean()),
            }
        torch.nn.utils.clip_grad_norm_(student_lora_params + d_head_params, args.grad_clip)
        opt_lora.step()
        opt_d_head.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:
            # LoRA + d_head 両方保存
            save_lora_state(student_transformer, out_dir, f"shortcut_step{step:05d}")
            torch.save(student_transformer.d_head.state_dict(),
                       out_dir / f"shortcut_d_head_step{step:05d}.pt")
            print(f"[save] shortcut_step{step:05d}", 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, "shortcut_final")
    torch.save(student_transformer.d_head.state_dict(), out_dir / "shortcut_d_head_final.pt")
    log_f.close()


if __name__ == "__main__":
    main()