SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
72 lines
3.2 KiB
Python
72 lines
3.2 KiB
Python
"""Training methods on Wan2.1-1.3B — all over the SHARED loops, train ≡ serve.
|
||
|
||
Verifies: finetune learns (loss ↓), DMD2 real/fake scores, NFT likelihood-free C2 + group reuse +
|
||
old-as-behavior-policy, self-forcing on the causal chunk loop + slab-KV.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from v2._enums import ConsistencyLevel
|
||
from v2.recipes.wan21 import build_wan21_card
|
||
from v2.recipes.wan_causal import build_wan_causal_card
|
||
from v2.training import (
|
||
WeightRole,
|
||
build_diffusion_nft,
|
||
build_dmd2,
|
||
build_finetune,
|
||
build_self_forcing,
|
||
)
|
||
|
||
BATCH = {"prompts": ["a cat", "a dog"], "seeds": [1, 2]}
|
||
|
||
|
||
def test_finetune_learns_and_is_c1():
|
||
ft = build_finetune(build_wan21_card())
|
||
l0, _ = ft.train_step(BATCH, 0)
|
||
l1, m1 = ft.train_step(BATCH, 1)
|
||
assert np.isfinite(l1["loss"]) and l1["loss"] <= l0["loss"] + 1e-9 # real SGD ⇒ non-increasing
|
||
assert m1["grad_norm/student"] > 0.0
|
||
assert ft.consistency_level() is ConsistencyLevel.C1 and not ft.manages_optimization()
|
||
|
||
|
||
def test_dmd2_real_and_fake_scores_finite_with_two_optimizers():
|
||
dm = build_dmd2(build_wan21_card(), generator_update_interval=1)
|
||
lm, m = dm.train_step(BATCH, 0)
|
||
assert np.isfinite(lm["dmd_loss"]) and np.isfinite(lm["critic_loss"])
|
||
assert m["grad_norm/critic"] > 0.0 and m["grad_norm/student"] > 0.0 # student + critic
|
||
|
||
|
||
def test_diffusion_nft_is_likelihood_free_c2_with_group_reuse():
|
||
nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=4, num_inner_timesteps=2)
|
||
assert nft.manages_optimization() and nft.consistency_level() is ConsistencyLevel.C2
|
||
assert nft.old_sync.role is WeightRole.OLD_POLICY # sample from OLD, not student
|
||
lm, m = nft.managed_train_step(BATCH, 0)
|
||
assert np.isfinite(lm["policy_loss"]) and np.isfinite(m["old_deviate"])
|
||
assert m["advantage_std"] >= 0.0
|
||
# the shared prompt encodes ONCE per K-sample group (the 24× text-encode reduction)
|
||
assert nft.old.caches.stats()["feature"]["hits"] > 0
|
||
|
||
|
||
def test_nft_old_policy_decays_toward_student_over_iterations():
|
||
nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=2)
|
||
_, m0 = nft.managed_train_step(BATCH, 0)
|
||
_, m200 = nft.managed_train_step(BATCH, 200)
|
||
assert m0["old_decay"] == 0.0 # decay_type=1: min(step*0.001, 0.5)
|
||
assert m200["old_decay"] == 0.2
|
||
assert m0["old_weights_version"] != m200["old_weights_version"] # version bumped each sync
|
||
|
||
|
||
def test_self_forcing_drives_causal_chunk_loop_and_slab_kv():
|
||
sf = build_self_forcing(build_wan_causal_card(training_mode=True), generator_update_interval=1)
|
||
assert sf.student_loop_id == "chunk_rollout"
|
||
lm, m = sf.train_step(BATCH, 0)
|
||
assert np.isfinite(lm["dmd_loss"]) and np.isfinite(lm["critic_loss"])
|
||
assert sf.student.caches.stats()["slab_kv"]["used_bytes"] > 0 # KV slabs written during rollout
|
||
|
||
|
||
def test_methods_reference_the_same_loops_the_engine_serves():
|
||
assert build_finetune(build_wan21_card()).student_loop_id == "diffusion_denoise"
|
||
assert build_dmd2(build_wan21_card()).student_loop_id == "diffusion_denoise"
|
||
assert build_self_forcing(build_wan_causal_card()).student_loop_id == "chunk_rollout"
|