SolitaryThinker
b3ddf6014d
Stress-tests the v2 Card/Loop/Program design with a UniRL/PromptRL-style joint RL recipe: a prompt-refiner LM expert and a flow generator expert, two separate experts driven by two loop types in one request, both updated simultaneously from a single RL reward. - loop/sampler.py: flow_sde_step_with_logprob — FlowGRPO SDE rollout sampler (per-step Gaussian log-prob), distinct from the deterministic ODE serve step. - request/params.py: gated sde_rollout/sde_noise_scale on DiffusionParams so the serve path stays byte-identical (default ODE). - models/wan21/loop.py: gated SDE-rollout capture in WanDenoiseLoop.advance. - models/backend.py: ToyPromptRefiner — a real REINFORCE categorical policy (the Qwen role), separate weights from the generator. - models/unified/: the unified card+program — two disjoint experts (llm + transformer) on ar_decode + diffusion_denoise; the topological opposite of the Cosmos3 MoT card, same vocabulary. - training/methods/unified_rl.py: joint GRPO — one reward -> group advantage -> LM token policy gradient + DiT FlowGRPO PPO; two LRs; prompt-only/joint flag; reuses the shared diffusion loop for rollout. - training/weight_sync.py: WeightSyncPlan gains a component scope so the two experts version + cache-invalidate independently (LM sync never flushes the frozen text-encoder feature cache). - tests/test_unified_rl.py: 9 tests incl. likelihood-based C2 identity, the two-loop interleave parity gate, joint vs prompt-only. Full suite 91 passed. designv4.md: unified design doc reflecting v2 as built+tested, with the joint RL stress test as the validating case study (the design held — new card + new method, no new runtime primitive).
19 lines
747 B
Python
19 lines
747 B
Python
"""Training methods — all on Wan2.1-1.3B, all over the SHARED loops (design_v3 §10)."""
|
|
from __future__ import annotations
|
|
|
|
from .base import TrainingMethod, new_instance, predict_x0
|
|
from .diffusion_nft import DiffusionNFTMethod, build_diffusion_nft
|
|
from .dmd2 import DMD2Method, build_dmd2
|
|
from .finetune import FineTuneMethod, build_finetune
|
|
from .self_forcing import SelfForcingMethod, build_self_forcing
|
|
from .unified_rl import UnifiedRLMethod, build_unified_rl
|
|
|
|
__all__ = [
|
|
"TrainingMethod", "new_instance", "predict_x0",
|
|
"FineTuneMethod", "build_finetune",
|
|
"DMD2Method", "build_dmd2",
|
|
"DiffusionNFTMethod", "build_diffusion_nft",
|
|
"SelfForcingMethod", "build_self_forcing",
|
|
"UnifiedRLMethod", "build_unified_rl",
|
|
]
|