#!/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()