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
|
|
"""04 — Self-forcing: causal video distillation (Wan-causal).
|
|
|
|
"Train against your own KV-cached causal rollout." Self-forcing extends DMD2, but the student rollout
|
|
is the SHARED `chunk_rollout` loop (with the slab-KV cache) — exactly the causal loop the engine
|
|
*streams* at serve time. So the causal student is distilled under the runtime it will be served on
|
|
(the (recipe, runtime) pair). The loss is DMD2's distribution-matching, applied to the rollout latents.
|
|
|
|
Run: python3 v2_examples/training/04_self_forcing.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pathlib
|
|
import sys
|
|
|
|
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))
|
|
|
|
from v2.models.wan_causal import build_wan_causal_card
|
|
from v2.training.methods import build_self_forcing
|
|
|
|
|
|
def main() -> None:
|
|
method = build_self_forcing(build_wan_causal_card())
|
|
batch = {"prompts": ["a drone flight over dunes", "a river through a forest"], "seeds": [1, 2]}
|
|
|
|
print("method:", method.name, "| student_loop_id:", method.student_loop_id, "(the causal, KV-cached loop)")
|
|
last = None
|
|
for it in range(12):
|
|
loss, last = method.train_step(batch, it)
|
|
print(f"dmd_loss : {last['dmd_loss']:.4f} critic_loss : {last['critic_loss']:.4f}")
|
|
print(f"grad_norm : student={last.get('grad_norm/student', 0.0):.4f} critic={last['grad_norm/critic']:.4f}")
|
|
print("\nThe student trains on its OWN causal rollout (chunk_rollout + slab-KV) — same loop the "
|
|
"engine streams.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|