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).
82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
"""Model definitions — concrete (recipe, runtime) cards (design_v3 §4).
|
|
|
|
Phase 1: Wan2.1-1.3B (T2V), LTX2.3 (2-stage distilled), Wan-causal (self-forcing student).
|
|
``build_default_engine`` loads all three onto one engine (one resident instance per card).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from .bagel import build_bagel_card, build_bagel_program
|
|
from .cosmos3 import build_cosmos3_card, build_cosmos3_program
|
|
from .ltx2 import build_ltx2_card, build_ltx2_program
|
|
from .unified import build_unified_card, build_unified_program
|
|
from .wan21 import build_wan21_card, build_wan_t2v_program
|
|
from .wan_causal import build_wan_causal_card, build_wan_causal_program
|
|
|
|
__all__ = [
|
|
"build_wan21_card", "build_wan_t2v_program",
|
|
"build_ltx2_card", "build_ltx2_program",
|
|
"build_wan_causal_card", "build_wan_causal_program",
|
|
"build_cosmos3_card", "build_cosmos3_program",
|
|
"build_bagel_card", "build_bagel_program",
|
|
"build_unified_card", "build_unified_program",
|
|
"build_default_engine", "build_omni_engine", "build_unified_engine",
|
|
]
|
|
|
|
_BUILDERS = [
|
|
(build_wan21_card, build_wan_t2v_program),
|
|
(build_ltx2_card, build_ltx2_program),
|
|
(build_wan_causal_card, build_wan_causal_program),
|
|
]
|
|
|
|
# Phase-2 omni cards (MoT, one resident instance running AR + diffusion loops on shared weights)
|
|
_OMNI_BUILDERS = [
|
|
(build_cosmos3_card, build_cosmos3_program),
|
|
(build_bagel_card, build_bagel_program),
|
|
]
|
|
|
|
|
|
def build_default_engine(engine: Any = None) -> Any:
|
|
"""Register Wan2.1, LTX2.3, and Wan-causal onto one engine (one resident instance each)."""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
for build_card, build_program in _BUILDERS:
|
|
card = build_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_program())
|
|
return eng
|
|
|
|
|
|
def build_omni_engine(engine: Any = None) -> Any:
|
|
"""Register the phase-2 omni cards (Cosmos3 + BAGEL/lance) onto one engine.
|
|
|
|
Each is ONE resident MoT instance whose ``transformer`` runs BOTH an ar_decode loop and a
|
|
diffusion_denoise loop (shared weights) — the §16 'true omni/MoT serving' claim made native.
|
|
"""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
for build_card, build_program in _OMNI_BUILDERS:
|
|
card = build_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_program())
|
|
return eng
|
|
|
|
|
|
def build_unified_engine(engine: Any = None) -> Any:
|
|
"""Register the unified LM+generator card (UniRL/PromptRL): a prompt-refiner ``llm`` expert + a
|
|
flow ``transformer`` expert on the SAME engine, each driving its own loop. Two separate experts,
|
|
one request, both trainable under one RL reward (the §10 joint-RL stress test)."""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
card = build_unified_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_unified_program())
|
|
return eng
|