From b3ddf6014daefc2b8e4538ea1a90c31f0b5a704a Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 16 Jun 2026 13:37:18 -0700 Subject: [PATCH] [feat] UniRL/PromptRL joint LM+generator RL stress test + designv4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- designv4.md | 402 ++++++++++++++++++++++++++++++ v2/loop/sampler.py | 28 +++ v2/models/__init__.py | 18 +- v2/models/backend.py | 76 ++++++ v2/models/unified/__init__.py | 9 + v2/models/unified/card.py | 100 ++++++++ v2/models/unified/program.py | 47 ++++ v2/models/wan21/loop.py | 31 ++- v2/request/params.py | 4 + v2/tests/test_unified_rl.py | 177 +++++++++++++ v2/training/__init__.py | 5 +- v2/training/methods/__init__.py | 2 + v2/training/methods/unified_rl.py | 222 +++++++++++++++++ v2/training/weight_sync.py | 12 +- 14 files changed, 1120 insertions(+), 13 deletions(-) create mode 100644 designv4.md create mode 100644 v2/models/unified/__init__.py create mode 100644 v2/models/unified/card.py create mode 100644 v2/models/unified/program.py create mode 100644 v2/tests/test_unified_rl.py create mode 100644 v2/training/methods/unified_rl.py diff --git a/designv4.md b/designv4.md new file mode 100644 index 000000000..c8c6d07fa --- /dev/null +++ b/designv4.md @@ -0,0 +1,402 @@ +# FastVideo v4 — The Unified Design (as built in `v2/`, validated by joint LM+generator RL) + +> **What this document is.** `design_v3.md` is the north star — the argument for a model-native runtime +> for the (recipe, runtime) era. This document, v4, is the *unified* design: it folds the v3 thesis +> together with **what was actually built and tested** in the `v2/` package, and with the one workload +> that most stressed the design — **UniRL/PromptRL-style joint LM + generator reinforcement learning** +> (arXiv 2510.17937; PromptRL, arXiv 2602.01382). v4 is not aspirational. Every structural claim below +> is backed by code in `v2/` and a passing test (91 tests, 14 files, two independent runners — `pytest` +> and a zero-dependency `v2/run_tests.py`). The neural forwards are toy numpy stand-ins (no GPU/torch in +> this environment, §12); the **control flow, contracts, scheduling, caching, parity gates, and training +> math are real**, which is the whole point of the (recipe, runtime) separation: swap the component +> factory for a torch adapter and the loops/policies/scheduler/caches/training are unchanged. + +--- + +## Table of contents + +1. The thesis (carried, with validation status) +2. The **three** signature ideas (v3 had two; the third is what the stress test proved) +3. Planes and dependency order (as enforced in `v2/`) +4. The Model Plane — `ModelCard` · `Loop` · `Program` +5. The driven-loop contract — `init/next/advance/finalize` +6. Runtime & scheduler — one WorkUnit, one budget, the interleave gate +7. Caches & the consistency ladder (the C2 split) +8. Training & RL on the same loops +9. **The UniRL/PromptRL joint-RL stress test (the centerpiece)** +10. Serving & fleet — our own stack, Dynamo optional +11. What v2 is and is not (honesty) +12. Package layout (actual) +13. Falsifiers, open questions, reference-synthesis delta + +--- + +## 1. The thesis (carried, with validation status) + +Three facts about video/omni generation dictate the architecture, unchanged from v3: + +- **A deployable model is a post-training artifact.** A usable model is *created* by training (distillation + for latency, QAT for precision, distillation+self-forcing for causal/world models). Every inference + capability is therefore a **(recipe, runtime) pair**: the weights and the loop that produced-and-assumes + them are one versioned object — a `ModelCard`. ✅ *Built:* `v2/card/` (`RecipeSpec` + `ParitySpec` on every + card; 5 phase-1/2 cards + the unified card). +- **Video/omni systems are loop systems.** The work is iteration — denoise steps, AR tokens, chunked + rollout, VAE tiles, optimizer steps, media chunks — not a single `forward()`. A runtime that reduces + everything to `forward()` cannot schedule, batch, cancel, stream, reserve memory for, or capture the + behavior of what actually runs. ✅ *Built:* `v2/loop/` driven-loop contract; `v2/runtime/` WorkUnit + scheduler; every denoise step and AR token is a runtime-visible WorkUnit. +- **Omni models share weights across loop types within one request.** ✅ *Built and tested:* the Cosmos3 and + BAGEL cards run **one resident MoT instance** driving an `ar_decode` loop then a `diffusion_denoise` loop + on shared weights (`v2/tests/test_omni.py`). + +The single invariant, restated and now demonstrable in code: + +```text +Model cards own components, loops, recipes, and parity. # v2/card/, v2/models/*/card.py +Programs compose loops into tasks. # v2/program/, v2/models/*/program.py +The scheduler executes loop steps as WorkUnits under one budget. # v2/runtime/ +Caches are correct by key, not by hope. # v2/cache/ (CacheKey) +Training records behavior on the same loops it serves. # v2/training/ (engine never imports training) +Deployment places and routes; it does not define semantics.# v2/deploy/, v2/serving/ +Products stream artifacts; they do not reach into the model.# v2/request/, v2/serving/server.py +``` + +## 2. The three signature ideas + +v3 named two. The joint-RL stress test (§9) forced articulation of a third — the one that turns the design +from "an elegant serving runtime" into "the substrate for content-creation RL." + +### 2.1 The (recipe, runtime) pair is a first-class, versioned, typed object + +A model is not a checkpoint; it is a `ModelCard` owning components, loops, the recipe that produced the +weights, and the parity contract binding train-forward to serve-forward. You cannot ship the weights +without the loop they assume, or change the loop without re-proving parity. ✅ `v2/card/card.py`. + +### 2.2 Driven loops: the model owns control flow, the runtime owns execution + +A loop is a state machine — `init → (next → advance)* → finalize` — where the model describes *the next +step it needs* and the runtime decides *when and with whom that step runs*. Per-request state lives in a +typed `LoopState`, never in module globals, so interleaving requests through one instance cannot smear +state. ✅ `v2/loop/contracts.py`; the **interleave parity gate** (`v2/parity/interleave_gate.py`) proves +serial == interleaved bit-for-bit — the bet loop-inversion lives or dies on. + +### 2.3 (new) One vocabulary spans every weight-sharing topology — including joint multi-expert RL + +The Card/Loop/Program split is not specialized to one sharing pattern. The **same primitives** express: + +| Topology | Components | Loops | Sharing | Card | +|---|---|---|---|---| +| Single diffusion | `transformer` | `diffusion_denoise` | — | `wan21` | +| **MoT omni** | **one** `transformer` | `ar_decode` **+** `diffusion_denoise` | both loops → **same** component | `cosmos3`, `bagel` | +| **Joint LM+gen RL** | `llm` **+** `transformer` (separate) | `ar_decode`→`llm`, `diffusion_denoise`→`transformer` | **disjoint** experts, one request | `unified` | + +MoT (one module, two loops) and UniRL (two experts, two loops) are *topological opposites*, yet both are +ordinary `shared_weight_components` bindings on `LoopSpec`s — **no new primitive**. The third idea: *the +weight-sharing graph is data on the card, not structure in the runtime.* This is what lets the same engine +serve a 1.3B single-DiT, a 30B MoT, and a Qwen+FLUX joint-RL recipe without three runtimes. ✅ Proven by +`v2/tests/test_unified_rl.py::test_two_separate_experts_not_a_shared_mot`. + +## 3. Planes and dependency order + +```text + Products: Python · CLI · OpenAI API (v2/serving/server.py) · fleet (v2/serving, v2/deploy) + │ (thin: validate intent, make requests, subscribe to streams) + Request / Session / Artifact / Stream v2/request/ + │ + Program Plane (typed loop programs) v2/program/, v2/models/*/program.py + │ + ┌──────────── Model Plane (CENTER) ────────────┐ v2/card/, v2/models/*/card.py + │ ModelCard: components · loops · recipe · │ + │ parity · capabilities · caches · parallelism │ + └───────────────────────┬───────────────────────┘ + │ + ┌────────────────────────┼────────────────────────┐ + │ Runtime / Scheduler │ Training / RL │ same loops, different capture + │ v2/runtime/ WorkUnits │ v2/training/ rollout · │ + │ + GPU-time budget │ reward · weight-sync │ + └────────────────────────┼────────────────────────┘ + │ + Memory · Cache · Transport · Parallelism v2/memory, v2/cache, v2/transport, v2/parallel + │ + Deployment / Fleet (DeploymentCard → LocalFleet | Dynamo) v2/deploy/, v2/serving/ +``` + +**Enforced boundary (the load-bearing rule):** `training/` imports `card`/`loop`/`runtime`/`models`; **the +engine never imports `training`.** Rollout is the serve forward *plus capture* — a client of the engine, +not a fork of it. ✅ enforced by import structure; `v2/training/__init__.py` documents it. + +## 4. The Model Plane — `ModelCard` · `Loop` · `Program` + +The three-way split (the question the user asked v2 to explain and then stress-test): + +- **`ModelCard`** — the *(recipe, runtime)* pair. Owns `components` (`ComponentSpec`: a lazy factory + a + `load_id` pointing at the real torch module + `required_for`/`optional_for`/`resident_for`), `loops` + (`LoopSpec`: kind, work-unit kind, `shared_weight_components`, `cache_policy`, `loop_factory`), `recipe` + (`RecipeSpec`: method, `assumes_loop`, `assumes_precision`, `consistency_required`), `parity` + (`ParitySpec`: consistency levels, interleave-required), `capabilities`, `caches`, `precision`, + `parallelism`. `card.validate()` checks every loop's components and cache policies exist. *The card is the + center; everything else is a view over it.* ✅ `v2/card/card.py`, `v2/models/*/card.py`. +- **`Loop`** — the runtime semantics: a driven state machine (§5). The model owns the transitions; the + runtime owns the lifecycle. +- **`Program`** — composes loops + component-nodes into a task DAG (`ComponentNode` for kernel-free seams + like text-encode/pack/vae-decode; `ModelLoopNode` to drive a loop to completion and bind its result into + a slot). ✅ `v2/program/`, `v2/models/*/program.py`. + +**Naming & rationale** (carried from the explanation v2 gave): "Card" because it is the spec sheet for a +deployable artifact — recipe on one side, runtime on the other. "Loop" because the atomic *runtime* +semantic in this domain is iteration, not `forward`. "Program" because a request is a *composition* of +loops, and composition deserves its own typed object distinct from any single loop. Versus FastVideo's +`ComposedPipeline` (stages own both semantics and orchestration, fused) and vllm-omni's pipeline-spec + +deploy-YAML (good split, but the diffusion loop is one opaque stage) — v2 separates *all three*: semantics +(Loop), composition (Program), and the deployable unit (Card). + +## 5. The driven-loop contract + +```python +class Loop: # v2/loop/contracts.py (duck-typed; protocols, not ABCs) + def init(self, req, model, ctx) -> LoopState # build per-request state (seeded rng, latents…) + def next(self, st) -> WorkPlan | Done # describe the next step (kernel-free: a thunk) + def advance(self, st, result) -> LoopState # fold the result; capture behavior if ROLLOUT + def finalize(self, st) -> LoopResult # produce outputs + metrics + behavior +``` + +`next` is **kernel-free** — it builds a `WorkPlan` carrying a `run(model, override=None)` thunk, a +`ResourceRequest` (compute-seconds, resident bytes, peak activation), a `ShapeSignature` (for batch +coalescing), and optional `emits` (stream chunks). The runtime decides *when* to call `run`. `advance` +folds the returned `StepResult` back into `LoopState` and, under the `ROLLOUT` profile only, appends a +behavior slice to `st.trajectory`. This is how **train ≡ serve**: identical `next`/`advance`, the profile +flips capture on. ✅ `WanDenoiseLoop` (`v2/models/wan21/loop.py`), `ARDecodeLoop` +(`v2/models/omni/ar_loop.py`). + +**Policies decompose the step body.** CFG, flow-shift, precision, and expert-routing are *policies* +composed into a loop, not branches inside it — `ClassicCFG.combine`, `FlowShiftPolicy.build_schedule`, +`PrecisionPolicy.cast`, `NoRouting/MoT routing`. CFG is a policy over *one* shared denoise body, so the +same body serves cond-only (guidance=1, the RL case) and full CFG (serving). ✅ `v2/loop/policies.py`, +`v2/tests/test_policies.py`. + +## 6. Runtime & scheduler — one WorkUnit, one budget + +Every step of every loop — a denoise step, an AR token, (later) a VAE tile or a transfer — is a `WorkUnit` +priced in **one currency: GPU-seconds**. Admission reserves a refundable `StepTicket` against a compute +budget and resident/peak memory against typed pools *before* a step runs; infeasible requests fail fast +(`AdmissionInfeasible`) rather than busy-spin; budget is refunded on completion. ✅ `v2/runtime/`, +`v2/memory/`, `v2/tests/test_admission.py`. + +The **interleave gate** is the correctness spine: `run_serial(reqs)` must equal `run_interleaved(reqs)` +bit-for-bit, where interleaved round-robins one unit per request per tick across *all* loops (including +the two-loop unified program). A batch-of-1 gate is structurally blind to cross-request state smearing; +this one is not. ✅ `v2/parity/interleave_gate.py`; passes for all phase-1 models, the omni MoT cards, and +the unified two-loop program (`test_unified_serve_interleave_parity`). + +## 7. Caches & the consistency ladder + +**Caches are correct by key.** `CacheKey` carries every output-semantic field: `model_id`, `component_id`, +**per-component** `weights_version`, `adapter_versions`, `precision`, and content-hashed inputs. Per-class +pools (`feature`, `paged_kv`, …) match the granularity reality. The per-component version is load-bearing: +a transformer weight sync bumps only the transformer's version and invalidates only its cache entries — the +**frozen text-encoder's feature cache survives**, so a K-sample RL group encodes its shared prompt *once*. +✅ `v2/cache/keys.py`, `v2/models/common.py::cached_text_encode`, `v2/tests/test_cache.py`. + +**The consistency ladder** (parity as a typed gate, declared on the card): + +- **C0** structural · **C1** deterministic-seed reproducibility · **C2** rollout↔train identity · + **C3** batch-invariance · **C4** artifact-quality/preference (named, not yet built). +- **The C2 split** (the rung v3 added, now demonstrated by *two* methods that sit on opposite halves): + - **C2 likelihood-free** — DiffusionNFT: no log-probs; identity is *behavioral* (seeded final sample + + prediction-space MSE). ✅ `v2/training/methods/diffusion_nft.py`. + - **C2 likelihood-based** — UniRL: a *per-step log-prob identity*. Recomputing each SDE step's Gaussian + log-prob under the unchanged rollout weights reproduces the captured log-prob (⇒ PPO ratio == 1) to + float32 trajectory-storage precision. ✅ `test_c2_likelihood_identity_at_rollout_weights` (max + `|Δlogp| ≈ 3.3e-7`). *This is the parity gate the PPO importance ratio rests on* — if it failed, the + ratio would be measuring a kernel gap instead of a policy change. + +## 8. Training & RL on the same loops + +```text +serve : request → program → loop → WorkUnits → artifacts +rollout : prompt batch → program → loop → WorkUnits → BehaviorRecords → rewards → update +``` + +The loop kernel is shared; the only difference is capture and training policy. The moat — and the place a +serving-only runtime structurally cannot follow — is that **the rollout forward *is* the serve forward plus +capture**: every serving optimization (distilled samplers, caches, step batching) is automatically a +rollout optimization, and there is one numerics surface (the ladder *measures* the gap rather than a +correction layer *papering over* it). + +`v2/training/` is ~16× smaller than `fastvideo/train/` not because it does less *method math* — it is a +faithful CPU port (NFT is line-for-line vs the source `diffusion_nft.py`) — but because it carries none of +the GPU/FSDP/checkpoint/distributed infrastructure. The five+one methods: + +| Method | Consistency | Roles (WeightSyncPlan) | Notes | +|---|---|---|---| +| `finetune` | C1 | student | plain flow-match regression | +| `dmd2` | C2 (free) | student + fake-score critic + teacher | distribution-matching distillation | +| `diffusion_nft` | C2 (free) | student + **old** (decay-blended behavior) + reference | samples from *old*, not student | +| `self_forcing` | C2 (free) | student + teacher | causal/chunked student | +| **`unified_rl`** | **C2 (based)** | **student (llm+transformer) + reference (both)** | **§9 — joint LM+gen RL** | + +`WeightSyncPlan` ships a **role**, not "the weights" (student / EMA / old_policy / reference / teacher / +critic), and now a **component scope** so a sync versions and cache-invalidates *one* expert in isolation. +✅ `v2/training/weight_sync.py`. + +## 9. The UniRL/PromptRL joint-RL stress test (the centerpiece) + +> *The user's framing: "the LLM and the generator may be updated simultaneously by training recipes — +> in this case, reinforcement learning. This is critical to the future of FastVideo and content creation."* + +UniRL-Zero / PromptRL train a **prompt-refiner LM** (Qwen) and a **flow generator** (FLUX) *together* under +one RL reward: the LM rewrites the prompt, the rewritten prompt conditions the generator, a single reward +on the generated image produces a **group-relative advantage**, and that one advantage drives **two +updates** — a token policy gradient on the LM and a FlowGRPO PPO update on the diffusion transformer. They +are *separate experts* (not weight-shared), updated *simultaneously*. This is the hardest thing v2 had been +asked to express: a training rollout that spans **two different loop types**, captures **two kinds of +log-prob**, and lands **two component updates** from **one reward**. + +### 9.1 What was built (`v2/models/unified/`, `v2/training/methods/unified_rl.py`) + +- **`unified` card** — two disjoint experts: `llm` (`ToyPromptRefiner`, a categorical policy over + refinement actions — the Qwen role) bound to `ar_decode`, and `transformer` (`ToyDiT`, the FLUX role) + bound to `diffusion_denoise`. Serve program: `text_encode → refine(ar_decode) → apply_refinement → + denoise(diffusion_denoise) → vae_decode`. +- **`flow_sde_step_with_logprob`** (`v2/loop/sampler.py`) — the FlowGRPO **SDE rollout sampler** (stochastic + step + per-step Gaussian log-prob), distinct from the deterministic ODE `flow_match_euler_step` used at + serve time. Gated on `DiffusionParams.sde_rollout` (default `False`) so **the serve path is byte-for-byte + unchanged** — the 82 pre-existing tests never see it. +- **`UnifiedRLMethod`** — the joint GRPO loop: per prompt, the LM *samples* a refinement action (capturing + its categorical log-prob); K diffusion samples roll out under SDE (capturing per-step log-probs), with + `num_skip_refinement` samples using the *original* prompt (partial refinement) so the advantage contrasts + refined vs unrefined; one reward → group-relative advantage A = clip((r−mean)/(std+ε), ±5); the advantage + drives **(a)** the LM via REINFORCE (`−A·logπ(a) + β·KL`) and **(b)** the DiT via FlowGRPO PPO + (`max(−A·ρ, −A·clip(ρ,1±ε)) + β·KL`, ρ = exp(logp_cur − logp_rollout)); two learning rates + (LLM_LR ≫ DIT_LR), two `WeightSyncPlan`s (`("llm",)` and `("transformer",)`). + +### 9.2 What it exercises that nothing else in v2 did + +1. **A rollout that spans two loop types** — `ar_decode` (refine) *and* `diffusion_denoise` (generate), not + one. The training rollout drives the *same* shared loops the engine serves. +2. **Dual log-prob capture** — a categorical action log-prob (LM) *and* a Gaussian per-step log-prob (SDE + diffusion), both feeding one advantage. +3. **Likelihood-based C2** (§7) — the per-step PPO identity, the counterpart to NFT's likelihood-free C2. +4. **A multi-component joint update** — two weight-sync plans, versioned and cache-scoped *independently* + (the LM sync must not flush the frozen text-encoder's feature cache, and does not). +5. **Rollout sampler ≠ serve sampler** — a controlled §9.4 divergence (SDE for exploration at train time, + ODE for determinism at serve time), gated so it cannot leak into serving. +6. **A flag for prompt-only vs joint** — `joint=False` freezes the generator (ODE rollout, no DiT update); + only the LM learns. This is PromptRL's prompt-only ablation, as a one-line mode. + +### 9.3 What it proved (the test results) + +- **The LM actually learns** to pick the reward-favored refinement: `P(target action)` 0.125 → 0.95 (joint) + and 0.90 (prompt-only) over 40 iterations. +- **The generator actually moves** in joint mode (finite FlowGRPO grads, PPO ratio ≈ 1.0, KL finite) and is + **byte-for-byte frozen** in prompt-only mode. +- **The two experts version independently** (distinct weight versions; `text_encoder` stays at `v0`). +- **The two-loop serve program passes the interleave parity gate** unchanged. +- ✅ All in `v2/tests/test_unified_rl.py` (9 tests); full suite **91 passed**, zero regressions. + +### 9.4 The verdict: the design *held* + +The stress test required **no new primitive** — only (a) a second sampler in the sampler *library*, (b) two +boolean fields on `DiffusionParams`, (c) a gated capture branch in the existing `WanDenoiseLoop`'s +`advance`, (d) a new toy component, a new card, and a new method, and (e) one generalization: +`WeightSyncPlan` gained a `components` scope (a strict improvement that also makes the §7.1 cache-scoping +guarantee explicit). The loop contract, the runtime, the scheduler, the cache keys, the parity gate, the +program model, and the training-plane boundary were **untouched**. *A joint multi-expert RL recipe is a +new card + a new method, not a new engine* — which is exactly the claim §2.3 makes, now with a test behind +it. The refinements the test forced (the component-scoped weight sync; the explicit likelihood-based C2 +rung; `guidance_scale=1` at RL time to keep the PPO recompute faithful) are improvements, not patches. + +## 10. Serving & fleet — our own stack, Dynamo optional + +Per the explicit instruction "*don't completely rely on Dynamo; we still need our own version*," v2 ships a +complete serving/fleet stack and treats Dynamo as one optional backend: + +- **`AsyncEngine`** with role/stage pools, an `OpenAI`-compatible server on stdlib asyncio + (`v2/serving/server.py`, with HTTP timeouts/size limits and SSE streaming), and a `DisaggregatedRunner` + proven **bit-identical to inline** execution. +- **Connectors** with `chunk_ready` readiness + **credit-based flow control** (sglang-omni's model, not + vllm-omni's readiness-parking). +- **`LocalFleet`** (our own placement/routing) and a `DynamoWorkerAdapter` (`v2/deploy/`) — the core exports + a `DeploymentCard` + cost model *to* a fleet planner; it is never *defined by* one. +- Hardening from the serving review is in-tree (`v2/tests/test_serving_fixes.py`): pool-capacity leaks on + cancel/error, duplicate-request-id deadlock, credit leaks, cost-model aliasing, SSE aclose on disconnect. + +## 11. What v2 is and is not (honesty) + +- **Is:** a faithful, CPU-testable realization of the v3 contracts — Card/Loop/Program, WorkUnit scheduling, + typed caches, the consistency ladder, train≡serve, omni MoT, joint LM+gen RL, and a full serving/fleet + stack — ~7,500 LOC of package + ~1,400 LOC of tests, 91 passing. +- **Is not:** a GPU runtime. The neural forwards are deterministic numpy toys (`v2/models/backend.py`); + there are no real weights, kernels, FSDP, or NCCL. C3 (batch-invariance) and C4 (quality/preference) are + named, not built. **The deferred next step (the user's standing request) is to port the component + factories to torch adapters** wrapping the real `fastvideo.models` modules — at which point, by + construction, the loops/policies/scheduler/caches/parity/training code does not change. + +## 12. Package layout (actual `v2/`) + +```text +v2/ + card/ ModelCard, ComponentSpec, LoopSpec, RecipeSpec, ParitySpec, instance, load_card + loop/ contracts (LoopState/WorkPlan/StepResult/Done), policies (cfg/flowshift/precision/routing), sampler + program/ ComponentNode, ModelLoopNode, Program, compiler + runtime/ Engine (run/run_serial/run_interleaved), scheduler, admission, StepTicket + cache/ keys (CacheKey, content_hash), per-class pools (feature, paged_kv) + memory/ allocator, reservations, refundable budget + transport/ manifests + backends (in-proc for the mini) + parallel/ ParallelPlan, mesh, validation + parity/ interleave_gate, consistency ladder, compare_outputs + extend/ observers, interceptors + request/ requests, params (DiffusionParams + sde_rollout), tasks (TaskType), streams, artifacts + models/ backend (toy components) + common (cached_text_encode) + wan21/ ltx2/ wan_causal/ # phase 1: T2V, distilled, self-forcing student + omni/ cosmos3/ bagel/ # phase 2: MoT (one module, two loops) + unified/ # §9: two experts, two loops, joint RL + training/ rollout, behavior, rewards, weight_sync, methods/{finetune,dmd2,diffusion_nft,self_forcing,unified_rl} + serving/ AsyncEngine, pools, DisaggregatedRunner, connectors, OpenAI server + deploy/ DeploymentCard, LocalFleet, DynamoWorkerAdapter + tests/ 14 files, 91 tests run via `pytest v2/tests/` OR `python3 v2/run_tests.py` +``` + +**Enforced boundaries:** `card/` imports no product/runtime; `runtime/` executes `card/` loops but defines +no semantics; `training/` requires behavior records but forks no loop and is never imported by the engine; +`serving/`/`deploy/` place and route but define no semantics. `parity/` is a first-class package, not a test +folder — it is how the (recipe, runtime) pair is kept honest. + +## 13. Falsifiers, open questions, reference-synthesis delta + +The v3 falsifiers stand (step-level scheduling must beat request-level on a real duty-cycle trace; the +general WorkUnit scheduler may be over-general for non-step kinds; cost-model admission is a modeling bet; +the clean-slate premise is out of scope by request; quality/C4 is unmeasured). The stress test **retired one +open question** — *"does the Card/Loop/Program split generalize beyond serving and MoT, or will joint +multi-expert RL force a redesign?"* — answered: it generalizes, no redesign (§9.4). Two new questions it +**opened**: + +- **Throughput of two-loop rollouts.** A joint rollout interleaves an `ar_decode` and a `diffusion_denoise` + loop per sample; the homogeneous-group batching win (K identical diffusion samples) is intact, but the LM + refine step is a serial prefix. *Open:* on GPU, does the refine step want its own role-pool (LM served on + separate hardware from the generator, à la disaggregated prefill), and does the WorkUnit cost model price + the two loop kinds correctly enough to co-schedule them? — measure on the torch port. +- **BehaviorRecord size for joint RL.** Dual log-prob capture (token + per-step SDE) plus the SDE + trajectory is heavier than NFT's seeded-sample capture. *Open:* confirm the opt-in instrument sizing holds + for the LM+gen case at real resolution. + +**Reference-synthesis delta (added to v3's table):** + +| Source | Take | Constrain / reject | +|---|---|---| +| **UniRL-Zero / PromptRL** (arXiv 2510.17937, 2602.01382) | Joint LM-refiner + flow-generator RL under one reward; FlowGRPO SDE-prefix/ODE-tail with per-step log-probs; group-relative advantage → token-PG **and** PPO; partial-refinement contrast; prompt-only vs joint ablation; two LRs | Expresses the two experts and the joint update as a **card + method**, not a bespoke trainer; the SDE rollout sampler is a *library* entry gated behind `sde_rollout`, never the serve path; the PPO ratio rests on the **likelihood-based C2** parity gate, not on hope | + +--- + +## Final position + +v3 argued that the atomic unit is the (recipe, runtime) pair and that a model-native, loop-driven runtime is +what the post-training era requires. v4 is the same argument, now *built and measured* in `v2/`: the +contracts hold under serving, under MoT omni, and — the hard case — under **joint LM + generator +reinforcement learning**, where one reward simultaneously updates two separate experts driven by two +different loop types in one request. That workload, which UniRL/PromptRL identify as central to the future +of content creation, fit inside the design as *a new card and a new method* with no new runtime primitive. +The design's bet is therefore not that it is elegant — it is that **the weight-sharing topology and the +training recipe are data on a card, so a new frontier recipe is a card, not a rewrite.** The remaining work +is the GPU port, where, by construction, the loops, scheduler, caches, parity, and training plane do not +change — only the component factories do. diff --git a/v2/loop/sampler.py b/v2/loop/sampler.py index af520f07f..2cca590fe 100644 --- a/v2/loop/sampler.py +++ b/v2/loop/sampler.py @@ -38,3 +38,31 @@ def flow_match_euler_step(x_t: np.ndarray, velocity: np.ndarray, def add_noise(x0: np.ndarray, noise: np.ndarray, sigma: float) -> np.ndarray: """Forward flow-match interpolation: x_σ = (1-σ)·x0 + σ·noise.""" return (1.0 - sigma) * x0 + sigma * noise + + +def flow_sde_step_with_logprob(x_t, velocity, sigma_t: float, sigma_next: float, *, + noise=None, prev_sample=None, noise_scale: float = 0.7): + """FlowGRPO-style stochastic step + Gaussian log-prob (PromptRL/UniRL §6; FlowGRPO lineage). + + The RL *rollout* sampler (vs the deterministic ODE ``flow_match_euler_step`` used at serve time). + Injects noise for exploration and returns the log-prob of the realized sample under the step's + Gaussian — the quantity the FlowGRPO PPO ratio uses. ``prev_sample=None`` samples a new step; + passing ``prev_sample`` (the rollout sample) recomputes its log-prob under the *current* velocity + (the ratio's numerator at update time). Returns ``(prev_sample, log_prob, mean, eff_std)``. + """ + x_t = np.asarray(x_t, dtype=np.float64) + velocity = np.asarray(velocity, dtype=np.float64) + s = min(float(sigma_t), 0.9999) + dt = float(sigma_next) - float(sigma_t) # negative (σ decreases) + std = float(np.sqrt(s / (1.0 - s)) * noise_scale) + denom = 2.0 * max(s, 1e-6) + mean = x_t * (1.0 + std ** 2 / denom * dt) + velocity * (1.0 + std ** 2 * (1.0 - s) / denom) * dt + eff_std = max(std * np.sqrt(max(-dt, 1e-12)), 1e-6) + if prev_sample is None: + n = noise if noise is not None else np.zeros_like(x_t) + prev_sample = mean + eff_std * np.asarray(n, dtype=np.float64) + prev_sample = np.asarray(prev_sample, dtype=np.float64) + var = eff_std ** 2 + log_prob = float(np.mean(-((prev_sample - mean) ** 2) / (2.0 * var) + - np.log(eff_std) - 0.5 * np.log(2.0 * np.pi))) + return prev_sample.astype("float32"), log_prob, mean.astype("float32"), float(eff_std) diff --git a/v2/models/__init__.py b/v2/models/__init__.py index 98d731c78..a5599a991 100644 --- a/v2/models/__init__.py +++ b/v2/models/__init__.py @@ -10,6 +10,7 @@ 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 @@ -19,7 +20,8 @@ __all__ = [ "build_wan_causal_card", "build_wan_causal_program", "build_cosmos3_card", "build_cosmos3_program", "build_bagel_card", "build_bagel_program", - "build_default_engine", "build_omni_engine", + "build_unified_card", "build_unified_program", + "build_default_engine", "build_omni_engine", "build_unified_engine", ] _BUILDERS = [ @@ -63,3 +65,17 @@ def build_omni_engine(engine: Any = None) -> Any: 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 diff --git a/v2/models/backend.py b/v2/models/backend.py index 7418155af..cec8f81af 100644 --- a/v2/models/backend.py +++ b/v2/models/backend.py @@ -145,6 +145,82 @@ class ToyMoTDiT(ToyDiT): return (rng.standard_normal((4, 8)) * 0.1).astype("float32") +class ToyPromptRefiner: + """Toy LM prompt-refiner — UniRL/PromptRL's Qwen role (a *separate expert*, not MoT-shared). + + A categorical policy over a small set of refinement "actions" (stand-in for the LM emitting a + refined prompt). Sampling an action picks an embedding offset that shifts the diffusion + conditioning; the realized reward then drives the policy by REINFORCE. This is the *second + trainable expert* in the unified-RL stress test: one reward → a token-policy-gradient update + here AND a FlowGRPO update on the DiT, under one advantage (design_v3 §10; PromptRL §6). + + Real gradient: loss = −A·logπ(a); ∇_logits = −A·(onehot(a) − softmax). So the toy LM actually + learns to prefer the reward-favored action — the LM-side analogue of ToyDiT.mse_grad_step. + """ + + def __init__(self, n_actions: int = 8, dim: int = TEXT_DIM, seq: int = TEXT_SEQ, seed: int = 0): + self.n = int(n_actions) + self.dim, self.seq = int(dim), int(seq) + self.logits = np.zeros(self.n, dtype="float64") # the trainable policy params + rng = np.random.default_rng(seed) + # each action ⇒ a fixed conditioning offset (the "content" of that refined prompt) + self._emb = (rng.standard_normal((self.n, self.seq, self.dim)) * 0.1).astype("float32") + + def _probs(self) -> np.ndarray: + z = self.logits - self.logits.max() + e = np.exp(z) + return e / e.sum() + + def sample_refinement(self, rng) -> tuple[int, float]: + """Sample an action (exploration) and return (action, logπ(action)).""" + p = self._probs() + a = int(rng.choice(self.n, p=p)) + return a, float(np.log(p[a] + 1e-12)) + + def argmax_refinement(self) -> tuple[int, float]: + a = int(np.argmax(self.logits)) + return a, float(np.log(self._probs()[a] + 1e-12)) + + def action_logprob(self, action: int) -> float: + return float(np.log(self._probs()[int(action)] + 1e-12)) + + def ar_forward(self, tokens) -> int: + """und-pathway hook so the declared ``ar_decode`` loop is runnable (greedy ⇒ interleave-safe).""" + return int(np.argmax(self.logits)) + + def refined_embed(self, base_embed, action: int) -> np.ndarray: + """Apply the chosen refinement to the base text embedding → the diffusion conditioning.""" + base = np.zeros((self.seq, self.dim), dtype="float32") if base_embed is None \ + else np.asarray(base_embed, dtype="float32") + return (base + self._emb[int(action)]).astype("float32") + + def pg_step(self, action: int, advantage: float, lr: float) -> float: + """REINFORCE on the categorical policy. Returns grad_norm (the per-method regression signal).""" + p = self._probs() + onehot = np.zeros(self.n, dtype="float64") + onehot[int(action)] = 1.0 + grad = -float(advantage) * (onehot - p) # ∇_logits (−A·logπ(a)) + self.logits = self.logits - float(lr) * grad # descent on −A·logπ ⇒ ascent on A·logπ + return float(np.linalg.norm(grad)) + + def kl_to(self, other: "ToyPromptRefiner") -> float: + """KL(π_self ‖ π_other) over the categorical — the LM-PG reference-KL penalty term.""" + p, q = self._probs(), other._probs() + return float(np.sum(p * (np.log(p + 1e-12) - np.log(q + 1e-12)))) + + def clone(self) -> "ToyPromptRefiner": + c = ToyPromptRefiner.__new__(ToyPromptRefiner) + c.n, c.dim, c.seq = self.n, self.dim, self.seq + c.logits, c._emb = self.logits.copy(), self._emb.copy() + return c + + def copy_from(self, other: "ToyPromptRefiner") -> None: + self.logits = other.logits.copy() + + def blend_from(self, other: "ToyPromptRefiner", decay: float) -> None: + self.logits = decay * self.logits + (1.0 - decay) * other.logits + + class ToyVAE: """Tiny deterministic VAE. encode: video→latent (mean-pool + channel proj); decode: latent→video.""" diff --git a/v2/models/unified/__init__.py b/v2/models/unified/__init__.py new file mode 100644 index 000000000..ef37666fa --- /dev/null +++ b/v2/models/unified/__init__.py @@ -0,0 +1,9 @@ +"""Unified LM+generator model (UniRL/PromptRL): a prompt-refiner LM expert + a flow generator expert, +both trainable under one RL reward. The stress test that the Card/Loop/Program split holds for +joint multi-expert RL (design_v3 §4, §10).""" +from __future__ import annotations + +from .card import N_REFINE_ACTIONS, build_unified_card +from .program import apply_refinement_node, build_unified_program + +__all__ = ["build_unified_card", "build_unified_program", "apply_refinement_node", "N_REFINE_ACTIONS"] diff --git a/v2/models/unified/card.py b/v2/models/unified/card.py new file mode 100644 index 000000000..e9e30fded --- /dev/null +++ b/v2/models/unified/card.py @@ -0,0 +1,100 @@ +"""Unified LM+generator ModelCard — the UniRL/PromptRL stress test (design_v3 §4, §10). + +This is the topological *opposite* of the Cosmos3 omni card, expressed in the SAME vocabulary: + + * Cosmos3 MoT: ONE ``transformer`` bound to BOTH ``ar_decode`` and ``diffusion_denoise`` + (shared weights — splitting them severs shared state). + * Unified UniRL: TWO separate experts — an ``llm`` prompt-refiner (bound to ``ar_decode``) and a + ``transformer`` flow generator (bound to ``diffusion_denoise``). The LM refines the prompt; its + refinement conditions the generator. They do NOT share weights (UniRL keeps Qwen and FLUX as + distinct experts), yet BOTH are trainable policies updated under one RL reward. + +The stress-test claim: the Card/Loop/Program split that expresses MoT weight-sharing *also* expresses +two-separate-experts-joint-RL with no new primitive — only a second ``shared_weight_components`` target +and a method (``unified_rl``) that emits two WeightSyncPlans. The loops, caches, parity gate, rollout +capture, and serve path are unchanged. +""" +from __future__ import annotations + +from ..._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind +from ...card import ( + CacheContract, + CapabilityMatrix, + ComponentSpec, + CostModel, + LoopSpec, + ModelCard, + ParallelismContract, + ParitySpec, + PrecisionContract, + RecipeSpec, +) +from ...loop.policies import ClassicCFG, FlowShiftPolicy, NoRouting, PrecisionPolicy +from ...parallel import ParallelPlan +from ..backend import ToyDiT, ToyPromptRefiner, ToyTextEncoder, ToyVAE, _seed_from +from ..omni import ARDecodeLoop +from ..wan21.loop import WanDenoiseLoop + +N_REFINE_ACTIONS = 8 + + +def build_unified_card(model_id: str = "unirl-qwenflux") -> ModelCard: + seed = _seed_from(model_id) + ar_cost = CostModel(kind=WorkUnitKind.AR_TOKEN, base_seconds=5e-5, per_unit_seconds=1e-7) + dn_cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) + cfg, flow = ClassicCFG(), FlowShiftPolicy(shift=3.0) + precision, expert = PrecisionPolicy(), NoRouting("transformer") + + def refine_factory(): + # the LM expert is the ar loop's "transformer" — a *different* component than the generator + return ARDecodeLoop(loop_id="ar_decode", transformer_id="llm", cost=ar_cost, max_tokens=1) + + def denoise_factory(): + return WanDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow, + precision=precision, expert=expert, cost=dn_cost) + + components = { + "text_encoder": ComponentSpec("text_encoder", kind="text_encoder", + factory=lambda inst: ToyTextEncoder(), + required_for={"reason", "t2v"}), + "llm": ComponentSpec( # the prompt-refiner expert (Qwen role) + "llm", kind="dit", + load_id="fastvideo.models.llm:QwenPromptRefiner", + factory=lambda inst: ToyPromptRefiner(n_actions=N_REFINE_ACTIONS, seed=seed + 1), + resident_for=["ar_decode"], required_for={"reason", "t2v"}), + "transformer": ComponentSpec( # the flow generator expert (FLUX role) + "transformer", kind="dit", + load_id="fastvideo.models.dits.flux:FluxTransformer", + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["diffusion_denoise"], required_for={"t2v"}), + "vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + } + loops = { + "ar_decode": LoopSpec("ar_decode", kind=LoopKind.AR_DECODE, work_unit_kind=WorkUnitKind.AR_TOKEN, + step_cost_model=ar_cost, shared_weight_components=["llm"], + cache_policy=["paged_kv"], loop_factory=refine_factory), + "diffusion_denoise": LoopSpec("diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=dn_cost, + shared_weight_components=["transformer"], cache_policy=["feature"], + loop_factory=denoise_factory), + } + card = ModelCard( + model_id=model_id, family="unified", components=components, loops=loops, + capabilities=CapabilityMatrix.of( + Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, + Capability.VAE_DECODE, Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise", + assumes_precision="float32", consistency_required=ConsistencyLevel.C2), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1, ConsistencyLevel.C2], + interleave_required=True), + caches={ + "feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True), + "paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12, + reuse_across_requests=False), + }, + precision=PrecisionContract(default_dtype="float32", training_precision="float32"), + parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], + default_plan=ParallelPlan.single()), + ) + return card.validate() diff --git a/v2/models/unified/program.py b/v2/models/unified/program.py new file mode 100644 index 000000000..81b5d734c --- /dev/null +++ b/v2/models/unified/program.py @@ -0,0 +1,47 @@ +"""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() diff --git a/v2/models/wan21/loop.py b/v2/models/wan21/loop.py index cbdac3508..b27cb4d85 100644 --- a/v2/models/wan21/loop.py +++ b/v2/models/wan21/loop.py @@ -21,7 +21,7 @@ from ...loop.contracts import ( StepResult, WorkPlan, ) -from ...loop.sampler import flow_match_euler_step +from ...loop.sampler import flow_match_euler_step, flow_sde_step_with_logprob from ...request.streams import StreamChunk from ..backend import LATENT_CHANNELS @@ -59,6 +59,9 @@ class WanDenoiseLoop: st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale) st.scratch["stream_video"] = bool(req.outputs.stream.get("video")) + # FlowGRPO RL rollout: switch the sampler to SDE-with-logprob (else deterministic ODE serve). + st.scratch["sde"] = bool(getattr(req.diffusion, "sde_rollout", False)) + st.scratch["sde_noise_scale"] = float(getattr(req.diffusion, "sde_noise_scale", 0.7)) st.plugin_state["cfg"] = {} return st @@ -75,6 +78,7 @@ class WanDenoiseLoop: pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] scale = st.scratch["guidance_scale"] cfg, precision = self.cfg, self.precision + sde, noise_scale, rng = st.scratch.get("sde", False), st.scratch.get("sde_noise_scale", 0.7), st.rng def run(model, override=None): if override is not None and "noise_pred" in override: @@ -83,7 +87,15 @@ class WanDenoiseLoop: dit = model.component(expert_id) preds = {b: dit(x, pe if b == "cond" else ne, sigma_t) for b in branches} velocity = cfg.combine(preds, scale, sctx, cfg_state) - x_next = flow_match_euler_step(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next) + velocity = precision.cast(velocity) + if sde: # FlowGRPO rollout: stochastic + log-prob + noise = rng.standard_normal(x.shape) + x_next, logp, _m, _s = flow_sde_step_with_logprob( + precision.cast(x), velocity, sigma_t, sigma_next, noise=noise, noise_scale=noise_scale) + return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), + "latents": x_next.astype("float32"), + "sde_logprob": logp, "prev": np.asarray(x, dtype="float32")}) + x_next = flow_match_euler_step(precision.cast(x), velocity, sigma_t, sigma_next) return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")}) @@ -106,11 +118,16 @@ class WanDenoiseLoop: def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["video"] = result.output["latents"] if st.profile == ExecutionProfile.ROLLOUT: - st.trajectory.append({ - "step": st.step_idx, "sigma": st.sigmas[st.step_idx], - "velocity": np.asarray(result.output["noise_pred"]).copy(), - "latents": np.asarray(st.latents["video"]).copy(), - }) + i = st.step_idx + rec = {"step": i, "sigma": st.sigmas[i], + "velocity": np.asarray(result.output["noise_pred"]).copy(), + "latents": np.asarray(st.latents["video"]).copy()} + if "sde_logprob" in result.output: # FlowGRPO: capture the PPO log-prob slice + rec.update(sde_logprob=result.output["sde_logprob"], + prev=np.asarray(result.output["prev"]).copy(), + sample=np.asarray(result.output["latents"]).copy(), + sigma_t=st.sigmas[i], sigma_next=st.sigmas[i + 1]) + st.trajectory.append(rec) st.step_idx += 1 return st diff --git a/v2/request/params.py b/v2/request/params.py index 7f0cc031e..421974b41 100644 --- a/v2/request/params.py +++ b/v2/request/params.py @@ -37,6 +37,10 @@ class DiffusionParams: 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): diff --git a/v2/tests/test_unified_rl.py b/v2/tests/test_unified_rl.py new file mode 100644 index 000000000..01f5e1351 --- /dev/null +++ b/v2/tests/test_unified_rl.py @@ -0,0 +1,177 @@ +"""UniRL/PromptRL joint LM+generator RL — the design stress test (design_v3 §4, §10; arXiv 2510.17937). + +These tests assert the *design holds*: a joint two-expert RL recipe fits inside the existing +Card/Loop/Program + training-plane vocabulary with no new primitive. Specifically: + + * the unified card is two SEPARATE experts (llm + transformer) on two loops — the topological + opposite of the Cosmos3 MoT card (one shared module on two loops), same vocabulary; + * the serve path is unchanged ODE; SDE-with-logprob capture is gated on ``sde_rollout``; + * likelihood-BASED C2: a per-step log-prob identity (the other half of NFT's likelihood-free C2); + * one reward → two updates (LM token-PG + DiT FlowGRPO-PPO) under one group-relative advantage; + * two WeightSyncPlans version + cache-scope the experts independently; + * prompt-only mode freezes the generator; + * the two-loop serve program still passes the interleave parity gate. +""" +from __future__ import annotations + +import numpy as np + +from v2._enums import ExecutionProfile, LoopKind +from v2.loop.sampler import flow_sde_step_with_logprob +from v2.models import build_unified_engine +from v2.models.common import cached_text_encode +from v2.models.unified import build_unified_card +from v2.parity import assert_interleave_parity +from v2.request import DiffusionParams, TaskType, make_request +from v2.training.methods import build_unified_rl +from v2.training.rollout import rollout_loop + +C2_TOL = 1e-5 # float32 trajectory-storage precision (bf16/fp32 replay buffers in practice) + + +def _inst(eng): + return eng._registry["unirl-qwenflux"][0] + + +def _rollout(inst, prompt, *, sde, seed=7, steps=4): + emb = cached_text_encode(inst, prompt) + neg = cached_text_encode(inst, "") + req = make_request(TaskType.T2V, inst.card.model_id, prompt, + diffusion=DiffusionParams(num_steps=steps, seed=seed, guidance_scale=1.0, + sde_rollout=sde)) + return rollout_loop(inst, "diffusion_denoise", req, + slots={"text_embeds": emb, "neg_text_embeds": neg}, + profile=ExecutionProfile.ROLLOUT), emb + + +# --- card / program shape --------------------------------------------------------- # + +def test_unified_card_validates_and_serves(): + eng = build_unified_engine() + req = make_request(TaskType.T2V, "unirl-qwenflux", "a red fox in snow") + res = eng.run(req) + assert res.artifacts["video"].frames is not None + assert np.asarray(res.artifacts["video"].frames).ndim == 4 + + +def test_two_separate_experts_not_a_shared_mot(): + """UniRL keeps the refiner and generator as DISTINCT experts — the two loops bind different + components. (Contrast Cosmos3, where ar_decode + diffusion_denoise bind the SAME transformer.)""" + card = build_unified_card() + assert card.loops["ar_decode"].kind == LoopKind.AR_DECODE + assert card.loops["diffusion_denoise"].kind == LoopKind.DIFFUSION_DENOISE + assert card.loops["ar_decode"].shared_weight_components == ["llm"] + assert card.loops["diffusion_denoise"].shared_weight_components == ["transformer"] + eng = build_unified_engine() + inst = _inst(eng) + assert inst.component("llm") is not inst.component("transformer") # genuinely separate weights + + +# --- serve ODE vs rollout SDE (gated, §9.4) --------------------------------------- # + +def test_serve_is_ode_rollout_is_sde(): + eng = build_unified_engine() + inst = _inst(eng) + res_ode, _ = _rollout(inst, "a fox", sde=False) + res_sde, _ = _rollout(inst, "a fox", sde=True) + assert all("sde_logprob" not in r for r in res_ode.behavior) # serve sampler: no log-probs + sde_steps = [r for r in res_sde.behavior if "sde_logprob" in r] + assert sde_steps, "SDE rollout must capture per-step log-probs" + for r in sde_steps: # the PPO slice is complete + assert {"sde_logprob", "prev", "sample", "sigma_t", "sigma_next"} <= set(r) + + +def test_sde_rollout_is_deterministic_given_seed(): + eng = build_unified_engine() + inst = _inst(eng) + a, _ = _rollout(inst, "a fox", sde=True, seed=123) + b, _ = _rollout(inst, "a fox", sde=True, seed=123) + assert np.array_equal(a.outputs["latents"], b.outputs["latents"]) # seeded ⇒ bit-identical + assert [r["sde_logprob"] for r in a.behavior if "sde_logprob" in r] == \ + [r["sde_logprob"] for r in b.behavior if "sde_logprob" in r] + + +def test_c2_likelihood_identity_at_rollout_weights(): + """Likelihood-BASED C2: recomputing each step's log-prob under the UNCHANGED rollout weights + reproduces the captured log-prob (⇒ PPO ratio == 1). This is the parity gate this method + lives on, the counterpart to NFT's likelihood-FREE C2.""" + eng = build_unified_engine() + inst = _inst(eng) + res, emb = _rollout(inst, "a fox", sde=True, seed=7) + dit = inst.component("transformer") + max_diff = 0.0 + for r in res.behavior: + if "sde_logprob" not in r: + continue + v = dit(r["prev"], emb, r["sigma_t"]) + _, logp, _, _ = flow_sde_step_with_logprob(r["prev"], v, r["sigma_t"], r["sigma_next"], + prev_sample=r["sample"]) + max_diff = max(max_diff, abs(logp - r["sde_logprob"])) + assert max_diff < C2_TOL, f"C2 likelihood identity broken: {max_diff:.2e}" + + +# --- joint vs prompt-only RL ------------------------------------------------------ # + +def _train(joint, target, iters=40, **kw): + m = build_unified_rl(build_unified_card(), joint=joint, target_action=target, + num_samples_per_prompt=4, num_skip_refinement=1, **kw) + batch = {"prompts": ["a fox", "blue sky", "a city at night"], "seeds": [1, 2, 3]} + p_start = float(m.llm._probs()[target]) + last = None + for it in range(iters): + _, last = m.managed_train_step(batch, it) + return m, p_start, float(m.llm._probs()[target]), last + + +def test_joint_rl_updates_both_experts(): + m, p0, p1, last = _train(joint=True, target=3) + # 1) the LM learned to pick the reward-favored refinement action + assert p1 > p0 + 0.2, f"LM did not learn the target action: {p0:.3f} -> {p1:.3f}" + # 2) the generator moved (FlowGRPO PPO produced real grads) + assert last["grad_norm/transformer"] > 0.0 + assert "dit_pg_loss" in last and np.isfinite(last["dit_pg_loss"]) + # 3) PPO ratio is near 1 (small on-policy steps), KL finite + assert 0.5 < last["ppo_ratio_mean"] < 2.0 + assert np.isfinite(last["kl_div_loss"]) + # 4) two experts versioned INDEPENDENTLY; the frozen text-encoder cache is NOT invalidated + inst = m.student + assert last["llm_weights_version"] != last["transformer_weights_version"] + assert inst.version_of("text_encoder") == "v0" # §7.1 cache scope preserved + assert inst.version_of("llm") != "v0" and inst.version_of("transformer") != "v0" + + +def test_prompt_only_freezes_the_generator(): + card = build_unified_card() + m = build_unified_rl(card, joint=False, target_action=5, + num_samples_per_prompt=4, num_skip_refinement=1) + w0 = m.dit.w_x.copy() + batch = {"prompts": ["a fox", "blue sky", "a city"], "seeds": [1, 2, 3]} + p0 = float(m.llm._probs()[5]) + last = None + for it in range(40): + _, last = m.managed_train_step(batch, it) + p1 = float(m.llm._probs()[5]) + assert p1 > p0 + 0.2, "prompt-only: LM still must learn" + assert np.array_equal(w0, m.dit.w_x), "prompt-only: generator weights must be frozen" + assert last["transformer_weights_version"] == "frozen" + assert "dit_pg_loss" not in last + + +def test_one_reward_drives_both_via_one_advantage(): + """The structural claim: a single reward per sample → one group-relative advantage → BOTH a + token-PG (LM) and a FlowGRPO-PPO (DiT) update. Joint moves the DiT; prompt-only does not — + yet the LM learns in both, from the same advantage.""" + _, _, p1_joint, _ = _train(joint=True, target=2) + _, _, p1_prompt, _ = _train(joint=False, target=2) + assert p1_joint > 0.4 and p1_prompt > 0.4 # LM learns either way + + +# --- the two-loop serve program still passes the core parity gate ----------------- # + +def test_unified_serve_interleave_parity(): + eng = build_unified_engine() + reqs = [make_request(TaskType.T2V, "unirl-qwenflux", p, + diffusion=DiffusionParams(num_steps=4, seed=s)) + for p, s in [("alpha", 11), ("beta", 22), ("alpha", 11)]] + divs = assert_interleave_parity(eng, reqs) + assert not divs, f"unified two-loop program failed interleave parity: {divs}" diff --git a/v2/training/__init__.py b/v2/training/__init__.py index 758b0dc48..886fe55b1 100644 --- a/v2/training/__init__.py +++ b/v2/training/__init__.py @@ -12,10 +12,12 @@ from .methods import ( 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 @@ -23,7 +25,8 @@ from .weight_sync import WeightRole, WeightSyncPlan __all__ = [ "TrainingMethod", "FineTuneMethod", "DMD2Method", "DiffusionNFTMethod", "SelfForcingMethod", - "build_finetune", "build_dmd2", "build_diffusion_nft", "build_self_forcing", + "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", ] diff --git a/v2/training/methods/__init__.py b/v2/training/methods/__init__.py index 0e166a2c8..22f264c94 100644 --- a/v2/training/methods/__init__.py +++ b/v2/training/methods/__init__.py @@ -6,6 +6,7 @@ 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", @@ -13,4 +14,5 @@ __all__ = [ "DMD2Method", "build_dmd2", "DiffusionNFTMethod", "build_diffusion_nft", "SelfForcingMethod", "build_self_forcing", + "UnifiedRLMethod", "build_unified_rl", ] diff --git a/v2/training/methods/unified_rl.py b/v2/training/methods/unified_rl.py new file mode 100644 index 000000000..7979dc4e1 --- /dev/null +++ b/v2/training/methods/unified_rl.py @@ -0,0 +1,222 @@ +"""UnifiedRLMethod — joint LM + generator RL (UniRL / PromptRL §6; design_v3 §10). + +The stress test for the whole design. One reward, computed once per sample, drives TWO trainable +experts under one group-relative advantage: + + * the ``llm`` prompt-refiner — a **token policy gradient** (REINFORCE on the refinement action, + ``−A·logπ(a) + β·KL(π‖π_ref)``); and + * the ``transformer`` flow generator — **FlowGRPO PPO** over the SDE rollout's per-step log-probs + (``max(−A·ρ, −A·clip(ρ, 1±ε)) + β·KL`` with ``ρ = exp(logp_cur − logp_rollout)``). + +What it exercises that nothing else in v2 does: + 1. a training rollout that spans **two loop types** — the ``ar_decode`` refine *and* the + ``diffusion_denoise`` rollout — not one; + 2. **dual log-prob capture** — a categorical action log-prob (LM) and a Gaussian per-step + log-prob (SDE diffusion) — feeding one advantage; + 3. **likelihood-based C2** — per-step log-prob identity (the OTHER half of the C2 split; NFT is + likelihood-*free*); + 4. a **multi-component joint update** — two WeightSyncPlans (one per expert), versioned and + cache-invalidated independently (the LM sync must not flush the text-encoder feature cache); + 5. the rollout SDE sampler ≠ the serve ODE sampler — a controlled §9.4 divergence, gated by + ``DiffusionParams.sde_rollout`` so the serve path is untouched. + +``joint=False`` is PromptRL's *prompt-only* mode: the generator is frozen (ODE rollout, no DiT +update); only the LM learns. ``num_skip_refinement`` samples per group use the *original* prompt +(partial refinement) so the advantage contrasts refined vs unrefined — the LM's learning signal. + +This is a stress test, not a redesign: it adds a method and a card, no new loop/runtime primitive. +""" +from __future__ import annotations + +import numpy as np + +from ..._enums import ConsistencyLevel, ExecutionProfile +from ...loop.sampler import flow_sde_step_with_logprob +from ...models.common import cached_text_encode +from ...request import DiffusionParams, TaskType, make_request +from ..rollout import rollout_loop +from ..weight_sync import WeightRole, WeightSyncPlan +from .base import TrainingMethod, new_instance + + +class UnifiedRLMethod(TrainingMethod): + name = "unified_rl" + consistency = ConsistencyLevel.C2 # likelihood-BASED (per-step log-prob identity) + student_loop_id = "diffusion_denoise" + + def __init__(self, student_instance, reference_instance, *, joint: bool = True, + num_samples_per_prompt: int = 4, num_skip_refinement: int = 1, rollout_steps: int = 4, + llm_lr: float = 0.4, dit_lr: float = 0.02, beta: float = 0.01, adv_clip_max: float = 5.0, + ppo_clip: float = 1e-2, sde_noise_scale: float = 0.7, target_action: int = 3, + action_bonus: float = 2.0, **kw): + super().__init__(student_instance, lr=dit_lr, **kw) + self.reference = reference_instance + self.joint = joint + self.K = num_samples_per_prompt + self.num_skip = max(0, min(num_skip_refinement, self.K - 1)) + self.rollout_steps = rollout_steps + self.llm_lr, self.dit_lr = llm_lr, dit_lr + self.beta, self.adv_clip_max, self.ppo_clip = beta, adv_clip_max, ppo_clip + self.sde_noise_scale = sde_noise_scale + self.target_action, self.action_bonus = target_action, action_bonus + # one plan PER trainable expert — versioned + cache-scoped independently (design_v3 §7.1, §10) + self.llm_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("llm",)) + self.dit_sync = WeightSyncPlan(role=WeightRole.STUDENT, components=("transformer",)) + # reference = frozen copy of BOTH experts at init (the two KL anchors) + self.reference.component("llm").copy_from(self.llm) + self.reference.component("transformer").copy_from(self.dit) + + # --- expert handles --------------------------------------------------------- # + @property + def llm(self): + return self.student.component("llm") + + @property + def dit(self): + return self.student.component("transformer") + + def manages_optimization(self) -> bool: + return True + + def get_grad_clip_targets(self, iteration: int = 0) -> dict: + t = {"llm": self.llm} + if self.joint: + t["transformer"] = self.dit + return t + + # --- toy reward: a diffusion signal + an LM action bonus -------------------- # + def _reward(self, latent: np.ndarray, action: int | None) -> float: + """Reward depends on BOTH experts: lower-spread latents (generator) + the LM picking the + reward-favored ``target_action`` (refiner). So one scalar trains both.""" + diffusion_signal = float(np.tanh(-np.std(np.asarray(latent, dtype="float64")))) + lm_signal = self.action_bonus if (action is not None and int(action) == self.target_action) else 0.0 + return diffusion_signal + lm_signal + + def _rollout_sample(self, prompt: str, emb, neg_emb, seed: int): + """Drive the SHARED diffusion loop, conditioned on the (possibly refined) embedding. + SDE rollout (with per-step log-prob capture) when joint; deterministic ODE when prompt-only.""" + # guidance_scale=1 ⇒ the CFG combine collapses to the cond branch, so the PPO recompute + # (a single cond forward) replays the rollout velocity exactly — the likelihood-based C2 + # identity (FLUX is guidance-distilled: one effective forward, no CFG at RL time). + req = make_request(TaskType.T2V, self.student.card.model_id, prompt, + diffusion=DiffusionParams(num_steps=self.rollout_steps, seed=seed, + guidance_scale=1.0, sde_rollout=self.joint, + sde_noise_scale=self.sde_noise_scale)) + slots = {"text_embeds": np.asarray(emb, dtype="float32"), + "neg_text_embeds": np.asarray(neg_emb, dtype="float32")} + return rollout_loop(self.student, self.student_loop_id, req, slots=slots, + profile=ExecutionProfile.ROLLOUT) + + def _dit_ppo(self, behavior, emb, advantage): + """FlowGRPO PPO over one rollout's SDE trajectory: per-step ratio + KL, then a toy DiT step. + Returns (mean_ppo_loss, mean_kl, mean_ratio, grad_norm).""" + ref_dit = self.reference.component("transformer") + ppo_losses, kls, ratios, gnorms = [], [], [], [] + for rec in behavior or []: + if "sde_logprob" not in rec: + continue + prev, sample = rec["prev"], rec["sample"] + st, sn = float(rec["sigma_t"]), float(rec["sigma_next"]) + v_cur = self.dit(prev, emb, st) + _, logp_cur, mean_cur, eff_std = flow_sde_step_with_logprob( + prev, v_cur, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) + v_ref = ref_dit(prev, emb, st) + _, _, mean_ref, _ = flow_sde_step_with_logprob( + prev, v_ref, st, sn, prev_sample=sample, noise_scale=self.sde_noise_scale) + ratio = float(np.exp(np.clip(logp_cur - float(rec["sde_logprob"]), -10.0, 10.0))) + clipped = float(np.clip(ratio, 1.0 - self.ppo_clip, 1.0 + self.ppo_clip)) + ppo = max(-advantage * ratio, -advantage * clipped) # FlowGRPO PPO (per element) + kl = float(np.mean((mean_cur - mean_ref) ** 2) / (2.0 * max(eff_std ** 2, 1e-8))) + ppo_losses.append(ppo) + kls.append(kl) + ratios.append(ratio) + # toy generator update: nudge toward the rollout velocity scaled by advantage sign + v_target = np.asarray(rec["velocity"], dtype="float32") + lr_eff = self.dit_lr * float(np.clip(advantage, -1.0, 1.0)) + _, gn = self.dit.mse_grad_step(prev, emb, st, v_target, lr_eff) + gnorms.append(gn) + m = lambda xs: float(np.mean(xs)) if xs else 0.0 # noqa: E731 + return m(ppo_losses), m(kls), m(ratios), m(gnorms) + + def managed_train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]: + self.iteration = iteration + prompts = batch["prompts"] + seeds = batch.get("seeds", list(range(len(prompts)))) + neg_base = cached_text_encode(self.student, "") # shared negative embed + lm_losses, dit_losses, kls, ratios = [], [], [], [] + lm_gnorms, dit_gnorms, all_reward, all_adv = [], [], [], [] + + for pi, prompt in enumerate(prompts): + base_emb = cached_text_encode(self.student, prompt) # encoded ONCE for the group + rng = np.random.default_rng((seeds[pi] + 1) * 7919 + iteration) + action, logp_action = self.llm.sample_refinement(rng) # the LM's RL action (sampled) + refined_emb = self.llm.refined_embed(base_emb, action) + + samples = [] # (latent, behavior, emb, reward, refined?) + for k in range(self.K): + refined = k >= self.num_skip # partial refinement + emb = refined_emb if refined else base_emb + act = action if refined else None + res = self._rollout_sample(prompt, emb, neg_base, seed=seeds[pi] * 1000 + k) + latent = np.asarray(res.outputs["latents"], dtype="float32") + reward = self._reward(latent, act) + samples.append((latent, res.behavior, emb, reward, refined)) + + rewards = np.array([s[3] for s in samples], dtype="float64") + adv = np.clip((rewards - rewards.mean()) / (rewards.std() + 1e-4), + -self.adv_clip_max, self.adv_clip_max) # group-relative advantage + all_reward.extend(rewards.tolist()) + all_adv.extend(adv.tolist()) + + # --- LM token policy gradient: advantage of the REFINED samples (vs original) ------- # + refined_adv = [adv[k] for k in range(self.K) if samples[k][4]] + lm_adv = float(np.mean(refined_adv)) if refined_adv else 0.0 + lm_kl = self.llm.kl_to(self.reference.component("llm")) + gn_lm = self.llm.pg_step(action, lm_adv, self.llm_lr) # REAL REINFORCE step + lm_losses.append(-lm_adv * logp_action + self.beta * lm_kl) + lm_gnorms.append(gn_lm) + + # --- generator FlowGRPO PPO (joint only): one update per refined sample ------------- # + if self.joint: + for k in range(self.K): + if not samples[k][4]: + continue + ppo, kl, ratio, gn = self._dit_ppo(samples[k][1], samples[k][2], float(adv[k])) + dit_losses.append(ppo + self.beta * kl) + kls.append(kl) + ratios.append(ratio) + dit_gnorms.append(gn) + + # publish both experts' weights — independent versions + cache scopes (design_v3 §10) + lm_ver = self.llm_sync.apply(self.llm, self.llm, self.student) + dit_ver = self.dit_sync.apply(self.dit, self.dit, self.student) if self.joint else "frozen" + + p_target = float(self.llm._probs()[self.target_action]) # LM learning signal + metrics = { + "lm_pg_loss": float(np.mean(lm_losses)), + "grad_norm/llm": float(np.mean(lm_gnorms)), + "reward_mean": float(np.mean(all_reward)), + "advantage_std": float(np.std(all_adv)), + "refine_target_prob": p_target, + "joint": float(self.joint), + "llm_weights_version": lm_ver, + "transformer_weights_version": dit_ver, + } + loss_map = {"lm_pg_loss": metrics["lm_pg_loss"]} + if self.joint: + metrics.update({ + "dit_pg_loss": float(np.mean(dit_losses)) if dit_losses else 0.0, + "kl_div_loss": float(np.mean(kls)) if kls else 0.0, + "ppo_ratio_mean": float(np.mean(ratios)) if ratios else 1.0, + "grad_norm/transformer": float(np.mean(dit_gnorms)) if dit_gnorms else 0.0, + }) + loss_map["dit_pg_loss"] = metrics["dit_pg_loss"] + return loss_map, metrics + + def train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]: + return self.managed_train_step(batch, iteration) + + +def build_unified_rl(card, *, joint: bool = True, **kw) -> UnifiedRLMethod: + """Two roles: a trainable student (llm + transformer) and a frozen reference (both experts).""" + return UnifiedRLMethod(new_instance(card), new_instance(card), joint=joint, **kw) diff --git a/v2/training/weight_sync.py b/v2/training/weight_sync.py index 1f3997cc3..f93546025 100644 --- a/v2/training/weight_sync.py +++ b/v2/training/weight_sync.py @@ -36,6 +36,10 @@ class WeightSyncPlan: layout_adapters: dict[str, Any] = field(default_factory=dict) transport: str = "in_proc" # in_proc | shm | nccl | rdma decay: float = 0.0 # for EMA / OLD_POLICY blend + # which component(s) this plan syncs — versioned + cache-invalidated in isolation (design_v3 §7.1). + # Joint RL (UniRL) ships ONE plan per trainable expert: ``("transformer",)`` and ``("llm",)`` — + # so the LM weight sync does NOT flush the frozen text-encoder's feature cache, and vice-versa. + components: tuple[str, ...] = ("transformer",) def apply(self, src_dit: Any, dst_dit: Any, dst_instance: Any = None) -> str: """Execute the lifecycle. Returns the new weights_version published on dst_instance.""" @@ -45,10 +49,10 @@ class WeightSyncPlan: dst_dit.blend_from(src_dit, self.decay) # decay·dst + (1-decay)·src else: dst_dit.copy_from(src_dit) # hard copy (student push / reference init) - # 4) bump version + 5) invalidate the TRANSFORMER's caches/graphs + 6) publish. - # Scoped to the transformer (the synced weights) so frozen text/vision encoders keep - # their feature caches (design_v3 §7.1). + # 4) bump version + 5) invalidate ONLY this plan's component caches/graphs + 6) publish. + # Scoped to ``self.components`` so frozen text/vision encoders keep their feature caches, + # and the two experts of a joint-RL update version independently (design_v3 §7.1, §10). version = f"w{next(_ver)}" if dst_instance is not None: - dst_instance.set_weights_version(version, components=["transformer"]) + dst_instance.set_weights_version(version, components=list(self.components)) return version