From 8d1c8fa01a4b6fa3c2ab41c8fa8b49afb91518e2 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 7 Jun 2026 13:20:16 -0700 Subject: [PATCH] [feat] cosmos3 PR4: text reasoning (und pathway + lm_head), bit-exact MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the reasoning (VLM text-generation) capability. The framework reasoner uses only the und (causal) pathway weights (no _moe_gen) + embed_tokens / norm / lm_head — all already present in the native DiT — so a text-only forward + lm_head reproduces it with no new model code. - `cosmos3_generate_reasoner_text(transformer, input_ids, max_new_tokens)`: greedy text reasoning via the und causal backbone + lm_head (text-only prefill, re-prefill per step; KV-cache fast path is a later optimization). Mirrors the framework `generate_reasoner_text` (text-only path). - Parity (`test_cosmos3_reasoning_parity.py`, framework = oracle): * prefill logits (native text-only forward + lm_head vs framework `reasoner_forward` + lm_head, per-position): max abs diff = 0.0, mean = 0.0; * greedy generation vs framework `generate_reasoner_text(do_sample=False)`: token-for-token identical across seeds. Text reasoning is fully supported + bit-exact. Image-conditioned reasoning additionally needs the Qwen3-VL `vision_encoder` (a transformers ViT, tracked separately). Full suite: 146 passed, 0 skipped. --- .../basic/cosmos3/cosmos3_pipeline.py | 54 ++++++++ .../cosmos3/test_cosmos3_reasoning_parity.py | 123 ++++++++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index 07e38dcc7..5079d9fcd 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -109,6 +109,60 @@ def cosmos3_tokenize_caption( return list(token_ids) +# =========================================================================== +# Reasoning (VLM text generation) — und (causal) pathway + lm_head +# =========================================================================== +def cosmos3_generate_reasoner_text( + transformer: Any, + input_ids: list[int], + max_new_tokens: int, + *, + eos_token_id: int | list[int] | None = None, +) -> list[int]: + """Greedy text reasoning via the und (causal) backbone + ``lm_head``. + + Mirrors the framework ``generate_reasoner_text`` (text-only prefill, greedy): + only the und-pathway weights (no ``_moe_gen``) + ``embed_tokens`` / ``norm`` / + ``lm_head`` participate; the generation pathway and the VFM multimodal + embedders are bypassed (no vision/sound/action tokens). Token-for-token + identical to the framework reasoner (``test_cosmos3_reasoning_parity``). + + Re-prefills each step (no KV cache) — correctness-first; a KV-cache fast path + is a later optimization. Returns the newly generated token ids. + """ + device = next(transformer.parameters()).device + ids = [int(x) for x in input_ids] + eos: set[int] = set() + if eos_token_id is not None: + eos = {int(eos_token_id)} if isinstance(eos_token_id, int) else {int(x) for x in eos_token_id} + + new_tokens: list[int] = [] + for _ in range(int(max_new_tokens)): + n = len(ids) + pos = torch.arange(n).unsqueeze(0).expand(3, -1).contiguous().to(device) + out = transformer( + text_ids=torch.tensor(ids, device=device, dtype=torch.long), + text_indexes=torch.arange(n, device=device), + position_ids=pos, + sequence_length=n, + split_lens=[n], + attn_modes=["causal"], + vision_tokens=[], + vision_token_shapes=[], + vision_sequence_indexes=torch.empty(0, dtype=torch.long, device=device), + vision_timesteps=torch.empty(0, device=device), + vision_mse_loss_indexes=torch.empty(0, dtype=torch.long, device=device), + vision_noisy_frame_indexes=[], + ) + logits = transformer.lm_head(out["last_hidden_state"][n - 1]) # [vocab] + nxt = int(logits.argmax().item()) + ids.append(nxt) + new_tokens.append(nxt) + if nxt in eos: + break + return new_tokens + + # =========================================================================== # VAE encode/decode bridge (normalize / denormalize, matching the framework) # =========================================================================== diff --git a/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py b/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py new file mode 100644 index 000000000..346782437 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 text reasoning vs the framework. + +The omni model's reasoning (VLM text generation) uses ONLY the und (causal) +pathway weights (no ``_moe_gen``) + ``embed_tokens`` / ``norm`` / ``lm_head``; +the generation pathway and all multimodal embedders are bypassed. FastVideo's +native ``Cosmos3VFMTransformer`` already contains exactly those (the und branch +of the dual-pathway forward + ``lm_head``), so a text-only forward + ``lm_head`` +reproduces the framework reasoner. + +This pins: + * **prefill logits** — native text-only forward + ``lm_head`` vs the framework + ``language_model.model.reasoner_forward`` + ``lm_head`` (per-position); and + * **greedy generation** — ``cosmos3_generate_reasoner_text`` vs the framework + ``generate_reasoner_text(do_sample=False)`` (token-for-token). + +Framework is the parity ORACLE (CPU/float32 via SDPA monkey-patch). Text-only +(image-conditioned reasoning additionally needs the Qwen3-VL ``vision_encoder``, +tracked separately). + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py -q -s +""" +from __future__ import annotations + +import pytest +import torch + +cosmos_framework = pytest.importorskip( + "cosmos_framework", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) + +from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # noqa: E402 + cosmos3_generate_reasoner_text, +) + +from .test_cosmos3_dit_parity import _copy_weights # noqa: E402 +from .test_cosmos3_dit_parity_mrope import ( # noqa: E402 + _build_tiny_cosmos3_mrope, + _build_tiny_fastvideo_dit_mrope, +) +from .test_cosmos3_reference_forward import _apply_sdpa_patches # noqa: E402 + +pytestmark = [pytest.mark.local] +_apply_sdpa_patches() + + +def _framework_prefill_logits(vfm, input_ids: torch.Tensor) -> torch.Tensor: + """Framework reasoner text-only prefill logits ``[1, T, vocab]`` (the oracle). + + Mirrors ``_impl_generate_reasoner_text`` prefill: ``model.reasoner_forward`` + then ``lm_head`` (here over ALL positions, not just the last).""" + from cosmos_framework.model.vfm.mot.unified_mot import ReasonerKVCache + + causal_lm = vfm.language_model + model = causal_lm.model + cache = ReasonerKVCache.empty(num_layers=len(model.layers)) + hidden = model.reasoner_forward(input_ids.unsqueeze(0), cache=cache) # [1,T,hidden] + return causal_lm.lm_head(hidden) # [1,T,vocab] + + +def _native_prefill_logits(dit, input_ids: torch.Tensor) -> torch.Tensor: + """Native text-only forward + ``lm_head`` -> ``[T, vocab]``.""" + n = input_ids.numel() + pos = torch.arange(n).unsqueeze(0).expand(3, -1).contiguous() + out = dit( + text_ids=input_ids.to(torch.long), + text_indexes=torch.arange(n), + position_ids=pos, + sequence_length=n, + split_lens=[n], + attn_modes=["causal"], + vision_tokens=[], + vision_token_shapes=[], + vision_sequence_indexes=torch.empty(0, dtype=torch.long), + vision_timesteps=torch.empty(0), + vision_mse_loss_indexes=torch.empty(0, dtype=torch.long), + vision_noisy_frame_indexes=[], + ) + return dit.lm_head(out["last_hidden_state"]) # [T, vocab] + + +def _diffs(a, b): + d = (a - b).abs() + return d.max().item(), d.mean().item() + + +class TestCosmos3ReasoningParity: + + def _build(self, seed_model=42, num_layers=2): + vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers) + dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers) + _copy_weights(vfm, dit) + return vfm, dit + + @pytest.mark.parametrize("n_text", [4, 8, 12]) + def test_reasoner_prefill_logits_match_framework(self, n_text): + vfm, dit = self._build() + torch.manual_seed(n_text) + input_ids = torch.randint(0, 60, (n_text,)) + with torch.no_grad(): + fw = _framework_prefill_logits(vfm, input_ids)[0] # [T, vocab] + fv = _native_prefill_logits(dit, input_ids) # [T, vocab] + assert fw.shape == fv.shape, f"shape fw={fw.shape} fv={fv.shape}" + mx, mn = _diffs(fv, fw) + print(f"\n[reasoner_prefill n={n_text}] logits max abs diff = {mx:.3e} mean abs diff = {mn:.3e}") + torch.testing.assert_close(fv, fw, atol=1e-4, rtol=1e-3) + # The greedy decisions (argmax per position) must agree exactly. + assert torch.equal(fv.argmax(-1), fw.argmax(-1)) + + @pytest.mark.parametrize("seed", [0, 1, 2, 3]) + def test_greedy_reasoning_matches_framework(self, seed): + vfm, dit = self._build() + torch.manual_seed(seed) + input_ids = torch.randint(0, 60, (6,)) + fw = vfm.generate_reasoner_text( + input_ids.unsqueeze(0), max_new_tokens=8, do_sample=False, return_only_new_tokens=True, + )[0].tolist() + fv = cosmos3_generate_reasoner_text(dit, input_ids.tolist(), max_new_tokens=8) + print(f"\n[greedy_reason seed={seed}] framework={fw} native={fv}") + assert fv == fw, f"token mismatch: native={fv} framework={fw}"