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).
48 lines
2.5 KiB
Python
48 lines
2.5 KiB
Python
"""Unified LM+generator serve program (design_v3 §15b; UniRL/PromptRL inference path):
|
|
|
|
text_encode → refine (ar_decode on llm) → apply_refinement → denoise (diffusion_denoise) → vae_decode
|
|
|
|
The serve path is deterministic: the LM refines greedily (argmax action) and the generator runs the
|
|
ODE Euler sampler. The ``unified_rl`` *training* method drives the same two loops but in the ROLLOUT
|
|
profile — the LM *samples* an action and the generator uses the SDE sampler with per-step log-prob
|
|
capture (design_v3 §9.4: rollout sampler ≠ serve sampler, a controlled divergence). Same loops, two
|
|
profiles — the §10 "train ≡ serve through one loop" invariant, now spanning two loop *types*.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind
|
|
from ..common import text_encode_node_fn
|
|
from ..omni import vae_decode_node
|
|
from .card import N_REFINE_ACTIONS
|
|
|
|
|
|
def apply_refinement_node(refine_slot: str, n_actions: int = N_REFINE_ACTIONS):
|
|
"""ComponentNode fn: the LM's chosen refinement action reshapes the diffusion conditioning."""
|
|
def fn(instance, slots, request, ctx) -> None:
|
|
llm = instance.component("llm")
|
|
out = slots.get(refine_slot)
|
|
toks = out.get("tokens", []) if isinstance(out, dict) else []
|
|
action = (int(toks[0]) % n_actions) if toks else 0
|
|
slots["text_embeds"] = llm.refined_embed(slots.get("text_embeds"), action)
|
|
slots["refine_action"] = action
|
|
return fn
|
|
|
|
|
|
def build_unified_program() -> Program:
|
|
return Program(
|
|
program_id="unified.unirl", kind=ProgramKind.INLINE,
|
|
nodes=[
|
|
ComponentNode("text_encode", fn=text_encode_node_fn,
|
|
writes=("text_embeds", "neg_text_embeds")),
|
|
ModelLoopNode("refine", loop_id="ar_decode", output_slot="refine_out",
|
|
reads=("text_embeds",), writes=("refine_out",)),
|
|
ComponentNode("apply_refine", fn=apply_refinement_node("refine_out"),
|
|
reads=("refine_out", "text_embeds"), writes=("text_embeds", "refine_action")),
|
|
ModelLoopNode("denoise", loop_id="diffusion_denoise", output_slot="denoise_out",
|
|
reads=("text_embeds",), writes=("denoise_out",)),
|
|
ComponentNode("vae_decode", fn=vae_decode_node("denoise_out", "video"),
|
|
reads=("denoise_out",), writes=("video",)),
|
|
],
|
|
output_artifacts={"video": "video", "latents": "denoise_out"},
|
|
).validate()
|