[feat] cosmos3 PR4: vision_encoder parity (transformers Qwen3-VL == framework)
Image-conditioned reasoning needs the Qwen3-VL vision_encoder. The checkpoint's vision_encoder is a standard transformers `Qwen3VLVisionModel` (the framework ships its own copy of the same model). Like the Qwen2 tokenizer, FastVideo reuses the transformers model (no diffusers) rather than re-porting a 27-layer ViT. Parity (`test_cosmos3_vision_encoder_parity.py`): transformers `Qwen3VLVisionModel` vs the framework `Qwen3VLVisionModel`, tiny config, weights copied across, identical (hidden_states, grid_thw) forward — embeds max abs diff = 0.0, mean = 0.0 across grids. Also verified on the REAL 1.15 GB checkpoint: both strict-load and produce identical [N, out_hidden] embeds (max=mean=0.0). All five omni modalities now have bit-exact framework parity: video gen (T2V/I2V/T2I), audio gen (t2vs), action gen, text reasoning, vision encoding. Remaining: wire image-conditioned reasoning (deepstack multimodal prefill joining the two already-proven components). Full suite: 149 passed.
This commit is contained in:
@@ -0,0 +1,96 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Numerical-parity test: Cosmos3 vision_encoder (transformers) vs the framework.
|
||||
|
||||
Image-conditioned reasoning needs the Qwen3-VL ``vision_encoder``. The Cosmos3
|
||||
checkpoint's ``vision_encoder`` is a standard ``transformers`` ``Qwen3VLVisionModel``
|
||||
(``architectures: ["Qwen3VLVisionModel"]``); the framework ships its own copy
|
||||
(``cosmos_framework.model.vfm.vlm.qwen3_vl.qwen3_vl.Qwen3VLVisionModel``). Like the
|
||||
Qwen2 tokenizer, FastVideo reuses the ``transformers`` model (no diffusers); this
|
||||
pins it bit-for-bit against the framework's implementation.
|
||||
|
||||
Both built tiny from the SAME vision config; ``transformers`` weights copied into
|
||||
the framework model; identical ``(hidden_states, grid_thw)`` forward. The
|
||||
framework is the parity ORACLE.
|
||||
|
||||
(Separately verified on the REAL 1.15 GB checkpoint: both strict-load and produce
|
||||
identical ``[N, out_hidden]`` embeds, max abs diff = 0.0.)
|
||||
|
||||
Run:
|
||||
cd <worktree> && <fv-cosmos3 python> -m pytest \
|
||||
tests/local_tests/cosmos3/test_cosmos3_vision_encoder_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.",
|
||||
)
|
||||
transformers = pytest.importorskip("transformers")
|
||||
|
||||
pytestmark = [pytest.mark.local]
|
||||
|
||||
_TINY_VISION_CFG = dict(
|
||||
depth=2,
|
||||
hidden_size=32,
|
||||
num_heads=2,
|
||||
intermediate_size=64,
|
||||
patch_size=16,
|
||||
temporal_patch_size=2,
|
||||
spatial_merge_size=2,
|
||||
in_channels=3,
|
||||
out_hidden_size=64,
|
||||
num_position_embeddings=64,
|
||||
deepstack_visual_indexes=[1],
|
||||
hidden_act="gelu_pytorch_tanh",
|
||||
initializer_range=0.02,
|
||||
)
|
||||
|
||||
|
||||
def _diffs(a, b):
|
||||
d = (a - b).abs()
|
||||
return d.max().item(), d.mean().item()
|
||||
|
||||
|
||||
def _to_tensor(out):
|
||||
return out[0] if isinstance(out, (tuple, list)) else out
|
||||
|
||||
|
||||
class TestCosmos3VisionEncoderParity:
|
||||
|
||||
@pytest.mark.parametrize(("grid_h", "grid_w"), [(2, 2), (4, 4), (4, 6)])
|
||||
def test_vision_encoder_matches_framework(self, grid_h, grid_w):
|
||||
from cosmos_framework.model.vfm.vlm.qwen3_vl.configuration_qwen3_vl import (
|
||||
Qwen3VLVisionConfig as FwVisionConfig,
|
||||
)
|
||||
from cosmos_framework.model.vfm.vlm.qwen3_vl.qwen3_vl import (
|
||||
Qwen3VLVisionModel as FwVisionModel,
|
||||
)
|
||||
from transformers import Qwen3VLVisionModel as TfVisionModel
|
||||
from transformers.models.qwen3_vl.configuration_qwen3_vl import (
|
||||
Qwen3VLVisionConfig as TfVisionConfig,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
tf_model = TfVisionModel(TfVisionConfig(**_TINY_VISION_CFG)).eval()
|
||||
fw_model = FwVisionModel(FwVisionConfig(**_TINY_VISION_CFG)).eval()
|
||||
# transformers is the unit under test; copy its weights into the framework
|
||||
# oracle (identical key surface — both are the same Qwen3-VL ViT).
|
||||
missing, unexpected = fw_model.load_state_dict(tf_model.state_dict(), strict=False)
|
||||
assert not missing and not unexpected, f"key mismatch: missing={missing[:3]} unexpected={unexpected[:3]}"
|
||||
|
||||
in_dim = _TINY_VISION_CFG["in_channels"] * _TINY_VISION_CFG["temporal_patch_size"] * (
|
||||
_TINY_VISION_CFG["patch_size"] ** 2)
|
||||
seq = grid_h * grid_w
|
||||
hidden_states = torch.randn(seq, in_dim)
|
||||
grid_thw = torch.tensor([[1, grid_h, grid_w]], dtype=torch.long)
|
||||
|
||||
with torch.no_grad():
|
||||
tf_out = _to_tensor(tf_model(hidden_states, grid_thw))
|
||||
fw_out = _to_tensor(fw_model(hidden_states, grid_thw))
|
||||
assert tf_out.shape == fw_out.shape, f"shape tf={tf_out.shape} fw={fw_out.shape}"
|
||||
mx, mn = _diffs(tf_out, fw_out)
|
||||
print(f"\n[vision_encoder grid={grid_h}x{grid_w}] embeds max abs diff = {mx:.3e} mean abs diff = {mn:.3e}")
|
||||
torch.testing.assert_close(tf_out, fw_out, atol=1e-5, rtol=1e-4)
|
||||
Reference in New Issue
Block a user