[feat] UniRL/PromptRL joint LM+generator RL stress test + designv4
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).
This commit is contained in:
SolitaryThinker
committed by
Peiyuan Zhang
Peiyuan Zhang
parent
a9e5f6ee7a
commit
b3ddf6014d
+402
@@ -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.
|
||||
@@ -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)
|
||||
|
||||
+17
-1
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
+24
-7
@@ -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
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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}"
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user