SolitaryThinker
3ec511419f
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).
64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
"""Request parameters: AR sampling vs diffusion params + output/capture spec.
|
|
|
|
design_v3 §12: "AR ``sampling`` vs ``diffusion`` params, an ``OutputSpec``
|
|
(requested modalities + streaming + capture flags), and per-node overrides."
|
|
design.md §6.1: ``sampling``/``diffusion`` are *defaults*; ``node_params`` override
|
|
per graph node (validated against the PipelineSpec's node-parameter schema).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SamplingParams:
|
|
"""AR decode knobs (used by ar_decode loops: reasoner/thinker/talker — phase 2)."""
|
|
max_tokens: int = 512
|
|
temperature: float = 1.0
|
|
top_p: float = 1.0
|
|
top_k: int = 0
|
|
stop: tuple[str, ...] = ()
|
|
seed: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiffusionParams:
|
|
"""Denoise knobs. ``guidance`` may be a scalar or per-modality dict (joint A/V)."""
|
|
num_steps: int = 50
|
|
guidance_scale: float = 5.0
|
|
# Per-modality guidance for joint denoise (LTX-2 A/V, Cosmos3 t2vs). Empty => use scalar.
|
|
guidance_per_modality: dict[str, float] = field(default_factory=dict)
|
|
shift: float = 5.0 # flow-match sigma shift
|
|
sigmas: tuple[float, ...] | None = None # explicit schedule (distilled few-step)
|
|
negative_prompt: str = ""
|
|
height: int = 480
|
|
width: int = 832
|
|
num_frames: int = 81
|
|
fps: int = 16
|
|
seed: int | None = None
|
|
# FlowGRPO-style RL rollout: use the SDE sampler (with per-step log-prob capture) instead of the
|
|
# deterministic ODE serve sampler (PromptRL/UniRL §6). Default False ⇒ the serve path is unchanged.
|
|
sde_rollout: bool = False
|
|
sde_noise_scale: float = 0.7
|
|
|
|
|
|
class CaptureMode(str, Enum):
|
|
"""What the loop captures while running (design_v3 §9.4, §10)."""
|
|
NONE = "none" # serve forward: no-grad, no capture
|
|
BEHAVIOR = "behavior" # rollout forward: emit BehaviorRecord slices for RL/distill
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OutputSpec:
|
|
modalities: frozenset[str] = frozenset({"video"})
|
|
# streaming flags per modality: True | chunk size (ms / frames). {} => no streaming.
|
|
stream: dict[str, object] = field(default_factory=dict)
|
|
return_latents: bool = False
|
|
return_trajectory: bool = False # full denoise trajectory (parity / debug)
|
|
capture: CaptureMode = CaptureMode.NONE
|
|
|
|
@property
|
|
def streaming(self) -> bool:
|
|
return bool(self.stream)
|