[refactor]: remove training semantics from v2 inference contracts
This commit is contained in:
+6
-7
@@ -1,15 +1,14 @@
|
||||
"""v2 — a scoped, CPU-testable realization of the model-native runtime (see v2/README.md).
|
||||
"""v2 - the FastVideo inference runtime (see v2/README.md).
|
||||
|
||||
> A model card is a (recipe, runtime) pair with a parity obligation.
|
||||
> The model owns loop semantics; the runtime owns loop lifecycle.
|
||||
> One resident instance runs many loops; one scheduler runs their steps in one currency.
|
||||
> Caches are correct by key; parity is correct by test; the interleave gate is non-negotiable.
|
||||
> Training records behavior on the same loops it serves.
|
||||
> Caches are correct by key; parity is correct by test.
|
||||
|
||||
Phase 1 supports Wan2.1-1.3B (T2V) and LTX2.3 (2-stage distilled), plus four training
|
||||
methods on Wan2.1-1.3B (finetuning, DMD2, DiffusionNFT, self-forcing). The spine is
|
||||
omni-ready (multi-loop ModelInstance, ar_decode/chunk_rollout loop kinds) for the phase-2
|
||||
Cosmos3 + vllm-omni omni ports.
|
||||
v2 is inference-only. Training, finetuning, distillation, RL, and optimizer
|
||||
loops belong to ``fastvideo/train`` or ``fastvideo/training``. v2 only records
|
||||
checkpoint provenance in recipe metadata so inference can bind weights to the
|
||||
right loop and precision policy.
|
||||
|
||||
The core is numpy-only and CPU-testable; heavy Wan/LTX neural forwards become lazy torch
|
||||
adapters (see ``v2/platform/backends/``) that are off the test path.
|
||||
|
||||
@@ -21,8 +21,6 @@ if TYPE_CHECKING: # avoid runtime card -> parallel coupling
|
||||
pass
|
||||
|
||||
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Precision, parallelism, parity, cache, checkpoint contracts #
|
||||
# --------------------------------------------------------------------------- #
|
||||
@@ -31,7 +29,6 @@ class PrecisionContract:
|
||||
default_dtype: str = "float32"
|
||||
component_overrides: dict[str, str] = field(default_factory=dict)
|
||||
quantization_scheme: str | None = None # "nvfp4" | "int8" | None
|
||||
training_precision: str = "float32" # distinct from serving precision
|
||||
|
||||
def dtype_for(self, component_id: str) -> str:
|
||||
return self.component_overrides.get(component_id, self.default_dtype)
|
||||
@@ -52,7 +49,7 @@ class CacheContract:
|
||||
eviction: str = "lru" # "lru" | "fifo" | "none"
|
||||
reuse_across_requests: bool = True # paged/feature=True; slab depends on mode
|
||||
per_component: dict[str, int] = field(default_factory=dict)
|
||||
training_mode_disables_recycle: bool = False # chunk-KV training mode
|
||||
disable_recycle: bool = False # keep full causal/world context instead of sliding-window recycle
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -133,7 +130,7 @@ class LoopSpec:
|
||||
loop_id: str
|
||||
kind: LoopKind
|
||||
work_unit_kind: WorkUnitKind
|
||||
behavior_schema: type | None = None # what to capture for RL (None if not training-relevant)
|
||||
behavior_schema: type | None = None # optional inference trajectory/debug payload schema
|
||||
extension_schema: type | None = None # per-model LoopState extension (Cosmos3PackedSeq, etc.)
|
||||
cache_policy: list[str] = field(default_factory=list) # cache class names this loop draws from
|
||||
graph_capture: str = "eager" # eager | breakable_cudagraph
|
||||
|
||||
+4
-10
@@ -17,7 +17,6 @@ class LoopKind(str, Enum):
|
||||
VAE_TILE = "vae_tile" # tiled VAE encode/decode
|
||||
ENCODER = "encoder" # one-shot encoder (text/vision) — degenerate single-step loop
|
||||
AUDIO_DECODE = "audio_decode" # vocoder / codec decode
|
||||
TRAIN_FORWARD = "train_forward" # grad-enabled forward for a training method
|
||||
|
||||
|
||||
class WorkUnitKind(str, Enum):
|
||||
@@ -34,18 +33,16 @@ class WorkUnitKind(str, Enum):
|
||||
ENCODER_CHUNK = "encoder_chunk"
|
||||
VAE_TILE = "vae_tile"
|
||||
AUDIO_CHUNK = "audio_chunk"
|
||||
REWARD_BATCH = "reward_batch" # RL
|
||||
LOGPROB_BATCH = "logprob_batch" # RL (likelihood-based)
|
||||
TRANSFER = "transfer"
|
||||
CACHE_IO = "cache_io"
|
||||
GRAPH_CAPTURE = "graph_capture"
|
||||
|
||||
|
||||
class ConsistencyLevel(str, Enum):
|
||||
"""The consistency ladder. RL methods declare their required rung."""
|
||||
"""The inference parity ladder."""
|
||||
C0 = "C0" # component parity (VAE/encoder/block/scheduler-step in isolation)
|
||||
C1 = "C1" # loop parity (full denoise trajectory / AR logits, fixed seed)
|
||||
C2 = "C2" # behavioral identity (train-forward ≡ serve-forward on the RL objective's quantity)
|
||||
C2 = "C2" # behavioral identity for inference-visible trajectories
|
||||
C3 = "C3" # distribution parity (rollout distribution under allowed nondeterminism)
|
||||
C4 = "C4" # artifact quality (SSIM / reward agreement / human preference)
|
||||
|
||||
@@ -55,10 +52,9 @@ class ConsistencyLevel(str, Enum):
|
||||
|
||||
|
||||
class ExecutionProfile(str, Enum):
|
||||
"""Three forwards, one loop definition — differ only in grad mode + capture."""
|
||||
"""Inference execution profiles."""
|
||||
SERVE = "serve" # no-grad, graphed, cached, possibly quantized
|
||||
ROLLOUT = "rollout" # serve profile + behavior capture
|
||||
TRAIN = "train" # grad, checkpointed, FSDP-gathered
|
||||
ROLLOUT = "rollout" # serve profile + trajectory capture / stochastic rollout
|
||||
|
||||
|
||||
class Capability(str, Enum):
|
||||
@@ -75,5 +71,3 @@ class Capability(str, Enum):
|
||||
STREAMING_VIDEO_CONTINUATION = "streaming_video_continuation"
|
||||
VAE_ENCODE = "vae_encode"
|
||||
VAE_DECODE = "vae_decode"
|
||||
POLICY_ROLLOUT = "policy_rollout" # can serve as an RL rollout engine
|
||||
LOGPROB_RECOMPUTE = "logprob_recompute"
|
||||
|
||||
@@ -47,13 +47,11 @@ def flow_sde_step_with_logprob(x_t,
|
||||
noise=None,
|
||||
prev_sample=None,
|
||||
noise_scale: float = 0.7):
|
||||
"""FlowGRPO-style stochastic step + Gaussian log-prob.
|
||||
"""Stochastic diffusion step + Gaussian log-prob metadata.
|
||||
|
||||
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)``.
|
||||
This is the stochastic counterpart to deterministic ``flow_match_euler_step``. ``prev_sample=None``
|
||||
samples a new step; passing ``prev_sample`` recomputes its log-prob under the current velocity.
|
||||
Returns ``(prev_sample, log_prob, mean, eff_std)``.
|
||||
"""
|
||||
x_t = np.asarray(x_t, dtype=np.float64)
|
||||
velocity = np.asarray(velocity, dtype=np.float64)
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
"""The consistency ladder + numeric comparison helpers.
|
||||
|
||||
C0 component · C1 loop · C2 behavioral · C3 distribution · C4 artifact.
|
||||
C0 component · C1 loop · C2 trajectory behavior · C3 distribution · C4 artifact.
|
||||
|
||||
The C2 split is load-bearing: likelihood-based methods compare per-step log-probs;
|
||||
likelihood-free methods (DiffusionNFT) compare seeded final-sample + prediction-space
|
||||
identity (old_deviate / ref-MSE) — there are NO log-probs to match.
|
||||
The ladder is for inference parity. C0/C1 are component and loop parity, C2 can compare
|
||||
captured trajectory records, C3 covers stochastic distribution parity, and C4 covers artifact quality.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
@@ -132,8 +132,7 @@ class ParallelWorkflow(Workflow):
|
||||
|
||||
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) — the other non-linear shape."""
|
||||
a scorer, return the best. This is a feedback loop across servable models, not a training loop."""
|
||||
|
||||
def __init__(self, workflow_id, generator_stage, *, scorer, n: int = 4, score_key: str = "latents"):
|
||||
super().__init__(workflow_id, [generator_stage])
|
||||
|
||||
@@ -35,8 +35,8 @@ 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. Default False ⇒ the serve path is unchanged.
|
||||
# Stochastic sampling: use the SDE sampler instead of the deterministic ODE serve sampler.
|
||||
# Default False means the normal 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).
|
||||
@@ -47,7 +47,7 @@ class DiffusionParams:
|
||||
class CaptureMode(str, Enum):
|
||||
"""What the loop captures while running."""
|
||||
NONE = "none" # serve forward: no-grad, no capture
|
||||
BEHAVIOR = "behavior" # rollout forward: emit BehaviorRecord slices for RL/distill
|
||||
BEHAVIOR = "behavior" # trajectory/debug capture for inference diagnostics
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
"""Shared lightweight type aliases for v2.
|
||||
|
||||
The core (card/loop/runtime/cache/parity/program/request/training) is numpy-only
|
||||
The core (card/loop/runtime/cache/parity/program/request) is numpy-only
|
||||
and CPU-testable, so tensors are typed structurally as ``TensorLike``: a
|
||||
``numpy.ndarray`` on the CPU test path, a ``torch.Tensor`` on GPU. The core never
|
||||
imports torch — only model-component adapters do, lazily (see
|
||||
|
||||
+2
-15
@@ -11,7 +11,6 @@ from typing import Any
|
||||
from v2.recipes import build_default_engine, build_omni_engine
|
||||
from v2.recipes.wan21 import build_wan21_card
|
||||
from v2.core.request import DiffusionParams, OutputSpec, Request, SamplingParams, TaskType, make_request
|
||||
from v2.training import build_diffusion_nft
|
||||
|
||||
|
||||
def _t2v(mid: str, prompt: str, seed: int, steps: int = 4, **kw: Any) -> Request:
|
||||
@@ -40,17 +39,6 @@ def example_c_causal_streaming(eng) -> None:
|
||||
f"streamed_chunks={out.metrics.get('stream_chunks', 0)}")
|
||||
|
||||
|
||||
def example_d_rl_rollout() -> None:
|
||||
print("\n(d) RL rollout (DiffusionNFT): the SAME denoise loop + behavior capture (train ≡ serve)")
|
||||
nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=4, num_inner_timesteps=2)
|
||||
loss, m = nft.managed_train_step({"prompts": ["a red car", "a blue boat"], "seeds": [1, 2]}, 0)
|
||||
fc = nft.old.caches.stats()["feature"]
|
||||
print(f" policy_loss={loss['policy_loss']:.3f} kl={loss['kl_div_loss']:.5f} "
|
||||
f"reward_mean={m['reward_mean']:.3f} consistency={nft.consistency_level().value} (likelihood-free)")
|
||||
print(f" shared-prompt feature-cache reuse: {fc['hits']} hits / {fc['misses']} misses "
|
||||
f"(K samples encode the prompt once — the 24× reduction)")
|
||||
|
||||
|
||||
def example_g_omni_mot() -> None:
|
||||
print("\n(g) Omni / MoT (§16): ONE resident instance runs AR + diffusion loops on shared weights")
|
||||
eng = build_omni_engine()
|
||||
@@ -78,7 +66,7 @@ async def _serving_demo() -> None:
|
||||
import asyncio
|
||||
|
||||
from v2.serving.deploy import DynamoWorkerAdapter, FakeDynamoRuntime, LocalFleet, build_deployment_card
|
||||
from v2.recipes.wan21 import build_wan21_card, build_wan_t2v_program
|
||||
from v2.recipes.wan21 import build_wan_t2v_program
|
||||
from v2.runtime import AsyncEngine, PoolSet, wan_t2v_disaggregated
|
||||
from v2.serving import OmniOpenAIServer
|
||||
|
||||
@@ -143,12 +131,11 @@ def main() -> None:
|
||||
example_a_text_to_video(eng)
|
||||
example_b_ltx2_two_stage(eng)
|
||||
example_c_causal_streaming(eng)
|
||||
example_d_rl_rollout()
|
||||
example_g_omni_mot()
|
||||
example_h_serving_and_fleet()
|
||||
print("\n" + "=" * 78)
|
||||
print("All examples ran on CPU with numpy toy components. The architecture (cards, driven")
|
||||
print("loops, scheduler, caches, parity, training-on-shared-loops) is real; the neural")
|
||||
print("loops, scheduler, caches, parity, and serving) is real; the neural")
|
||||
print("forwards are toys. On a GPU box, swap ComponentSpec.factory for the torch adapters.")
|
||||
print("=" * 78)
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
Two tuple-keyed registries (``COMPONENTS``, ``KERNELS``) + a ``Platform`` that detects ``(device,
|
||||
arch)`` and resolves both, with the numpy reference as the terminal fallback rung and parity oracle.
|
||||
Lets CPU, GPU, and other backends coexist: a model's loops, policies, scheduler, caches, and
|
||||
training code never name a device; they call through the component/kernel seams and the resolved
|
||||
inference code never names a device; it calls through the component/kernel seams and the resolved
|
||||
``Platform`` decides the implementation.
|
||||
|
||||
Importing this package is cheap and cycle-free (registry + platform classes only); the concrete
|
||||
|
||||
@@ -45,7 +45,7 @@ class AccelDiT:
|
||||
return self._inner(*args, **kwargs)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
# delegate the trainable surface (clone/copy_from/blend_from/mse_grad_step) to the inner toy
|
||||
# delegate the weight/version surface (clone/copy_from/blend_from) to the inner toy
|
||||
return getattr(self._inner, name)
|
||||
|
||||
|
||||
|
||||
+6
-109
@@ -2,8 +2,8 @@
|
||||
|
||||
There is no GPU/torch/weights in this environment, so the heavy Wan/LTX neural forwards are
|
||||
represented by small, deterministic numpy components. They exercise the *real* loop control
|
||||
flow, CFG/flow-shift policies, scheduler steps, cache reuse, parity gates, and training-method
|
||||
math with real numbers — just with a toy network instead of a 1.3B DiT.
|
||||
flow, CFG/flow-shift policies, scheduler steps, cache reuse, and parity gates with real numbers -
|
||||
just with a toy network instead of a 1.3B DiT.
|
||||
|
||||
The contract these toys honor is what matters:
|
||||
* deterministic given (weights, inputs) → bit-reproducible (the interleave/parity gates rely on this);
|
||||
@@ -12,8 +12,8 @@ The contract these toys honor is what matters:
|
||||
|
||||
On a GPU box, ``ComponentSpec.factory`` swaps these for lazy torch adapters wrapping the real
|
||||
``fastvideo.models`` modules + weights (see each model's ``components.py``); the loops, policies,
|
||||
scheduler, caches, parity, and training code are unchanged. That is the whole point of the
|
||||
(recipe, runtime) separation.
|
||||
scheduler, caches, and parity are unchanged. That is the whole point of the (recipe, runtime)
|
||||
separation.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -105,7 +105,7 @@ class ToyDiT:
|
||||
audio = np.tanh(0.9 * au + 0.1 * float(sigma) + 0.3 * a_cond).astype("float32")
|
||||
return video, audio
|
||||
|
||||
# --- minimal trainable surface (so training methods do real optimizer steps) --------- #
|
||||
# --- weight/version surface for serving hot-reload tests ----------------------------- #
|
||||
def clone(self) -> ToyDiT:
|
||||
c = ToyDiT.__new__(ToyDiT)
|
||||
c.C, c.s_text, c.s_ctx = self.C, self.s_text, self.s_ctx
|
||||
@@ -113,40 +113,13 @@ class ToyDiT:
|
||||
return c
|
||||
|
||||
def blend_from(self, other: ToyDiT, decay: float) -> None:
|
||||
"""EMA / decay-blended-old-policy update: self ← decay·self + (1-decay)·other."""
|
||||
"""Blend resident weights toward another component version."""
|
||||
self.w_x = (decay * self.w_x + (1.0 - decay) * other.w_x).astype("float32")
|
||||
self.w_t = (decay * self.w_t + (1.0 - decay) * other.w_t).astype("float32")
|
||||
|
||||
def copy_from(self, other: ToyDiT) -> None:
|
||||
self.w_x, self.w_t = other.w_x.copy(), other.w_t.copy()
|
||||
|
||||
def mse_grad_step(self,
|
||||
latent: np.ndarray,
|
||||
text_embed,
|
||||
sigma: float,
|
||||
target: np.ndarray,
|
||||
lr: float,
|
||||
context=None,
|
||||
weight: float = 1.0) -> tuple[float, float]:
|
||||
"""One exact SGD step minimizing ``weight·MSE(forward(latent), target)`` w.r.t. ``w_x``.
|
||||
|
||||
Returns (loss, grad_norm). The grad_norm is the #1396-style per-method regression signal.
|
||||
This is a *real* gradient (chain rule through tanh), so the toy student actually learns —
|
||||
the structural stand-in for the FSDP/optimizer step a GPU trainer would run."""
|
||||
latent = np.asarray(latent, dtype=np.float32)
|
||||
target = np.asarray(target, dtype=np.float32)
|
||||
z = self._pre_tanh(latent, text_embed, sigma, context)
|
||||
pred = np.tanh(z)
|
||||
err = pred - target
|
||||
loss = float(weight * np.mean(err**2))
|
||||
n = err.size
|
||||
g = weight * (2.0 * err / n) * (1.0 - pred**2)
|
||||
axes = (list(range(1, g.ndim)), list(range(1, latent.ndim)))
|
||||
grad_wx = np.tensordot(g, latent, axes=axes) # [C_out, C_in]
|
||||
grad_norm = float(np.linalg.norm(grad_wx))
|
||||
self.w_x = (self.w_x - lr * grad_wx).astype("float32")
|
||||
return loss, grad_norm
|
||||
|
||||
|
||||
class ToyTokenizer:
|
||||
"""Deterministic byte-ish tokenizer for the omni AR pathway (phase 2)."""
|
||||
@@ -185,82 +158,6 @@ 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 (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.ndarray = 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.ndarray = 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 ToyTalker(ToyMoTDiT):
|
||||
"""Toy Talker — Qwen-Omni stage 1 (AR over a speech-codec vocab, conditioned on the Thinker).
|
||||
|
||||
|
||||
@@ -40,7 +40,7 @@ from collections.abc import Callable, Iterable
|
||||
# Op names — the stable identifiers a loop dispatches on (not free functions). #
|
||||
# --------------------------------------------------------------------------- #
|
||||
FLOW_MATCH_STEP = "flow_match_step" # deterministic flow-match Euler solver step (ODE serve)
|
||||
FLOW_SDE_STEP = "flow_sde_step" # FlowGRPO stochastic step + log-prob (RL rollout)
|
||||
FLOW_SDE_STEP = "flow_sde_step" # stochastic diffusion step + log-prob metadata
|
||||
|
||||
|
||||
def _always_available() -> bool:
|
||||
|
||||
@@ -38,12 +38,13 @@ def build_bagel_card(model_id: str = "bagel-mot") -> ModelCard:
|
||||
return ARDecodeLoop(loop_id="generate_text", transformer_id="transformer", max_tokens=6)
|
||||
|
||||
def image_factory():
|
||||
return WanDenoiseLoop(loop_id="generate_image",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
return WanDenoiseLoop(
|
||||
loop_id="generate_image",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
|
||||
components = {
|
||||
"tokenizer":
|
||||
@@ -93,7 +94,7 @@ def build_bagel_card(model_id: str = "bagel-mot") -> ModelCard:
|
||||
"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"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
"""Shared component-node helpers (text encode with feature cache, etc.).
|
||||
|
||||
The content-hash feature cache lets a K-sample RL group encode its shared prompt once instead of
|
||||
K times.
|
||||
The content-hash feature cache lets repeated requests reuse prompt embeddings safely.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
@@ -37,12 +37,13 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
|
||||
return ARDecodeLoop(loop_id="ar_decode", transformer_id="transformer", max_tokens=6)
|
||||
|
||||
def dn_factory():
|
||||
return WanDenoiseLoop(loop_id="diffusion_denoise",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
return WanDenoiseLoop(
|
||||
loop_id="diffusion_denoise",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
|
||||
components = {
|
||||
"tokenizer":
|
||||
@@ -89,8 +90,7 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT,
|
||||
Capability.TEXT_TO_VIDEO_SOUND, Capability.VAE_DECODE,
|
||||
Capability.POLICY_ROLLOUT),
|
||||
Capability.TEXT_TO_VIDEO_SOUND, Capability.VAE_DECODE),
|
||||
recipe=RecipeSpec(method="base",
|
||||
assumes_loop="diffusion_denoise",
|
||||
assumes_precision="float32",
|
||||
@@ -100,7 +100,7 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
|
||||
"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"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -107,7 +107,7 @@ def _build_flux2_card(model_id: str, *, text_encoder_load_id: str, text_encoder_
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1],
|
||||
tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
sampling_defaults=sampling_defaults,
|
||||
)
|
||||
|
||||
@@ -135,7 +135,7 @@ def build_ltx2_card(model_id: str = "ltx2-2stage-distilled") -> ModelCard:
|
||||
negative_prompt=""),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
@@ -211,7 +211,7 @@ def build_ltx2_base_card(model_id: str = "ltx2-single-stage") -> ModelCard:
|
||||
negative_prompt=LTX2_NEG),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
@@ -301,7 +301,7 @@ def build_ltx2_3_card(model_id: str = "ltx2.3-distilled") -> ModelCard:
|
||||
}),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -127,7 +127,7 @@ def build_matrixgame2_card(model_id: str = "matrix-game-2.0-base-distilled",
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1],
|
||||
tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
sampling_defaults=sampling_defaults or SamplingDefaults(num_steps=3,
|
||||
guidance_scale=1.0,
|
||||
|
||||
@@ -8,7 +8,7 @@ talker token, and vocoder chunk is a runtime-visible WorkUnit.
|
||||
A third weight-sharing topology, in the same Card/Loop/Program vocabulary:
|
||||
|
||||
* Cosmos3/BAGEL — ONE shared MoT module, two loops (weight sharing within a request).
|
||||
* unified/UniRL — TWO disjoint experts, two loops, jointly RL-trained.
|
||||
* multi-expert inference — TWO disjoint experts, two loops, one request.
|
||||
* Qwen-Omni — THREE disjoint experts (thinker, talker, vocoder), three loops, cascaded:
|
||||
the talker conditions on the thinker's tokens+hidden state, the vocoder on the
|
||||
talker's speech tokens; streaming codec→waveform.
|
||||
@@ -51,10 +51,7 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard:
|
||||
prompt_slot="talker_prompt_tokens")
|
||||
|
||||
def vocoder_factory():
|
||||
return VocoderLoop(loop_id="vocoder",
|
||||
vocoder_id="vocoder",
|
||||
chunk_tokens=2,
|
||||
speech_slot="speech_tokens")
|
||||
return VocoderLoop(loop_id="vocoder", vocoder_id="vocoder", chunk_tokens=2, speech_slot="speech_tokens")
|
||||
|
||||
components = {
|
||||
"tokenizer":
|
||||
@@ -115,8 +112,7 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard:
|
||||
family="qwen_omni",
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH,
|
||||
Capability.POLICY_ROLLOUT),
|
||||
capabilities=CapabilityMatrix.of(Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH),
|
||||
recipe=RecipeSpec(method="base",
|
||||
assumes_loop="vocoder",
|
||||
assumes_precision="float32",
|
||||
@@ -125,7 +121,7 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard:
|
||||
caches={
|
||||
"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"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
+16
-14
@@ -53,12 +53,13 @@ def build_wan21_i2v_card(model_id: str = "wan2.1-i2v-1.3b",
|
||||
precision, expert = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True), NoRouting("transformer")
|
||||
|
||||
def loop_factory():
|
||||
return WanDenoiseLoop(loop_id="i2v_denoise",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
return WanDenoiseLoop(
|
||||
loop_id="i2v_denoise",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
|
||||
components = {
|
||||
"text_encoder":
|
||||
@@ -109,7 +110,7 @@ def build_wan21_i2v_card(model_id: str = "wan2.1-i2v-1.3b",
|
||||
consistency_required=ConsistencyLevel.C1),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
|
||||
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
sampling_defaults=SamplingDefaults(num_steps=40,
|
||||
guidance_scale=5.0,
|
||||
@@ -138,12 +139,13 @@ def build_wan22_i2v_a14b_card(model_id: str = "wan2.2-i2v-a14b",
|
||||
expert = BoundaryTimestepRouting(high_noise="transformer", low_noise="transformer_2", boundary=boundary)
|
||||
|
||||
def loop_factory():
|
||||
return WanDenoiseLoop(loop_id="i2v_denoise",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
return WanDenoiseLoop(
|
||||
loop_id="i2v_denoise",
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
expert=expert,
|
||||
)
|
||||
|
||||
def _dit(cid: str) -> ComponentSpec:
|
||||
return ComponentSpec(cid,
|
||||
@@ -198,7 +200,7 @@ def build_wan22_i2v_a14b_card(model_id: str = "wan2.2-i2v-a14b",
|
||||
consistency_required=ConsistencyLevel.C1),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
|
||||
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
sampling_defaults=SamplingDefaults(num_steps=40,
|
||||
guidance_scale=3.5,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Wan2.1 T2V program: text_encode → diffusion_denoise → vae_decode.
|
||||
"""Wan2.1 T2V program: text_encode -> diffusion_denoise -> vae_decode.
|
||||
|
||||
The text-encode node uses the content-hash feature cache: the same prompt across requests (e.g. a
|
||||
K=24 RL group) encodes once and reuses.
|
||||
The text-encode node uses the content-hash feature cache: the same prompt across requests encodes
|
||||
once and reuses.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
@@ -32,7 +32,7 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
|
||||
num_chunks: int = 7,
|
||||
chunk_size: int = 3,
|
||||
steps_per_chunk: int = 4,
|
||||
training_mode: bool = False) -> ModelCard:
|
||||
preserve_full_context: bool = False) -> ModelCard:
|
||||
# Self-forcing DMD reference (fastvideo SelfForcingWanT2V480PConfig + causal_denoising.py): the
|
||||
# distilled student is CFG-FREE (guidance 1.0, single forward/step) and denoises with the 4-step DMD
|
||||
# schedule (dmd_denoising_steps=[1000,750,500,250], warped — reproduced by FlowShiftPolicy(5.0) over 4
|
||||
@@ -44,14 +44,15 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
|
||||
precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True)
|
||||
|
||||
def loop_factory():
|
||||
return ChunkRolloutLoop(loop_id="chunk_rollout",
|
||||
num_chunks=num_chunks,
|
||||
chunk_size=chunk_size,
|
||||
steps_per_chunk=steps_per_chunk,
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
)
|
||||
return ChunkRolloutLoop(
|
||||
loop_id="chunk_rollout",
|
||||
num_chunks=num_chunks,
|
||||
chunk_size=chunk_size,
|
||||
steps_per_chunk=steps_per_chunk,
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
)
|
||||
|
||||
components = {
|
||||
"text_encoder":
|
||||
@@ -89,7 +90,7 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.STREAMING_VIDEO_CONTINUATION,
|
||||
Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
|
||||
Capability.VAE_DECODE),
|
||||
recipe=RecipeSpec(method="self_forcing",
|
||||
parents=["wan2.1-1.3b"],
|
||||
assumes_loop="chunk_rollout",
|
||||
@@ -103,10 +104,10 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
|
||||
CacheContract(cache_class="slab_kv",
|
||||
max_bytes=1 << 26,
|
||||
reuse_across_requests=False,
|
||||
per_component={"window": (num_chunks if training_mode else 2)},
|
||||
training_mode_disables_recycle=training_mode),
|
||||
per_component={"window": (num_chunks if preserve_full_context else 2)},
|
||||
disable_recycle=preserve_full_context),
|
||||
},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
precision=PrecisionContract(default_dtype="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
sampling_defaults=SamplingDefaults(
|
||||
num_steps=4, guidance_scale=1.0, height=480, width=832, num_frames=81, fps=16,
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
"""ChunkRolloutLoop — causal/AR video.
|
||||
|
||||
Loop shape: outer loop over latent-frame chunks × inner denoise per chunk, carrying a per-chunk
|
||||
KV/context via the slab-KV cache (sink/local window in inference; full history in self-forcing
|
||||
training mode). Each completed chunk is streamable (``StepResult.emit``), and the rollout profile
|
||||
captures per-chunk behavior for the self-forcing distillation method.
|
||||
Loop shape: outer loop over latent-frame chunks x inner denoise per chunk, carrying a per-chunk
|
||||
KV/context via the slab-KV cache. The cache can use a sink/local window or preserve full request
|
||||
history. Each completed chunk is streamable (``StepResult.emit``), and the rollout profile can
|
||||
capture per-chunk inference trajectory records.
|
||||
|
||||
State is entirely in ``LoopState`` (the slab list + chunk position), so interleaving causal
|
||||
rollouts of two sessions cannot smear context — and the slab-KV pool is namespaced per request.
|
||||
@@ -151,9 +151,7 @@ class ChunkRolloutLoop:
|
||||
seq=st.scratch["chunk_idx"],
|
||||
data=x,
|
||||
preview=False))
|
||||
res = ResourceRequest(
|
||||
resident_bytes=int(x.nbytes),
|
||||
peak_activation_bytes=int(x.nbytes * len(branches)))
|
||||
res = ResourceRequest(resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * len(branches)))
|
||||
return WorkPlan(loop_id=self.loop_id,
|
||||
instance_id=st.instance_id,
|
||||
kind=WorkUnitKind.CHUNK_STEP,
|
||||
|
||||
Vendored
+6
-6
@@ -4,7 +4,7 @@ No single unified block pool: cache classes differ by 150-500x in natural granul
|
||||
each gets its own statically budgeted pool behind one ``CacheHandle``. Four classes:
|
||||
* ``FeatureCache`` — content-hash keyed, partitioned by adapter+weights (text/vision encoders)
|
||||
* ``ResidualCache`` — cache-dit residuals, scoped per request AND per CFG branch
|
||||
* ``SlabKVCache`` — chunk-KV slabs (self-forcing / world models); training mode disables recycle
|
||||
* ``SlabKVCache`` — chunk-KV slabs for causal/world models; can disable recycle
|
||||
* ``PagedKVCache`` — paged text-KV stub for ar_decode (phase-2 omni); minority case, lazy
|
||||
|
||||
KV is the minority case — a pure bidirectional deployment (Wan/LTX T2V) allocates none of it.
|
||||
@@ -71,7 +71,7 @@ class FeatureCache(_Pool):
|
||||
self.used_bytes -= nb
|
||||
|
||||
def invalidate_weights(self, version: str) -> None:
|
||||
"""RL update_weights bumps weight epoch → drop entries from older epochs (wholesale)."""
|
||||
"""Drop entries from older resident-weight epochs."""
|
||||
drop = [h for h, (_v, _nb, k) in self._store.items() if k.weights_version != version]
|
||||
for h in drop:
|
||||
_v, nb, _k = self._store.pop(h)
|
||||
@@ -124,21 +124,21 @@ class Slab:
|
||||
class SlabKVCache(_Pool):
|
||||
"""Chunk-KV slabs for causal/world-model rollout.
|
||||
|
||||
``training_mode`` disables mid-rollout recycling so activation-checkpoint recompute
|
||||
doesn't double-advance the cache (self-forcing).
|
||||
``disable_recycle`` preserves the full request-local context. Otherwise the
|
||||
cache behaves as a sliding window with the declared ``window`` size.
|
||||
"""
|
||||
|
||||
def __init__(self, policy: CachePolicy):
|
||||
super().__init__(policy)
|
||||
self._store: dict[str, list[Slab]] = {}
|
||||
self.window = max(1, policy.per_component.get("window", 1 << 30))
|
||||
self.training_mode = policy.training_mode_disables_recycle
|
||||
self.disable_recycle = policy.disable_recycle
|
||||
|
||||
def append(self, namespace: str, slab: Slab) -> None:
|
||||
slabs = self._store.setdefault(namespace, [])
|
||||
slabs.append(slab)
|
||||
self.used_bytes += _nbytes(slab.k) + _nbytes(slab.v)
|
||||
if not self.training_mode and len(slabs) > self.window:
|
||||
if not self.disable_recycle and len(slabs) > self.window:
|
||||
dropped = slabs.pop(0) # sliding-window recycle (inference only)
|
||||
self.used_bytes -= _nbytes(dropped.k) + _nbytes(dropped.v)
|
||||
|
||||
|
||||
Vendored
+3
-4
@@ -2,8 +2,7 @@
|
||||
|
||||
If a field can change output semantics, it is in the key (incorrect reuse is worse than no reuse).
|
||||
The serving hazard this kills: a request that shares a prompt but differs in te-LoRA stack must not
|
||||
serve stale embeddings — so the key is *partitioned* by ``adapter_versions``, not flushed. An RL
|
||||
``update_weights`` bumps ``weights_version`` and invalidates wholesale.
|
||||
serve stale embeddings, so the key is *partitioned* by ``adapter_versions``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -15,7 +14,7 @@ from typing import Any
|
||||
def content_hash(obj: Any) -> str:
|
||||
"""Stable content hash for feature-cache keys (text/vision embeddings).
|
||||
|
||||
Lets K rollout samples of one prompt reuse a single text encode.
|
||||
Lets repeated requests for one prompt reuse a single text encode.
|
||||
"""
|
||||
h = hashlib.sha256()
|
||||
if isinstance(obj, str):
|
||||
@@ -78,4 +77,4 @@ class CachePolicy:
|
||||
eviction: str = "lru" # "lru" | "fifo" | "none"
|
||||
reuse_across_requests: bool = True
|
||||
per_component: dict[str, int] = field(default_factory=dict)
|
||||
training_mode_disables_recycle: bool = False
|
||||
disable_recycle: bool = False
|
||||
|
||||
Vendored
+2
-2
@@ -2,7 +2,7 @@
|
||||
|
||||
Static partitioning makes cross-class fragmentation impossible (jumbo slab traffic cannot
|
||||
strand text-KV pages and vice versa). Each class gets a budget carved at init from the card's
|
||||
CacheContracts. ``invalidate_weights`` implements the wholesale RL weight-epoch bump.
|
||||
CacheContracts.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -32,7 +32,7 @@ class CacheManager:
|
||||
eviction=cc.eviction,
|
||||
reuse_across_requests=cc.reuse_across_requests,
|
||||
per_component=dict(cc.per_component),
|
||||
training_mode_disables_recycle=cc.training_mode_disables_recycle,
|
||||
disable_recycle=cc.disable_recycle,
|
||||
))
|
||||
return cls(policies)
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ correctness-critical and bug-prone regardless of device:
|
||||
* the **eager-break** — a step that is model-declared non-capturable (host RNG / data-dependent
|
||||
branch, e.g. the FlowGRPO SDE step) or that an interceptor overrode this step never touches the
|
||||
graph cache (vLLM's ``@eager_break_during_capture``).
|
||||
* **invalidation** — a resident-weights version bump (RL weight sync) or a shape change changes the
|
||||
* **invalidation** — a resident-weights version bump or a shape change changes the
|
||||
key → recapture, so a graph never outlives the weights/shape it was captured for. This is
|
||||
AUTOMATIC: the resident-weight version is *in* the key, so a synced component's old graphs simply
|
||||
become unreachable and the next step recaptures.
|
||||
@@ -159,7 +159,7 @@ class GraphCapturer:
|
||||
return plan.graph_fn(instance, g.workspace)
|
||||
|
||||
def invalidate(self, component_ids: Any = None) -> int:
|
||||
"""Evict captured graphs whose resident weights just changed (RL weight sync), mirroring
|
||||
"""Evict captured graphs whose resident weights just changed, mirroring
|
||||
``CacheManager.invalidate_components``. Without this, a synced component's old graphs become
|
||||
unreachable (the version is in the key) but linger — a real GPU memory leak across FlowGRPO's
|
||||
per-iteration syncs. With ``component_ids=None``, clears all. Returns the count dropped."""
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
Reservation must be pre-flight: a scheduler that admits a diffusion step and then discovers it
|
||||
cannot allocate the VAE tile is wrong.
|
||||
|
||||
Sleep/wake is component-granular (tags are component names) for RL: drop DiT + caches,
|
||||
Sleep/wake is component-granular (tags are component names): drop DiT + caches,
|
||||
keep VAE/text-encoder resident (CuMem-style).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""Per-class cache pools: feature partition/invalidation, slab-KV window/training."""
|
||||
"""Per-class cache pools: feature partition/invalidation and slab-KV windows."""
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
@@ -23,7 +23,7 @@ def test_feature_cache_weight_invalidation_is_wholesale():
|
||||
fc = FeatureCache(CachePolicy("feature", max_bytes=1 << 20))
|
||||
k = CacheKey(model_id="m", component_id="te", weights_version="v0")
|
||||
fc.put(k, np.ones((2,), dtype="float32"))
|
||||
fc.invalidate_weights("v1") # RL update_weights bump
|
||||
fc.invalidate_weights("v1")
|
||||
assert fc.get(k) is None
|
||||
|
||||
|
||||
@@ -35,17 +35,16 @@ def test_feature_cache_budget_eviction():
|
||||
assert fc.used_bytes <= 8
|
||||
|
||||
|
||||
def test_slab_kv_inference_window_vs_training_mode():
|
||||
def test_slab_kv_window_vs_disable_recycle():
|
||||
inf = SlabKVCache(CachePolicy("slab_kv", reuse_across_requests=False, per_component={"window": 2}))
|
||||
for i in range(4):
|
||||
inf.append("r", Slab(i, np.zeros((3,), dtype="float32"), None))
|
||||
assert len(inf.get("r")) == 2 # sliding window recycles (inference)
|
||||
|
||||
tr = SlabKVCache(CachePolicy("slab_kv", training_mode_disables_recycle=True,
|
||||
per_component={"window": 2}))
|
||||
tr = SlabKVCache(CachePolicy("slab_kv", disable_recycle=True, per_component={"window": 2}))
|
||||
for i in range(4):
|
||||
tr.append("r", Slab(i, np.zeros((3,), dtype="float32"), None))
|
||||
assert len(tr.get("r")) == 4 # no mid-rollout recycle in training mode (self-forcing)
|
||||
assert len(tr.get("r")) == 4 # no mid-rollout recycle when full context is requested
|
||||
assert tr.used_bytes > 0
|
||||
tr.clear_namespace("r")
|
||||
assert tr.used_bytes == 0
|
||||
|
||||
@@ -89,7 +89,7 @@ def test_weight_version_bump_recaptures_and_evicts_stale():
|
||||
inst = _wan()
|
||||
_drive(inst, seed=1, steps=4)
|
||||
assert inst.graphs.stats["captures"] == 1 and inst.graphs.captured == 1
|
||||
inst.set_weights_version("v1", ["transformer"]) # RL weight sync: bumps version + evicts graph
|
||||
inst.set_weights_version("v1", ["transformer"]) # version bump evicts stale graphs
|
||||
assert inst.graphs.captured == 0 # stale graph dropped (no leak)
|
||||
_drive(inst, seed=1, steps=4)
|
||||
assert inst.graphs.stats["captures"] == 2 # recaptured under the new version...
|
||||
|
||||
@@ -43,12 +43,12 @@ def test_feature_key_partitions_by_adapter_stack():
|
||||
assert inst.caches.stats()["feature"]["misses"] == m + 1
|
||||
|
||||
|
||||
def test_transformer_weight_sync_keeps_text_encoder_cache():
|
||||
def test_transformer_weight_reload_keeps_text_encoder_cache():
|
||||
card = build_wan21_card()
|
||||
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
||||
cached_text_encode(inst, "p")
|
||||
assert inst.caches.stats()["feature"]["used_bytes"] > 0
|
||||
inst.set_weights_version("w1", components=["transformer"]) # RL-style transformer-only sync
|
||||
inst.set_weights_version("w1", components=["transformer"]) # transformer-only reload
|
||||
h = inst.caches.stats()["feature"]["hits"]
|
||||
cached_text_encode(inst, "p") # text-encoder embed survived ⇒ hit (not flushed)
|
||||
assert inst.caches.stats()["feature"]["hits"] == h + 1
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
"""Qwen-Omni thinker→talker→vocoder — the third weight-sharing topology (vllm-omni
|
||||
``qwen2_5_omni`` pipeline).
|
||||
|
||||
Asserts the Card/Loop/Program vocabulary holds for a THREE-expert, THREE-loop cascade — the topology
|
||||
neither MoT (one shared module, two loops) nor joint-RL (two disjoint experts, two loops) covers:
|
||||
Asserts the Card/Loop/Program vocabulary holds for a THREE-expert, THREE-loop cascade:
|
||||
|
||||
* three disjoint experts (thinker, talker, vocoder) on three loop *types* (ar_decode, ar_decode,
|
||||
audio_decode) in one request — every token and vocoder chunk a runtime-visible WorkUnit;
|
||||
@@ -102,4 +101,3 @@ def test_vocoder_streams_audio_chunks():
|
||||
# the audio-streaming delta is exactly the vocoder's per-chunk emits (AR text emits are common)
|
||||
assert on.metrics["stream_chunks"] - off.metrics["stream_chunks"] == n_chunks
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user