SolitaryThinker
94b84b03b9
Three more example folders alongside inference/, all CPU/numpy, self-contained (sys.path bootstrap), public API only, every script verified to run green. training/ (7) — one per method, all via the uniform method.train_step seam: 01 finetune · 02 dmd2 distillation · 03 diffusion_nft (likelihood-free RL, samples from the old policy, feature-cache reuse) · 04 self-forcing (causal chunk_rollout) · 05 joint LM+generator RL (UniRL; joint + prompt-only) · 06 N-way joint RL (per_expert vs shared credit) · 07 end-to-end workflow RL (T2I+I2V from one final-video reward). omni/ (4) — 01 Cosmos3 (reason→joint denoise, shared MoT) · 02 BAGEL (text→image, shared MoT; scheduler prices both WorkUnit kinds) · 03 Qwen-Omni (thinker→talker→vocoder, three separate experts, text+audio) · 04 interleave parity across AR + diffusion loop types. workflows/ (2) — 01 cross-model T2I→I2V workflow (image provably conditions the video) · 02 workflow as a first-class servable (requires/validate, address by id, register_workflows catalog, WorkflowRegistry). Each folder has a README indexing its scripts.
38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
||
"""02 — DMD2 distillation (Wan2.1-1.3B).
|
||
|
||
Distribution-matching distillation: three roles on the same card — a **student** (the few-step model
|
||
being trained), a trainable **fake-score critic**, and a frozen **teacher**. The generator loss is the
|
||
DMD score difference (teacher − critic) on the student's own rollout latents; the critic is trained to
|
||
score the student. Two grad streams (student + critic) ⇒ two grad-norm signals.
|
||
|
||
Run: python3 v2_examples/training/02_distillation_dmd2.py
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import pathlib
|
||
import sys
|
||
|
||
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))
|
||
|
||
from v2.models.wan21 import build_wan21_card
|
||
from v2.training.methods import build_dmd2
|
||
|
||
|
||
def main() -> None:
|
||
method = build_dmd2(build_wan21_card(), rollout_steps=3)
|
||
batch = {"prompts": ["a waterfall in a canyon", "a lantern festival"], "seeds": [1, 2]}
|
||
|
||
print("method:", method.name, "| roles: student + fake-score critic + frozen teacher")
|
||
last = None
|
||
for it in range(15):
|
||
loss, last = method.train_step(batch, it)
|
||
print(f"dmd_loss : {last['dmd_loss']:.4f} (generator: teacher − critic score difference)")
|
||
print(f"critic_loss : {last['critic_loss']:.4f} (critic learns to score the student's samples)")
|
||
print(f"grad_norm : student={last.get('grad_norm/student', 0.0):.4f} "
|
||
f"critic={last['grad_norm/critic']:.4f}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|