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).
33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
"""Training plane — rollout records behavior on the same loops it serves (design_v3 §10).
|
|
|
|
Dependency rule (enforced): ``training`` imports ``card``/``loop``/``runtime``/``models``; the
|
|
engine never imports ``training``. The rollout forward IS the serve forward plus capture.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from .behavior import BehaviorRecord, TrajectoryStatus
|
|
from .methods import (
|
|
DiffusionNFTMethod,
|
|
DMD2Method,
|
|
FineTuneMethod,
|
|
SelfForcingMethod,
|
|
TrainingMethod,
|
|
UnifiedRLMethod,
|
|
build_diffusion_nft,
|
|
build_dmd2,
|
|
build_finetune,
|
|
build_self_forcing,
|
|
build_unified_rl,
|
|
)
|
|
from .rewards import MultiRewardScorer, build_multi_reward_scorer
|
|
from .rollout import RolloutContext, Trajectory, TrajectoryBuffer, rollout_loop
|
|
from .weight_sync import WeightRole, WeightSyncPlan
|
|
|
|
__all__ = [
|
|
"TrainingMethod", "FineTuneMethod", "DMD2Method", "DiffusionNFTMethod", "SelfForcingMethod",
|
|
"UnifiedRLMethod",
|
|
"build_finetune", "build_dmd2", "build_diffusion_nft", "build_self_forcing", "build_unified_rl",
|
|
"BehaviorRecord", "TrajectoryStatus", "Trajectory", "TrajectoryBuffer", "RolloutContext",
|
|
"rollout_loop", "WeightSyncPlan", "WeightRole", "MultiRewardScorer", "build_multi_reward_scorer",
|
|
]
|