diff --git a/v2/__init__.py b/v2/__init__.py index 1bd401de4..77213b067 100644 --- a/v2/__init__.py +++ b/v2/__init__.py @@ -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. diff --git a/v2/core/card/specs.py b/v2/core/card/specs.py index c194093bf..795868dca 100644 --- a/v2/core/card/specs.py +++ b/v2/core/card/specs.py @@ -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 diff --git a/v2/core/enums.py b/v2/core/enums.py index 82e938747..2c9d31d76 100644 --- a/v2/core/enums.py +++ b/v2/core/enums.py @@ -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" diff --git a/v2/core/loop/sampler.py b/v2/core/loop/sampler.py index 79ecfef6c..e66b50a3d 100644 --- a/v2/core/loop/sampler.py +++ b/v2/core/loop/sampler.py @@ -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) diff --git a/v2/core/parity/ladder.py b/v2/core/parity/ladder.py index df437d089..ac07dc7f6 100644 --- a/v2/core/parity/ladder.py +++ b/v2/core/parity/ladder.py @@ -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 diff --git a/v2/core/program/workflow.py b/v2/core/program/workflow.py index 2a163232b..bc460a6c5 100644 --- a/v2/core/program/workflow.py +++ b/v2/core/program/workflow.py @@ -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]) diff --git a/v2/core/request/params.py b/v2/core/request/params.py index 10f761842..fefd0cab6 100644 --- a/v2/core/request/params.py +++ b/v2/core/request/params.py @@ -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) diff --git a/v2/core/types.py b/v2/core/types.py index da9fb8be9..43da217e0 100644 --- a/v2/core/types.py +++ b/v2/core/types.py @@ -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 diff --git a/v2/examples.py b/v2/examples.py index a4826dd41..6bbd4c1f9 100644 --- a/v2/examples.py +++ b/v2/examples.py @@ -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) diff --git a/v2/platform/__init__.py b/v2/platform/__init__.py index 4c431fafc..7ef766ef3 100644 --- a/v2/platform/__init__.py +++ b/v2/platform/__init__.py @@ -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 diff --git a/v2/platform/backends/accel.py b/v2/platform/backends/accel.py index bd5b293c5..7f96fdcf7 100644 --- a/v2/platform/backends/accel.py +++ b/v2/platform/backends/accel.py @@ -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) diff --git a/v2/platform/backends/toy.py b/v2/platform/backends/toy.py index dd4233570..54fa1829e 100644 --- a/v2/platform/backends/toy.py +++ b/v2/platform/backends/toy.py @@ -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). diff --git a/v2/platform/registry.py b/v2/platform/registry.py index e71a9d0cf..8f8bf7ae8 100644 --- a/v2/platform/registry.py +++ b/v2/platform/registry.py @@ -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: diff --git a/v2/recipes/bagel/card.py b/v2/recipes/bagel/card.py index c4ce0d055..105cad794 100644 --- a/v2/recipes/bagel/card.py +++ b/v2/recipes/bagel/card.py @@ -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() diff --git a/v2/recipes/common.py b/v2/recipes/common.py index 3ce1093e8..e6da6e9b5 100644 --- a/v2/recipes/common.py +++ b/v2/recipes/common.py @@ -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 diff --git a/v2/recipes/cosmos3/card.py b/v2/recipes/cosmos3/card.py index e30bc678f..366029e02 100644 --- a/v2/recipes/cosmos3/card.py +++ b/v2/recipes/cosmos3/card.py @@ -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() diff --git a/v2/recipes/flux2/card.py b/v2/recipes/flux2/card.py index 1c6fa320b..a4c844581 100644 --- a/v2/recipes/flux2/card.py +++ b/v2/recipes/flux2/card.py @@ -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, ) diff --git a/v2/recipes/ltx2/card.py b/v2/recipes/ltx2/card.py index d76f56121..42bfbb1e2 100644 --- a/v2/recipes/ltx2/card.py +++ b/v2/recipes/ltx2/card.py @@ -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() diff --git a/v2/recipes/matrixgame2/card.py b/v2/recipes/matrixgame2/card.py index 476cbf253..fcadd534b 100644 --- a/v2/recipes/matrixgame2/card.py +++ b/v2/recipes/matrixgame2/card.py @@ -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, diff --git a/v2/recipes/qwen_omni/card.py b/v2/recipes/qwen_omni/card.py index d9742d2f3..d6952a317 100644 --- a/v2/recipes/qwen_omni/card.py +++ b/v2/recipes/qwen_omni/card.py @@ -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() diff --git a/v2/recipes/wan21/i2v.py b/v2/recipes/wan21/i2v.py index 54fc78375..9115e09ce 100644 --- a/v2/recipes/wan21/i2v.py +++ b/v2/recipes/wan21/i2v.py @@ -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, diff --git a/v2/recipes/wan21/program.py b/v2/recipes/wan21/program.py index 780763bfa..c3181a129 100644 --- a/v2/recipes/wan21/program.py +++ b/v2/recipes/wan21/program.py @@ -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 diff --git a/v2/recipes/wan_causal/card.py b/v2/recipes/wan_causal/card.py index 14dc85880..cd8b9a61d 100644 --- a/v2/recipes/wan_causal/card.py +++ b/v2/recipes/wan_causal/card.py @@ -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, diff --git a/v2/recipes/wan_causal/loop.py b/v2/recipes/wan_causal/loop.py index 3a55470b1..327fa1891 100644 --- a/v2/recipes/wan_causal/loop.py +++ b/v2/recipes/wan_causal/loop.py @@ -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, diff --git a/v2/runtime/cache/classes.py b/v2/runtime/cache/classes.py index 5e6979479..0d18ae10e 100644 --- a/v2/runtime/cache/classes.py +++ b/v2/runtime/cache/classes.py @@ -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) diff --git a/v2/runtime/cache/keys.py b/v2/runtime/cache/keys.py index da773f305..760fb2473 100644 --- a/v2/runtime/cache/keys.py +++ b/v2/runtime/cache/keys.py @@ -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 diff --git a/v2/runtime/cache/manager.py b/v2/runtime/cache/manager.py index 9a0025637..37fc61d0a 100644 --- a/v2/runtime/cache/manager.py +++ b/v2/runtime/cache/manager.py @@ -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) diff --git a/v2/runtime/cudagraph.py b/v2/runtime/cudagraph.py index 62de6ec0a..584d3d4b3 100644 --- a/v2/runtime/cudagraph.py +++ b/v2/runtime/cudagraph.py @@ -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.""" diff --git a/v2/runtime/memory/allocator.py b/v2/runtime/memory/allocator.py index 026ab6a7b..d236b0001 100644 --- a/v2/runtime/memory/allocator.py +++ b/v2/runtime/memory/allocator.py @@ -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 diff --git a/v2/tests/test_cache.py b/v2/tests/test_cache.py index 5e23aa318..286f5b5b3 100644 --- a/v2/tests/test_cache.py +++ b/v2/tests/test_cache.py @@ -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 diff --git a/v2/tests/test_cudagraph.py b/v2/tests/test_cudagraph.py index 8453b6c34..a7740c7bc 100644 --- a/v2/tests/test_cudagraph.py +++ b/v2/tests/test_cudagraph.py @@ -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... diff --git a/v2/tests/test_review_fixes.py b/v2/tests/test_review_fixes.py index fc2f21ce9..216c9f897 100644 --- a/v2/tests/test_review_fixes.py +++ b/v2/tests/test_review_fixes.py @@ -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 diff --git a/v2/tests/test_thinker_talker.py b/v2/tests/test_thinker_talker.py index 44d78378e..fe2cc74b9 100644 --- a/v2/tests/test_thinker_talker.py +++ b/v2/tests/test_thinker_talker.py @@ -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 -