From de295403c7637698d115b855457b64b6fe920d86 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 16 Jun 2026 23:51:32 -0700 Subject: [PATCH] =?UTF-8?q?[feat]=20Adapter=20plane,=20non-linear=20workfl?= =?UTF-8?q?ows,=20RL=E2=86=92distill=20flywheel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three more capabilities on distinct untested surfaces (167 tests pass, no new runtime primitive). A — Adapter plane (§9.19): one base + swappable LoRA/ControlNet adapters, selected per request (DiffusionParams.adapters); AdapterDenoiseLoop applies each active adapter's velocity delta. Per-request selection changes output, multi-LoRA composes, ControlNet conditions on a control image, mixed-adapter requests interleave without smearing, hot-swap changes generation, cache key partitions by adapter stack (the adapter_versions field, previously declared-only). ToyLoRA/ToyControlNet; models/adapters/. 6 tests. B — Non-linear workflows (§9.17): ParallelWorkflow (fan-out: one input → N models → merged) and BestOfNWorkflow (generate N → score with the served reward card → return best; inference-time scaling). The shapes a linear chain can't express. 4 tests. D — RL→distill flywheel (§9.18): run_flywheel RL-improves the base (NFT), then distills FROM the RL'd model (DMD2 teacher = RL'd policy) into a faster card, recording the base→rl→distilled provenance chain in RecipeSpec.parents. The distilled student is measurably closer to the RL'd teacher than the base; the distilled card serves few-step. training/flywheel.py. 4 tests. designv4 §9.17–§9.19 + layout/closing/counts (167 tests, 29 files). --- designv4.md | 55 ++++++++++++++++++-- v2/models/__init__.py | 2 + v2/models/adapters/__init__.py | 8 +++ v2/models/adapters/card.py | 73 ++++++++++++++++++++++++++ v2/models/adapters/loop.py | 47 +++++++++++++++++ v2/models/adapters/program.py | 31 +++++++++++ v2/models/backend.py | 48 +++++++++++++++++ v2/program/__init__.py | 10 +++- v2/program/workflow.py | 66 ++++++++++++++++++++++++ v2/request/params.py | 3 ++ v2/tests/test_adapters.py | 88 ++++++++++++++++++++++++++++++++ v2/tests/test_flywheel.py | 62 ++++++++++++++++++++++ v2/tests/test_workflow_shapes.py | 73 ++++++++++++++++++++++++++ v2/training/__init__.py | 3 +- v2/training/flywheel.py | 52 +++++++++++++++++++ 15 files changed, 613 insertions(+), 8 deletions(-) create mode 100644 v2/models/adapters/__init__.py create mode 100644 v2/models/adapters/card.py create mode 100644 v2/models/adapters/loop.py create mode 100644 v2/models/adapters/program.py create mode 100644 v2/tests/test_adapters.py create mode 100644 v2/tests/test_flywheel.py create mode 100644 v2/tests/test_workflow_shapes.py create mode 100644 v2/training/flywheel.py diff --git a/designv4.md b/designv4.md index 89ac19da0..75375897c 100644 --- a/designv4.md +++ b/designv4.md @@ -5,7 +5,7 @@ > 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 (153 tests, 26 files, two independent runners — `pytest` +> is backed by code in `v2/` and a passing test (167 tests, 29 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 @@ -556,6 +556,44 @@ cleanest proof of model-owned control flow, recursive composition depth, the RL correctness, the reward plane composing with serving, and an exact AR speedup — each a new card / loop / method / controller, **no new runtime primitive**. +### 9.17 Non-linear workflow shapes — fan-out + best-of-N feedback + +Every workflow so far was a *linear chain*. Two non-linear shapes complete the composition vocabulary: +**`ParallelWorkflow`** (fan-out) runs N stages on the *same* input and merges their artifacts namespaced +by branch — one prompt → N models in parallel (variants, or video+audio+upscale) → a merged result, the +engine free to interleave the branches' steps. **`BestOfNWorkflow`** (feedback) generates N candidates, +scores each with a reward scorer, and returns the best — inference-time scaling / rejection sampling, +composing a generator with the served REWARD_BATCH card (§9.15) in a loop. Both reuse the `Workflow` +requires/validate machinery (fan-out requires all branch models; best-of-N requires the generator). +✅ `v2/tests/test_workflow_shapes.py` (4: fan-out merge, best-of-N returns the argmax candidate, +determinism, requires). + +### 9.18 The (recipe, runtime) flywheel — RL → distill → faster card + +design_v3 §16's central product claim, made concrete and *measured*: `run_flywheel` (1) RL-improves the +base (DiffusionNFT), then (2) distills **from the RL'd model** — a DMD2 student whose **teacher is the +RL'd policy** — into a faster few-step card, (3) recording the provenance chain `base → rl → distilled` +in `RecipeSpec.parents`. Both stages drive the same denoise loop the engine serves, so the flywheel is +real, not aspirational — and the distilled student is **measurably closer to the RL'd teacher than the +base is** (the distillation worked), the distilled card is a valid few-step servable, and the recipe +chain is an auditable typed fact ("this fast model came from that RL run"). ✅ `v2/tests/test_flywheel.py` +(4: chaining, distilled-approaches-teacher, provenance chain, the distilled card serves few-step). + +### 9.19 The adapter plane — per-request LoRA / ControlNet over one base + +The `CacheKey.adapter_versions` field existed (and was cache-tested), but the *serving* capability did +not. Now many adapters are declared over ONE resident base; a request selects which to apply +(`DiffusionParams.adapters`), and `AdapterDenoiseLoop` applies each active adapter's velocity delta. What +it proves: per-request selection changes the output (base ≠ LoRA-A ≠ LoRA-B), **multi-LoRA composes**, a +**ControlNet** conditions on a control image (different control → different video), mixed-adapter requests +**interleave without smearing** (the active set lives in the request/`LoopState`, never global), +**hot-swap** changes generation, and the cache key **partitions by the adapter stack** (adapted ≠ base, +version A ≠ B — the §7.1 guarantee). One base, N lightweight swappable adapters — the production-ubiquitous +pattern, as a card + an isolated loop. ✅ `v2/tests/test_adapters.py` (6). + +These three (§9.17–§9.19) add the non-linear workflow shapes, the headline RL→distill flywheel, and the +adapter-serving plane — again each a new workflow / method+helper / card+loop, **no new runtime primitive**. + ## 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 @@ -609,11 +647,15 @@ v2/ ltx2/ (audio_vae) # §9.11: joint A/V denoise (T2VS, per-modality guidance) adaptive/ # §9.12: CacheDiTDenoiseLoop (cache-dit skip + early-exit) reward/ # §9.15: reward-model card, REWARD_BATCH work units + speculative/ # §9.16: draft + target AR, speculative decoding + adapters/ # §9.19: one base + LoRA/ControlNet, per-request (AdapterDenoiseLoop) + program/ Program (one model) + Workflow / ParallelWorkflow / BestOfNWorkflow (cross-model, §9.6/§9.17) training/ rollout, behavior, rewards (+ServedRewardScorer), weight_sync (+WeightSyncController), + flywheel (RL→distill, §9.18), methods/{finetune,dmd2,diffusion_nft,self_forcing,unified_rl,joint_multi_rl,workflow_rl} serving/ AsyncEngine, pools, DisaggregatedRunner, connectors, OpenAI server deploy/ DeploymentCard, LocalFleet, DynamoWorkerAdapter - tests/ 26 files, 153 tests run via `pytest v2/tests/` OR `python3 v2/run_tests.py` + tests/ 29 files, 167 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 @@ -668,9 +710,12 @@ state, cancellation, streaming), under **end-to-end RL across a workflow** (a fi earlier model), under **heterogeneous WorkUnit co-scheduling** (VAE tiles interleaved with denoise steps), under **joint audio+video** (LTX-2 T2VS with per-modality guidance), under **content-adaptive compute** (cache-dit skip + early-exit), under **nested workflows** (recursive composition), under **hot weight-sync -while serving** (the RL flywheel's drain-correct lifecycle), and under a **served reward model** -(REWARD_BATCH). Every one of those fit inside the design as *a new card, method, loop, Workflow, session -driver, or controller*, with **no new runtime primitive** — and the only real bug any stress test surfaced +while serving** (the RL flywheel's drain-correct lifecycle), under a **served reward model** +(REWARD_BATCH), under **speculative decoding** (exact, lower-latency AR), under **non-linear workflows** +(fan-out + best-of-N feedback), under the **RL→distill flywheel** (a faster card with provenance), and +under the **adapter plane** (per-request LoRA/ControlNet over one base). Every one of those fit inside the +design as *a new card, method, loop, Workflow, session driver, or controller*, with **no new runtime +primitive** — and the only real bug any stress test surfaced (a no-op generator gradient) was a fix in the sampler *library*, not the runtime. The design's bet is therefore not that it is elegant — it is that **the weight-sharing topology, the composition graph, the training recipe, the reward, and the session/sync lifecycle are all data over cards, loops, workflows, and diff --git a/v2/models/__init__.py b/v2/models/__init__.py index e11cf1620..5fac000b2 100644 --- a/v2/models/__init__.py +++ b/v2/models/__init__.py @@ -8,6 +8,7 @@ from __future__ import annotations from typing import Any from .adaptive import build_adaptive_card +from .adapters import build_adapter_card, build_adapter_program from .bagel import build_bagel_card, build_bagel_program from .cosmos3 import build_cosmos3_card, build_cosmos3_program from .image_video import ( @@ -40,6 +41,7 @@ __all__ = [ "build_flux_t2i_card", "build_flux_t2i_program", "build_wan_i2v_card", "build_wan_i2v_program", "build_t2i_then_i2v_workflow", "build_t2i_i2v_extend_workflow", "register_workflows", "build_tiled_card", "build_tiled_program", "build_adaptive_card", + "build_adapter_card", "build_adapter_program", "build_default_engine", "build_omni_engine", "build_unified_engine", "build_image_video_engine", "build_tiled_engine", ] diff --git a/v2/models/adapters/__init__.py b/v2/models/adapters/__init__.py new file mode 100644 index 000000000..021f503ee --- /dev/null +++ b/v2/models/adapters/__init__.py @@ -0,0 +1,8 @@ +"""Adapter plane: one base + swappable LoRA/ControlNet adapters, selected per request (design_v3 §9.19).""" +from __future__ import annotations + +from .card import ADAPTERS, build_adapter_card +from .loop import AdapterDenoiseLoop +from .program import build_adapter_program + +__all__ = ["build_adapter_card", "build_adapter_program", "AdapterDenoiseLoop", "ADAPTERS"] diff --git a/v2/models/adapters/card.py b/v2/models/adapters/card.py new file mode 100644 index 000000000..e850f945f --- /dev/null +++ b/v2/models/adapters/card.py @@ -0,0 +1,73 @@ +"""Adapter-serving card — one base DiT + a set of swappable LoRA/ControlNet adapters (design_v3 §9.19). + +Many adapters are declared as components over one resident ``transformer``; a request picks which to +apply (``DiffusionParams.adapters``). Adapters are versioned independently (the cache key's +``adapter_versions``) and hot-swappable. +""" +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 ToyControlNet, ToyDiT, ToyLoRA, ToyTextEncoder, ToyVAE, _seed_from +from .loop import AdapterDenoiseLoop + +ADAPTERS = ("lora_anime", "lora_realistic", "control_pose") # the served adapter library + + +def build_adapter_card(model_id: str = "wan-adapters") -> ModelCard: + seed = _seed_from(model_id) + 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 loop_factory(): + return AdapterDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow, + precision=precision, expert=expert, cost=cost) + + components = { + "text_encoder": ComponentSpec("text_encoder", kind="text_encoder", + factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}), + "transformer": ComponentSpec("transformer", kind="dit", 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"}), + # the adapter library — lightweight, optional, applied per request: + "lora_anime": ComponentSpec("lora_anime", kind="adapter", + factory=lambda inst: ToyLoRA("lora_anime", scale=0.6, seed=seed + 1), + optional_for={"t2v"}), + "lora_realistic": ComponentSpec("lora_realistic", kind="adapter", + factory=lambda inst: ToyLoRA("lora_realistic", scale=0.6, seed=seed + 2), + optional_for={"t2v"}), + "control_pose": ComponentSpec("control_pose", kind="adapter", + factory=lambda inst: ToyControlNet("control_pose", scale=0.8, seed=seed + 3), + optional_for={"t2v"}), + } + loops = { + "diffusion_denoise": LoopSpec("diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost, + shared_weight_components=["transformer"], cache_policy=["feature"], + loop_factory=loop_factory), + } + return ModelCard( + model_id=model_id, family="wan", components=components, loops=loops, + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.VAE_DECODE), + recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise", + assumes_precision="float32", consistency_required=ConsistencyLevel.C1), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True), + caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)}, + precision=PrecisionContract(default_dtype="float32", training_precision="float32"), + parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], + default_plan=ParallelPlan.single()), + ).validate() diff --git a/v2/models/adapters/loop.py b/v2/models/adapters/loop.py new file mode 100644 index 000000000..f160efea7 --- /dev/null +++ b/v2/models/adapters/loop.py @@ -0,0 +1,47 @@ +"""AdapterDenoiseLoop — per-request LoRA / ControlNet over one resident base (design_v3 §9.19). + +The adapter plane: a request selects adapters (``DiffusionParams.adapters``); the loop applies each +active adapter's velocity delta on top of the base DiT's prediction, then re-integrates. Many adapters +are served over ONE resident base; the active set lives in the request/``LoopState`` (never global), so +two requests using different adapters interleave without smearing. Isolated ``WanDenoiseLoop`` subclass +(empty adapter set ⇒ identical to the base), so nothing else is touched. +""" +from __future__ import annotations + +import numpy as np + +from ...loop.contracts import Done +from ...loop.sampler import flow_match_euler_step +from ..wan21.loop import WanDenoiseLoop + + +class AdapterDenoiseLoop(WanDenoiseLoop): + def init(self, req, model, ctx): + st = super().init(req, model, ctx) + st.scratch["adapters"] = tuple(getattr(req.diffusion, "adapters", ()) or ()) + st.scratch["control"] = ctx.slots.get("control_signal") # ControlNet conditioning (optional) + return st + + def next(self, st): + plan = super().next(st) + adapters = st.scratch.get("adapters") or () + if isinstance(plan, Done) or not adapters: + return plan + base_run = plan.run + x = st.latents["video"] + control = st.scratch.get("control") + sigma_t, sigma_next = st.sigmas[st.step_idx], st.sigmas[st.step_idx + 1] + prec = self.precision + + def run(model, override=None): + res = base_run(model, override) # base velocity (CFG combine) + v = np.asarray(res.output["noise_pred"], dtype="float32") + for aid in adapters: # apply each active adapter's delta + v = v + model.component(aid).delta(x, control) + res.output["noise_pred"] = v + res.output["latents"] = flow_match_euler_step(prec.cast(x), v, sigma_t, sigma_next).astype("float32") + return res + + plan.run = run + plan.label = f"{plan.label}.adapt[{'+'.join(adapters)}]" + return plan diff --git a/v2/models/adapters/program.py b/v2/models/adapters/program.py new file mode 100644 index 000000000..1b946052f --- /dev/null +++ b/v2/models/adapters/program.py @@ -0,0 +1,31 @@ +"""Adapter-serving program: control_signal → text_encode → diffusion_denoise(+adapters) → vae_decode.""" +from __future__ import annotations + +import numpy as np + +from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind +from ..common import text_encode_node_fn as _text_encode + + +def _control_signal(instance, slots, request, ctx) -> None: + """A ControlNet's conditioning image (pose/depth/edge) from the request → a slot the loop reads.""" + img = request.image() + slots["control_signal"] = None if img is None else np.asarray(img.pixels, dtype="float32") + + +def _vae_decode(instance, slots, request, ctx) -> None: + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_adapter_program() -> Program: + return Program( + program_id="wan.adapters", kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("control_signal", fn=_control_signal, writes=("control_signal",)), + ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), + ModelLoopNode("denoise", loop_id="diffusion_denoise", output_slot="denoise_out", + reads=("text_embeds", "control_signal"), writes=("denoise_out",)), + ComponentNode("vae_decode", fn=_vae_decode, reads=("denoise_out",), writes=("video",)), + ], + output_artifacts={"video": "video", "latents": "denoise_out"}, + ).validate() diff --git a/v2/models/backend.py b/v2/models/backend.py index 9a0d32bb6..000e6a83c 100644 --- a/v2/models/backend.py +++ b/v2/models/backend.py @@ -262,6 +262,54 @@ class ToyVocoder: return np.concatenate(segs) if segs else np.zeros(0, dtype="float32") +class ToyLoRA: + """Toy LoRA adapter — a low-rank velocity delta on a base DiT, selected and applied per request. + + Many of these are served over ONE resident base (the adapter is tiny: a rank-r channel delta); + a request picks which to apply. Swappable/versioned independently (the cache key's + ``adapter_versions``). Stand-in for a real LoRA / DoRA / IP-Adapter (design_v3 §9.19).""" + kind = "lora" + + def __init__(self, adapter_id: str, channels: int = LATENT_CHANNELS, scale: float = 0.6, + rank: int = 2, seed: int = 0): + self.adapter_id = adapter_id + self.scale = float(scale) + rng = np.random.default_rng(seed) + self.a = (rng.standard_normal((channels, rank)) * 0.4).astype("float32") + self.b = (rng.standard_normal((rank, channels)) * 0.4).astype("float32") + + def delta(self, latent, control=None) -> np.ndarray: + w = self.a @ self.b # [C, C] low-rank weight delta + d = np.tensordot(w, np.asarray(latent, dtype="float32"), axes=([1], [0])) # [C, ...] + return (self.scale * np.tanh(d)).astype("float32") + + def update(self, seed: int) -> None: # hot-swap: replace the adapter's weights + rng = np.random.default_rng(seed) + self.a = (rng.standard_normal(self.a.shape) * 0.4).astype("float32") + + +class ToyControlNet: + """Toy ControlNet adapter — conditions the velocity on a control signal (pose/depth/edge stand-in). + Different control ⇒ different generation; selected per request like a LoRA (design_v3 §9.19).""" + kind = "controlnet" + + def __init__(self, adapter_id: str, channels: int = LATENT_CHANNELS, scale: float = 0.8, seed: int = 1): + self.adapter_id = adapter_id + self.scale = float(scale) + rng = np.random.default_rng(seed) + self.proj = (rng.standard_normal(channels) * 0.3).astype("float32") + + def delta(self, latent, control=None) -> np.ndarray: + latent = np.asarray(latent, dtype="float32") + c = float(np.mean(control)) if control is not None else 0.0 # the control image's signal + shape = (latent.shape[0],) + (1,) * (latent.ndim - 1) + return (self.scale * c * self.proj.reshape(shape)).astype("float32") + + def update(self, seed: int) -> None: + rng = np.random.default_rng(seed) + self.proj = (rng.standard_normal(self.proj.shape) * 0.3).astype("float32") + + def _spec_target_next(tokens) -> int: """The target AR model's greedy next-token. Length-dependent so the sequence stays varied (no trivial absorbing fixed point) — the substrate for a meaningful speculative-decoding accept rate.""" diff --git a/v2/program/__init__.py b/v2/program/__init__.py index d18f8b010..5530cff1f 100644 --- a/v2/program/__init__.py +++ b/v2/program/__init__.py @@ -13,8 +13,14 @@ from .specs import ( when_opt, when_task, ) -from .workflow import Workflow, WorkflowRegistry, WorkflowStage +from .workflow import ( + BestOfNWorkflow, + ParallelWorkflow, + Workflow, + WorkflowRegistry, + WorkflowStage, +) __all__ = ["Program", "ProgramKind", "ProgramNode", "ComponentNode", "ModelLoopNode", "Edge", "EdgeKind", "always", "when_task", "when_opt", - "Workflow", "WorkflowStage", "WorkflowRegistry"] + "Workflow", "WorkflowStage", "WorkflowRegistry", "ParallelWorkflow", "BestOfNWorkflow"] diff --git a/v2/program/workflow.py b/v2/program/workflow.py index 635659179..b130b76f5 100644 --- a/v2/program/workflow.py +++ b/v2/program/workflow.py @@ -94,6 +94,72 @@ class Workflow: return out +def _payload(artifact: Any) -> Any: + for attr in ("frames", "samples", "latent", "tensor", "text"): + v = getattr(artifact, attr, None) + if v is not None: + return v + return None + + +class ParallelWorkflow(Workflow): + """Fan-out: run every stage on the SAME initial input (conceptually in parallel — the engine can + interleave their steps), then merge their artifacts into one Output, namespaced by stage label. + The non-linear shape a linear chain can't express: one prompt → N models → a merged result (e.g. + variants, or video+audio+upscale in parallel). Design_v3 §13.""" + + def run(self, engine: Any, **initial: Any) -> Any: + if not self.stages: + raise ValueError(f"workflow {self.workflow_id!r} has no stages") + from ..request.artifacts import Output + merged: dict[str, Any] = {} + metrics: dict[str, float] = {} + rid = self.workflow_id + for stage in self.stages: + label = stage.label or stage.model_id + req = stage.make_request(dict(initial)) # each branch sees the ORIGINAL input + if req.model_id != stage.model_id: + raise ValueError(f"{self.workflow_id!r} stage {label!r}: model mismatch") + out = engine.run(req) + rid = out.request_id + for name, art in out.artifacts.items(): + merged[f"{label}:{name}"] = art # namespaced — branches don't collide + for k, v in out.metrics.items(): + metrics[f"{label}:{k}"] = v + return Output(request_id=rid, artifacts=merged, metrics=metrics) + + +class BestOfNWorkflow(Workflow): + """Inference-time scaling / rejection sampling: generate N candidates (varying seed), score each with + a reward scorer, return the best. A feedback loop across models — generator + reward (e.g. the served + REWARD_BATCH card, §9.15) — the other non-linear shape. Design_v3 §10/§13.""" + + def __init__(self, workflow_id, generator_stage, *, scorer, n: int = 4, score_key: str = "latents"): + super().__init__(workflow_id, [generator_stage]) + self.scorer = scorer # any .score(media, prompts) -> {"avg": ...} + self.n = int(n) + self.score_key = score_key + + def run(self, engine: Any, **initial: Any) -> Any: + import numpy as np + stage = self.stages[0] + base_seed = int(initial.get("seed", 0)) + cands = [] + for i in range(self.n): + req = stage.make_request({**initial, "seed": base_seed * 100 + i}) + if req.model_id != stage.model_id: + raise ValueError(f"{self.workflow_id!r}: model mismatch") + cands.append(engine.run(req)) + media = [_payload(c.artifacts[self.score_key]) for c in cands] + scores = np.asarray(self.scorer.score(media, [initial.get("prompt", "")] * len(cands))["avg"]) + best = int(np.argmax(scores)) + out = cands[best] + out.metrics["best_of_n"] = float(self.n) + out.metrics["best_index"] = float(best) + out.metrics["best_score"] = float(scores[best]) + return out + + class WorkflowRegistry: """Declarative ``workflow_id → builder`` catalog — the cross-model analog of the card builders in ``models/__init__.py`` (cf. vllm-omni's ``pipeline_registry``). Adding a custom pipeline is one diff --git a/v2/request/params.py b/v2/request/params.py index 421974b41..92d15fa49 100644 --- a/v2/request/params.py +++ b/v2/request/params.py @@ -41,6 +41,9 @@ class DiffusionParams: # deterministic ODE serve sampler (PromptRL/UniRL §6). Default False ⇒ the serve path is unchanged. sde_rollout: bool = False sde_noise_scale: float = 0.7 + # Active adapters for this request (LoRA / ControlNet ids on the base) — the adapter plane (§9.19). + # Empty ⇒ the base model, unmodified. Part of the cache key (adapter_versions) so adapted ≠ base. + adapters: tuple[str, ...] = () class CaptureMode(str, Enum): diff --git a/v2/tests/test_adapters.py b/v2/tests/test_adapters.py new file mode 100644 index 000000000..b761d9b80 --- /dev/null +++ b/v2/tests/test_adapters.py @@ -0,0 +1,88 @@ +"""The adapter plane — per-request LoRA / ControlNet over one base (design_v3 §7.1, §9.19). + +Many adapters served over ONE resident base; a request picks which to apply (`DiffusionParams.adapters`). +The `adapter_versions` cache-key field existed but the serving capability didn't. These assert: +per-request selection changes output, multi-LoRA composes, ControlNet conditions on a control image, +mixed-adapter requests interleave without smearing, hot-swap changes generation, and the cache key +partitions by the active adapter stack (adapted ≠ base, and version A ≠ version B). +""" +from __future__ import annotations + +import numpy as np + +from v2.cache import CacheManager +from v2.cache.keys import CacheKey +from v2.card import load_card +from v2.models.adapters import build_adapter_card, build_adapter_program +from v2.parity import assert_interleave_parity +from v2.request import DiffusionParams, TaskType, make_request +from v2.request.modalpart import ImagePart, TextPart +from v2.runtime import Engine + + +def _engine(): + card = build_adapter_card() + inst = load_card(card, cache_manager=CacheManager.from_card(card)) + eng = Engine() + eng.register(card.model_id, inst, build_adapter_program()) + return eng, inst + + +def _gen(eng, adapters=(), control=None, seed=3): + inputs = (TextPart("a portrait"),) + ((ImagePart(pixels=control),) if control is not None else ()) + req = make_request(TaskType.T2V, "wan-adapters", "a portrait", inputs=inputs, + diffusion=DiffusionParams(num_steps=6, seed=seed, adapters=adapters)) + return np.asarray(eng.run(req).artifacts["video"].frames) + + +def test_per_request_adapter_changes_output(): + eng, _ = _engine() + base = _gen(eng) + anime = _gen(eng, ("lora_anime",)) + realistic = _gen(eng, ("lora_realistic",)) + assert not np.array_equal(base, anime) + assert not np.array_equal(anime, realistic) + assert np.array_equal(anime, _gen(eng, ("lora_anime",))) # deterministic + + +def test_multi_lora_composes(): + eng, _ = _engine() + a = _gen(eng, ("lora_anime",)) + b = _gen(eng, ("lora_realistic",)) + ab = _gen(eng, ("lora_anime", "lora_realistic")) + assert not np.array_equal(ab, a) and not np.array_equal(ab, b) + + +def test_controlnet_conditions_on_the_control_image(): + eng, _ = _engine() + c_hi = _gen(eng, ("control_pose",), control=np.full((3, 1, 4, 6), 0.5, dtype="float32")) + c_lo = _gen(eng, ("control_pose",), control=np.full((3, 1, 4, 6), -0.5, dtype="float32")) + assert not np.array_equal(c_hi, c_lo) # different control ⇒ different video + + +def test_mixed_adapter_requests_interleave_without_smearing(): + """One base, different adapters per request — interleaving is bit-identical to serial (the active + adapter set lives in the request/LoopState, never global).""" + eng, _ = _engine() + reqs = [make_request(TaskType.T2V, "wan-adapters", "a portrait", + diffusion=DiffusionParams(num_steps=4, seed=s, adapters=ad)) + for s, ad in [(1, ()), (2, ("lora_anime",)), (3, ("lora_realistic",))]] + assert not assert_interleave_parity(eng, reqs) + + +def test_hot_swap_changes_generation(): + eng, inst = _engine() + before = _gen(eng, ("lora_anime",)) + inst.component("lora_anime").update(seed=999) # swap the adapter's weights + assert not np.array_equal(before, _gen(eng, ("lora_anime",))) + + +def test_cache_key_partitions_by_adapter_stack(): + base = CacheKey(model_id="m", component_id="transformer", weights_version="w1", + adapter_versions=(), precision="float32", input_hashes=()) + a = CacheKey(model_id="m", component_id="transformer", weights_version="w1", + adapter_versions=CacheKey.adapters({"lora_anime": "v1"}), precision="float32", input_hashes=()) + b = CacheKey(model_id="m", component_id="transformer", weights_version="w1", + adapter_versions=CacheKey.adapters({"lora_anime": "v2"}), precision="float32", input_hashes=()) + assert a != base # adapted ≠ base (no stale reuse) + assert a != b # adapter version A ≠ version B (hot-swap) diff --git a/v2/tests/test_flywheel.py b/v2/tests/test_flywheel.py new file mode 100644 index 000000000..f0a021aa3 --- /dev/null +++ b/v2/tests/test_flywheel.py @@ -0,0 +1,62 @@ +"""The (recipe, runtime) flywheel — RL → distill → faster card (design_v3 §16; §9.18). + +design_v3's central product claim, made concrete: RL the base (DiffusionNFT), then distill from the +RL'd model (a DMD2 student whose teacher IS the RL'd policy), producing a faster card whose recipe +records the provenance chain base → rl → distilled. Both stages drive the same denoise loop the engine +serves — the flywheel is real, not aspirational. +""" +from __future__ import annotations + +import numpy as np + +from v2.models import build_wan_t2v_program +from v2.models.wan21 import build_wan21_card +from v2.request import DiffusionParams, TaskType, make_request +from v2.runtime import Engine +from v2.training.flywheel import run_flywheel +from v2.training.methods.base import new_instance + + +def _w(inst): + return inst.component("transformer").w_x + + +def test_flywheel_chains_rl_then_distill(): + base = build_wan21_card() + base_w = _w(new_instance(base)).copy() + fw = run_flywheel(base, rl_iters=8, distill_iters=20) + assert not np.allclose(base_w, _w(fw["rl"])) # RL moved the base policy + assert np.allclose(_w(fw["teacher"]), _w(fw["rl"])) # the distill teacher IS the RL'd model + assert not np.allclose(base_w, _w(fw["distilled"])) # the student distilled (moved) + + +def test_distilled_student_approaches_the_rl_teacher(): + base = build_wan21_card() + base_w = _w(new_instance(base)).copy() + fw = run_flywheel(base, rl_iters=8, distill_iters=20) + teach_w, dist_w = _w(fw["teacher"]), _w(fw["distilled"]) + # the distilled student is closer to the RL'd teacher than the (un-distilled) base is + assert np.linalg.norm(dist_w - teach_w) < np.linalg.norm(base_w - teach_w) + + +def test_provenance_chain_is_recorded(): + base = build_wan21_card() + fw = run_flywheel(base, rl_iters=4, distill_iters=6) + assert fw["rl_card"].recipe.method == "diffusion_nft" + assert fw["rl_card"].recipe.parents == ["wan2.1-1.3b"] # rl ← base + assert fw["distilled_card"].recipe.method == "dmd2" + assert fw["distilled_card"].recipe.parents == ["wan2.1-1.3b-rl"] # distilled ← rl + # the full auditable chain: base → rl → distilled + assert fw["distilled_card"].model_id == "wan2.1-1.3b-rl-distilled" + + +def test_distilled_card_is_a_servable_few_step_pair(): + base = build_wan21_card() + fw = run_flywheel(base, rl_iters=4, distill_iters=6) + card, inst = fw["distilled_card"], fw["distilled"] + eng = Engine() + eng.register(card.model_id, inst, build_wan_t2v_program()) # the distilled (recipe,runtime) pair serves + out = eng.run(make_request(TaskType.T2V, card.model_id, "a fox", + diffusion=DiffusionParams(num_steps=4, seed=1))) # few-step + assert out.artifacts["video"].frames is not None + assert out.metrics["denoise_steps"] == 4.0 diff --git a/v2/tests/test_workflow_shapes.py b/v2/tests/test_workflow_shapes.py new file mode 100644 index 000000000..d4a75cb2d --- /dev/null +++ b/v2/tests/test_workflow_shapes.py @@ -0,0 +1,73 @@ +"""Non-linear workflow shapes — fan-out + best-of-N feedback (design_v3 §13; §9.17). + +Every workflow so far was a linear chain. These exercise the two non-linear shapes: + + * **Fan-out** (`ParallelWorkflow`): one input → N models in parallel → merged result (namespaced by + branch). The engine can interleave the branches' steps. + * **Best-of-N** (`BestOfNWorkflow`): generate N candidates, score each with a reward scorer, return + the best — inference-time scaling / rejection sampling, a feedback loop composing a generator with + the served REWARD_BATCH card (§9.15). +""" +from __future__ import annotations + +import numpy as np + +from v2.models import build_default_engine, build_reward_card +from v2.program import BestOfNWorkflow, ParallelWorkflow, WorkflowStage +from v2.request import DiffusionParams, TaskType, make_request +from v2.training.methods.base import new_instance +from v2.training.rewards import ServedRewardScorer + + +def _t2v(model_id): + return lambda s: make_request(TaskType.T2V, model_id, s["prompt"], + diffusion=DiffusionParams(num_steps=4, seed=s.get("seed", 0))) + + +def test_fan_out_runs_branches_and_merges(): + eng = build_default_engine() + wf = ParallelWorkflow("variants.fanout", [ + WorkflowStage("wan2.1-1.3b", _t2v("wan2.1-1.3b"), label="wan"), + WorkflowStage("ltx2.3-distilled", _t2v("ltx2.3-distilled"), label="ltx"), + ]) + out = wf.run(eng, prompt="a comet over the sea", seed=1) + # both branches' artifacts present, namespaced by label + assert "wan:video" in out.artifacts and "ltx:video" in out.artifacts + wan_v = np.asarray(out.artifacts["wan:video"].frames) + ltx_v = np.asarray(out.artifacts["ltx:video"].frames) + assert wan_v.ndim == 4 and ltx_v.ndim == 4 + # two different models on the same prompt → different videos (genuine fan-out, not duplication) + assert wan_v.shape != ltx_v.shape or not np.array_equal(wan_v, ltx_v) + assert wf.requires == ["wan2.1-1.3b", "ltx2.3-distilled"] + + +def _best_of_n(n=4): + eng = build_default_engine() + scorer = ServedRewardScorer(new_instance(build_reward_card(batch=2))) + wf = BestOfNWorkflow("wan.best_of_n", WorkflowStage("wan2.1-1.3b", _t2v("wan2.1-1.3b")), + scorer=scorer, n=n, score_key="latents") + return eng, scorer, wf + + +def test_best_of_n_returns_the_highest_scored_candidate(): + eng, scorer, wf = _best_of_n(n=4) + out = wf.run(eng, prompt="a fox", seed=2) + assert out.metrics["best_of_n"] == 4.0 + # re-score the n candidates independently and confirm the returned one is the argmax + cands = [eng.run(make_request(TaskType.T2V, "wan2.1-1.3b", "a fox", + diffusion=DiffusionParams(num_steps=4, seed=2 * 100 + i))) for i in range(4)] + scores = scorer.score([np.asarray(c.artifacts["latents"].latent) for c in cands], ["a fox"] * 4)["avg"] + assert out.metrics["best_score"] == float(np.max(scores)) + assert out.metrics["best_index"] == float(int(np.argmax(scores))) + + +def test_best_of_n_is_deterministic(): + eng, _, wf = _best_of_n(n=3) + a = wf.run(eng, prompt="a city", seed=5).metrics["best_index"] + b = wf.run(eng, prompt="a city", seed=5).metrics["best_index"] + assert a == b + + +def test_best_of_n_requires_only_the_generator(): + _, _, wf = _best_of_n() + assert wf.requires == ["wan2.1-1.3b"] # the reward scorer is a held instance, not a chained stage diff --git a/v2/training/__init__.py b/v2/training/__init__.py index 20d2e1856..d403d0a1d 100644 --- a/v2/training/__init__.py +++ b/v2/training/__init__.py @@ -6,6 +6,7 @@ engine never imports ``training``. The rollout forward IS the serve forward plus from __future__ import annotations from .behavior import BehaviorRecord, TrajectoryStatus +from .flywheel import derive_card, run_flywheel from .methods import ( DiffusionNFTMethod, DMD2Method, @@ -31,7 +32,7 @@ __all__ = [ "TrainingMethod", "FineTuneMethod", "DMD2Method", "DiffusionNFTMethod", "SelfForcingMethod", "UnifiedRLMethod", "JointMultiExpertRL", "WorkflowRLMethod", "build_finetune", "build_dmd2", "build_diffusion_nft", "build_self_forcing", "build_unified_rl", - "build_joint_multi_rl", "build_workflow_rl", + "build_joint_multi_rl", "build_workflow_rl", "run_flywheel", "derive_card", "BehaviorRecord", "TrajectoryStatus", "Trajectory", "TrajectoryBuffer", "RolloutContext", "rollout_loop", "WeightSyncPlan", "WeightSyncController", "WeightRole", "MultiRewardScorer", "ServedRewardScorer", "build_multi_reward_scorer", diff --git a/v2/training/flywheel.py b/v2/training/flywheel.py new file mode 100644 index 000000000..77ef4001f --- /dev/null +++ b/v2/training/flywheel.py @@ -0,0 +1,52 @@ +"""The (recipe, runtime) flywheel — RL → distill → a faster card (design_v3 §16; §9.18). + +design_v3's central product claim: *preference/reward data → RL → a faster distilled card → a better +product*, with the recipe carrying provenance so the chain is auditable. This wires it end-to-end on the +shared loops: + + 1. **RL** the base model (DiffusionNFT) → improved policy weights. + 2. **Distill** from the RL'd model: a DMD2 student with the **RL'd model as its teacher** learns to + match it in fewer steps. + 3. The distilled card's ``RecipeSpec.parents`` records the chain ``base → rl → distilled`` — the + (recipe, runtime) provenance that makes "this fast model came from that RL run" a typed fact. + +Both stages drive the SAME ``diffusion_denoise`` loop the engine serves (no second sampler) — so the +flywheel is real, not aspirational. +""" +from __future__ import annotations + +import dataclasses +from typing import Any + +from .methods import build_diffusion_nft, build_dmd2 + + +def derive_card(base_card: Any, model_id: str, *, method: str, parents: list[str]) -> Any: + """A child (recipe, runtime) card: same components/loops, new recipe with provenance ``parents``.""" + recipe = dataclasses.replace(base_card.recipe, method=method, parents=list(parents)) + return dataclasses.replace(base_card, model_id=model_id, recipe=recipe).validate() + + +def run_flywheel(base_card: Any, *, rl_iters: int = 8, distill_iters: int = 15, + batch: dict | None = None) -> dict: + """RL the base, then distill from the RL'd model. Returns the instances + the provenance cards.""" + batch = batch or {"prompts": ["a red car", "a blue boat"], "seeds": [1, 2]} + + # 1) RL — DiffusionNFT improves the base policy on the shared denoise loop + nft = build_diffusion_nft(base_card, num_video_per_prompt=4, num_inner_timesteps=2) + for it in range(rl_iters): + nft.train_step(batch, it) + rl_card = derive_card(base_card, base_card.model_id + "-rl", + method="diffusion_nft", parents=[base_card.model_id]) + + # 2) Distill FROM the RL'd model: the DMD2 teacher IS the RL'd policy + dmd = build_dmd2(base_card, rollout_steps=4) + dmd.teacher.component("transformer").copy_from(nft.student.component("transformer")) + for it in range(distill_iters): + dmd.train_step(batch, it) + distilled_card = derive_card(base_card, base_card.model_id + "-rl-distilled", + method="dmd2", parents=[rl_card.model_id]) + + return {"rl": nft.student, "rl_card": rl_card, + "distilled": dmd.student, "distilled_card": distilled_card, + "teacher": dmd.teacher, "base_card": base_card}