1551 lines
59 KiB
Python
1551 lines
59 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import weakref
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Provide the TP metadata required by vLLM parallel linear layers."""
|
|
from vllm.model_executor import parameter
|
|
from vllm.model_executor.layers import linear
|
|
|
|
from vllm_omni.diffusion.models.cosmos3 import (
|
|
transformer_cosmos3,
|
|
transformer_cosmos3_edge,
|
|
)
|
|
|
|
monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
|
|
|
|
@pytest.fixture
|
|
def accelerator_device() -> torch.device:
|
|
"""Provide an accelerator device, skipping when none is available."""
|
|
if current_omni_platform.get_device_count() == 0:
|
|
pytest.skip("Accelerator required for this test")
|
|
return current_omni_platform.get_torch_device(0)
|
|
|
|
|
|
def _tiny_cosmos3_config(**overrides):
|
|
config = {
|
|
"hidden_size": 8,
|
|
"num_hidden_layers": 0,
|
|
"num_attention_heads": 2,
|
|
"num_key_value_heads": 2,
|
|
"head_dim": 4,
|
|
"intermediate_size": 16,
|
|
"vocab_size": 32,
|
|
"latent_patch_size": 1,
|
|
"latent_channel": 2,
|
|
"rope_scaling": {"mrope_section": [1, 1, 0]},
|
|
}
|
|
config.update(overrides)
|
|
return config
|
|
|
|
|
|
def _tiny_cosmos3_edge_config(**overrides):
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
COSMOS3_EDGE_BACKBONE_TYPE,
|
|
)
|
|
|
|
config = _tiny_cosmos3_config(
|
|
latent_patch_size=2,
|
|
latent_channel=48,
|
|
temporal_compression_factor=4,
|
|
backbone_type=COSMOS3_EDGE_BACKBONE_TYPE,
|
|
qk_norm_for_text=False,
|
|
)
|
|
config.update(overrides)
|
|
return config
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("config_kind", "expected_language", "expected_gen"),
|
|
[
|
|
("flat", "transformer", "transformer"),
|
|
("transformer", "transformer", "transformer"),
|
|
("other_component", None, None),
|
|
("language_model", "language", None),
|
|
("gen_layers", None, "gen"),
|
|
("transformer_language_override", "language", "transformer"),
|
|
("transformer_gen_disabled", "transformer", None),
|
|
("independent_subcomponents", "language", "gen"),
|
|
],
|
|
)
|
|
def test_transformer_resolves_global_and_subcomponent_quant_configs(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
config_kind: str,
|
|
expected_language: str | None,
|
|
expected_gen: str | None,
|
|
) -> None:
|
|
"""Resolve pipeline-level defaults and Cosmos3 subcomponent overrides."""
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
|
|
|
received: list[tuple[str, object]] = []
|
|
configs = {
|
|
"transformer": object(),
|
|
"language": object(),
|
|
"gen": object(),
|
|
}
|
|
|
|
class _StubLanguageModel(nn.Module):
|
|
def __init__(self, *, quant_config, **kwargs) -> None:
|
|
del kwargs
|
|
super().__init__()
|
|
received.append(("language_model", quant_config))
|
|
self.layers = nn.ModuleList()
|
|
|
|
class _StubGenDecoderLayer(nn.Module):
|
|
def __init__(self, *, quant_config, **kwargs) -> None:
|
|
del kwargs
|
|
super().__init__()
|
|
received.append(("gen_layers", quant_config))
|
|
|
|
monkeypatch.setattr(Cosmos3VFMTransformer, "_language_model_cls", _StubLanguageModel)
|
|
monkeypatch.setattr(transformer_cosmos3, "Cosmos3GenDecoderLayer", _StubGenDecoderLayer)
|
|
|
|
if config_kind == "flat":
|
|
top_level_config = configs["transformer"]
|
|
elif config_kind == "transformer":
|
|
top_level_config = ComponentQuantizationConfig({"transformer": configs["transformer"]})
|
|
elif config_kind == "other_component":
|
|
top_level_config = ComponentQuantizationConfig({"vae": configs["transformer"]})
|
|
elif config_kind == "language_model":
|
|
top_level_config = ComponentQuantizationConfig({"language_model": configs["language"]})
|
|
elif config_kind == "gen_layers":
|
|
top_level_config = ComponentQuantizationConfig({"gen_layers": configs["gen"]})
|
|
elif config_kind == "transformer_language_override":
|
|
top_level_config = ComponentQuantizationConfig(
|
|
{"transformer": configs["transformer"], "language_model": configs["language"]}
|
|
)
|
|
elif config_kind == "transformer_gen_disabled":
|
|
top_level_config = ComponentQuantizationConfig({"transformer": configs["transformer"], "gen_layers": None})
|
|
else:
|
|
top_level_config = ComponentQuantizationConfig(
|
|
{"language_model": configs["language"], "gen_layers": configs["gen"]}
|
|
)
|
|
|
|
Cosmos3VFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1),
|
|
dtype=torch.float32,
|
|
quantization_config=top_level_config,
|
|
custom_pipeline_args={},
|
|
model_config={},
|
|
)
|
|
)
|
|
|
|
assert received == [
|
|
("language_model", None if expected_language is None else configs[expected_language]),
|
|
("gen_layers", None if expected_gen is None else configs[expected_gen]),
|
|
]
|
|
|
|
|
|
def test_transformer_component_quant_reaches_real_linear(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""#7323 regression: unwrap ComponentQuantizationConfig before ColumnParallelLinear.
|
|
|
|
Passing the router through unchanged makes ``resolve(language_model.*)`` miss
|
|
a ``transformer``-only map and raise ``ValueError: All linear layers should
|
|
support quant method``. A leaf config must reach the real linear path after
|
|
resolution. Attention backends are stubbed so this stays CPU-safe; MLP and
|
|
QKV linears remain real ``ColumnParallelLinear`` instances.
|
|
"""
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
|
|
Cosmos3GatedMLP,
|
|
Cosmos3VFMTransformer,
|
|
_resolve_cosmos3_quant_configs,
|
|
)
|
|
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
|
|
|
class _LeafQuantConfig(QuantizationConfig):
|
|
def get_name(self) -> str:
|
|
return "test_leaf"
|
|
|
|
def get_quant_method(self, layer, prefix): # noqa: ANN001
|
|
del layer, prefix
|
|
return UnquantizedLinearMethod()
|
|
|
|
@classmethod
|
|
def get_supported_act_dtypes(cls):
|
|
return [torch.float32]
|
|
|
|
def get_min_capability(self) -> int:
|
|
return 0
|
|
|
|
@classmethod
|
|
def from_config(cls, config): # noqa: ANN001
|
|
raise NotImplementedError
|
|
|
|
def get_config_filenames(self) -> list[str]:
|
|
return []
|
|
|
|
class _StubFrameworkAttention(nn.Module):
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
del args, kwargs
|
|
super().__init__()
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "FrameworkAttention", _StubFrameworkAttention)
|
|
|
|
leaf = _LeafQuantConfig()
|
|
router = ComponentQuantizationConfig({"transformer": leaf})
|
|
|
|
with pytest.raises(ValueError, match="All linear layers should support quant method"):
|
|
Cosmos3GatedMLP(
|
|
hidden_size=8,
|
|
intermediate_size=16,
|
|
quant_config=router,
|
|
prefix="language_model.layers.0.mlp",
|
|
)
|
|
|
|
language_qc, gen_qc = _resolve_cosmos3_quant_configs(router)
|
|
assert language_qc is leaf
|
|
assert gen_qc is leaf
|
|
|
|
model = Cosmos3VFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1),
|
|
dtype=torch.float32,
|
|
quantization_config=router,
|
|
custom_pipeline_args={},
|
|
model_config={},
|
|
)
|
|
)
|
|
mlp = model.language_model.layers[0].mlp
|
|
assert isinstance(mlp.gate_proj.quant_method, UnquantizedLinearMethod)
|
|
assert isinstance(model.gen_layers[0].mlp.gate_proj.quant_method, UnquantizedLinearMethod)
|
|
assert isinstance(
|
|
model.language_model.layers[0].self_attn.to_q.quant_method,
|
|
UnquantizedLinearMethod,
|
|
)
|
|
|
|
|
|
def test_transformer_pathway_default_keeps_nested_layer_override() -> None:
|
|
"""Pathway leaf plus a more-specific key must still longest-prefix match."""
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
|
|
Cosmos3GatedMLP,
|
|
_resolve_cosmos3_quant_configs,
|
|
)
|
|
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
|
|
|
class _TaggedLinearMethod(UnquantizedLinearMethod):
|
|
def __init__(self, tag: str) -> None:
|
|
super().__init__()
|
|
self.tag = tag
|
|
|
|
class _TaggedQuantConfig(QuantizationConfig):
|
|
def __init__(self, tag: str) -> None:
|
|
super().__init__()
|
|
self.tag = tag
|
|
|
|
def get_name(self) -> str:
|
|
return self.tag
|
|
|
|
def get_quant_method(self, layer, prefix): # noqa: ANN001
|
|
del layer, prefix
|
|
return _TaggedLinearMethod(self.tag)
|
|
|
|
@classmethod
|
|
def get_supported_act_dtypes(cls):
|
|
return [torch.float32]
|
|
|
|
def get_min_capability(self) -> int:
|
|
return 0
|
|
|
|
@classmethod
|
|
def from_config(cls, config): # noqa: ANN001
|
|
raise NotImplementedError
|
|
|
|
def get_config_filenames(self) -> list[str]:
|
|
return []
|
|
|
|
pathway = _TaggedQuantConfig("pathway")
|
|
mlp_override = _TaggedQuantConfig("mlp")
|
|
language_qc, gen_qc = _resolve_cosmos3_quant_configs(
|
|
ComponentQuantizationConfig(
|
|
{
|
|
"language_model": pathway,
|
|
"gen_layers": pathway,
|
|
"language_model.layers.0.mlp": mlp_override,
|
|
}
|
|
)
|
|
)
|
|
|
|
assert isinstance(language_qc, ComponentQuantizationConfig)
|
|
assert language_qc.resolve("language_model.layers.0.mlp.gate_proj") is mlp_override
|
|
assert language_qc.resolve("language_model.layers.0.self_attn.to_q") is pathway
|
|
assert language_qc.resolve("language_model.layers.1.mlp.gate_proj") is pathway
|
|
assert gen_qc is pathway
|
|
|
|
mlp0 = Cosmos3GatedMLP(
|
|
hidden_size=8,
|
|
intermediate_size=16,
|
|
quant_config=language_qc,
|
|
prefix="language_model.layers.0.mlp",
|
|
)
|
|
mlp1 = Cosmos3GatedMLP(
|
|
hidden_size=8,
|
|
intermediate_size=16,
|
|
quant_config=language_qc,
|
|
prefix="language_model.layers.1.mlp",
|
|
)
|
|
assert mlp0.gate_proj.quant_method.tag == "mlp"
|
|
assert mlp1.gate_proj.quant_method.tag == "pathway"
|
|
|
|
|
|
def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
|
|
compute_mrope_position_ids_action,
|
|
compute_mrope_position_ids_sound,
|
|
compute_mrope_position_ids_text,
|
|
compute_mrope_position_ids_vision,
|
|
)
|
|
|
|
text_ids, text_offset = compute_mrope_position_ids_text(num_tokens=3, temporal_offset=5)
|
|
assert text_ids.tolist() == [[5, 6, 7], [5, 6, 7], [5, 6, 7]]
|
|
assert text_offset == 8
|
|
|
|
vision_ids, vision_offset = compute_mrope_position_ids_vision(2, 2, 3, temporal_offset=10, fps=None)
|
|
assert vision_ids.shape == (3, 12)
|
|
assert vision_ids[0].tolist() == [10] * 6 + [11] * 6
|
|
assert vision_offset == 12
|
|
|
|
modulated_ids, modulated_offset = compute_mrope_position_ids_vision(
|
|
2,
|
|
1,
|
|
1,
|
|
temporal_offset=10,
|
|
fps=12.0,
|
|
base_fps=24.0,
|
|
temporal_compression_factor=4,
|
|
)
|
|
torch.testing.assert_close(modulated_ids[0], torch.tensor([10.0, 12.0]))
|
|
assert modulated_offset == 13
|
|
|
|
sound_ids, sound_offset = compute_mrope_position_ids_sound(3, temporal_offset=10, sound_latent_fps=25.0)
|
|
torch.testing.assert_close(sound_ids[0], torch.tensor([10.0, 10.96, 11.92]))
|
|
assert sound_offset == 12
|
|
|
|
action_ids, action_offset = compute_mrope_position_ids_action(3, temporal_offset=10, action_fps=None)
|
|
assert action_ids.tolist() == [[11, 12, 13], [0, 0, 0], [0, 0, 0]]
|
|
assert action_offset == 14
|
|
|
|
|
|
def test_timestep_embedder_stores_frequencies_in_fp32() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
|
|
TimestepEmbedder,
|
|
)
|
|
|
|
embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16)
|
|
|
|
assert embedder.freqs.dtype == torch.float32
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("key", "value"),
|
|
[
|
|
("qk_norm_for_diffusion", False),
|
|
("qk_norm_for_text", False),
|
|
("position_embedding_type", "rotary"),
|
|
("unified_3d_mrope_reset_spatial_ids", False),
|
|
("joint_attn_implementation", "one_way"),
|
|
],
|
|
)
|
|
def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
|
|
with pytest.raises(ValueError, match=f"{key}="):
|
|
Cosmos3VFMTransformer._validate_supported_config({key: value})
|
|
Cosmos3VFMTransformer._validate_supported_config({})
|
|
Cosmos3VFMTransformer._validate_supported_config(None)
|
|
|
|
|
|
def test_edge_config_resolves_nemotron_defaults() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeVFMTransformer,
|
|
)
|
|
|
|
model = Cosmos3EdgeVFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=_tiny_cosmos3_edge_config(
|
|
layer_norm_epsilon=1e-5,
|
|
rope_scaling={},
|
|
rope_parameters={"mrope_section": [24, 20, 20]},
|
|
),
|
|
dtype=torch.float32,
|
|
)
|
|
)
|
|
|
|
assert model.rms_norm_eps == 1e-5
|
|
assert model.rope_theta == 100_000_000
|
|
assert model.mrope_section == [24, 20, 20]
|
|
assert model.latent_channel_size == 48
|
|
assert model.latent_patch_size == 2
|
|
assert model.temporal_compression_factor == 4
|
|
|
|
|
|
def test_edge_config_requires_backbone_type() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeVFMTransformer,
|
|
)
|
|
|
|
config = _tiny_cosmos3_edge_config()
|
|
config.pop("backbone_type")
|
|
|
|
with pytest.raises(ValueError, match="must declare backbone_type"):
|
|
Cosmos3EdgeVFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=config,
|
|
dtype=torch.float32,
|
|
)
|
|
)
|
|
|
|
|
|
def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeVFMTransformer,
|
|
Cosmos3Relu2MLP,
|
|
)
|
|
|
|
model = Cosmos3EdgeVFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=_tiny_cosmos3_edge_config(
|
|
num_hidden_layers=1,
|
|
use_und_k_norm_for_gen=False,
|
|
),
|
|
dtype=torch.float32,
|
|
)
|
|
)
|
|
|
|
layer = model.language_model.layers[0]
|
|
gen_layer = model.gen_layers[0]
|
|
assert isinstance(layer.mlp, Cosmos3Relu2MLP)
|
|
assert not hasattr(layer.mlp, "gate_proj")
|
|
assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP)
|
|
assert not hasattr(gen_layer.mlp, "gate_proj")
|
|
assert not hasattr(layer.self_attn, "norm_q")
|
|
assert not hasattr(layer.self_attn, "norm_k")
|
|
assert layer.self_attn.k_norm_und_for_gen is None
|
|
|
|
|
|
def test_edge_creates_gen_k_norm_when_enabled() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeVFMTransformer,
|
|
)
|
|
|
|
model = Cosmos3EdgeVFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=_tiny_cosmos3_edge_config(
|
|
num_hidden_layers=1,
|
|
qk_norm_for_diffusion=True,
|
|
use_und_k_norm_for_gen=True,
|
|
),
|
|
dtype=torch.float32,
|
|
)
|
|
)
|
|
|
|
k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen
|
|
assert model.use_und_k_norm_for_gen is True
|
|
assert isinstance(k_norm, RMSNorm)
|
|
assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm)
|
|
assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm)
|
|
|
|
|
|
def test_edge_supports_no_gen_qk_norm_variant() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeVFMTransformer,
|
|
)
|
|
|
|
model = Cosmos3EdgeVFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config=_tiny_cosmos3_edge_config(
|
|
num_hidden_layers=1,
|
|
qk_norm_for_diffusion=False,
|
|
use_und_k_norm_for_gen=False,
|
|
),
|
|
dtype=torch.float32,
|
|
)
|
|
)
|
|
|
|
assert model.qk_norm_for_diffusion is False
|
|
assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None
|
|
assert model.gen_layers[0].cross_attention.qk_norm is False
|
|
assert not hasattr(model.gen_layers[0].cross_attention, "norm_q")
|
|
assert not hasattr(model.gen_layers[0].cross_attention, "norm_k")
|
|
|
|
|
|
def test_edge_validates_required_relu2_weights() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeVFMTransformer,
|
|
)
|
|
|
|
model = object.__new__(Cosmos3EdgeVFMTransformer)
|
|
nn.Module.__init__(model)
|
|
model.num_hidden_layers = 1
|
|
model.qk_norm_for_diffusion = True
|
|
model.use_und_k_norm_for_gen = True
|
|
|
|
complete = {
|
|
"transformer.language_model.layers.0.mlp.up_proj.weight",
|
|
"transformer.language_model.layers.0.mlp.down_proj.weight",
|
|
"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
|
|
"transformer.gen_layers.0.mlp.up_proj.weight",
|
|
"transformer.gen_layers.0.mlp.down_proj.weight",
|
|
}
|
|
model.validate_loaded_weights(complete)
|
|
|
|
mlp_weights = complete - {
|
|
"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
|
|
}
|
|
for missing_weight in mlp_weights:
|
|
with pytest.raises(ValueError, match="missing required weights"):
|
|
model.validate_loaded_weights(complete - {missing_weight})
|
|
|
|
missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"}
|
|
with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"):
|
|
model.validate_loaded_weights(missing_k_norm)
|
|
|
|
model.use_und_k_norm_for_gen = False
|
|
model.validate_loaded_weights(missing_k_norm)
|
|
|
|
|
|
def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeCausalAttention,
|
|
)
|
|
|
|
class ScaleProjection(nn.Module):
|
|
def __init__(self, scale: float) -> None:
|
|
super().__init__()
|
|
self.scale = scale
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x * self.scale
|
|
|
|
class IdentityOutput(nn.Module):
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
class RawKAttention(nn.Module):
|
|
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
|
|
del v
|
|
return q + k
|
|
|
|
attn = object.__new__(Cosmos3EdgeCausalAttention)
|
|
nn.Module.__init__(attn)
|
|
attn.hidden_size = 4
|
|
attn.num_heads = 2
|
|
attn.num_kv_heads = 2
|
|
attn.head_dim = 2
|
|
attn.num_heads_local = 2
|
|
attn.num_kv_heads_local = 2
|
|
attn.to_q = ScaleProjection(1.0)
|
|
attn.to_k = ScaleProjection(1.0)
|
|
attn.to_v = ScaleProjection(1.0)
|
|
attn.to_out = IdentityOutput()
|
|
attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5)
|
|
attn.attn = RawKAttention()
|
|
|
|
hidden_states = torch.tensor(
|
|
[[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]],
|
|
dtype=torch.float32,
|
|
)
|
|
freqs_cos = torch.ones(1, 2, 1, 2)
|
|
freqs_sin = torch.zeros(1, 2, 1, 2)
|
|
|
|
reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin)
|
|
attn.to_k.scale = 3.0
|
|
scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin)
|
|
|
|
torch.testing.assert_close(gen_k, scaled_gen_k)
|
|
torch.testing.assert_close(gen_v, scaled_gen_v)
|
|
assert not torch.allclose(reasoner_out, scaled_reasoner_out)
|
|
|
|
|
|
def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
|
|
Cosmos3EdgeLanguageModel,
|
|
Cosmos3EdgeVFMTransformer,
|
|
)
|
|
|
|
for transformer_cls, language_model_cls in (
|
|
(Cosmos3VFMTransformer, Cosmos3LanguageModel),
|
|
(Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel),
|
|
):
|
|
model = object.__new__(transformer_cls)
|
|
nn.Module.__init__(model)
|
|
model.language_model = nn.Module()
|
|
model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
|
|
model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
|
|
model.norm_moe_gen = nn.LayerNorm(2)
|
|
|
|
matched = [
|
|
name
|
|
for name, module in model.named_modules()
|
|
if any(condition(name, module) for condition in model._hsdp_shard_conditions)
|
|
]
|
|
assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
|
|
assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"]
|
|
assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"]
|
|
assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"]
|
|
|
|
model.latent_patch_size = 2
|
|
model.latent_channel_size = 3
|
|
latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
|
|
torch.testing.assert_close(
|
|
model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5),
|
|
latents,
|
|
)
|
|
|
|
|
|
def test_gen_sp_plan_auto_pads_hidden_states_and_rope_together() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
|
|
prepare_plan = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
|
|
assert set(prepare_plan) == {0, 1, 2}
|
|
assert all(spec.auto_pad for spec in prepare_plan.values())
|
|
assert [prepare_plan[index].split_dim for index in range(3)] == [1, 1, 1]
|
|
|
|
|
|
def test_gen_sp_prepare_resets_branch_local_padding_metadata() -> None:
|
|
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare
|
|
|
|
ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=12121)
|
|
hidden = torch.zeros(1, 4, 8)
|
|
rope = torch.zeros(1, 4, 1, 2)
|
|
with override_forward_context(ctx):
|
|
output = Cosmos3GenSPPrepare()(hidden, rope, rope)
|
|
|
|
assert output[0] is hidden
|
|
assert output[1] is rope
|
|
assert output[2] is rope
|
|
assert ctx.sp_padding_size == 0
|
|
assert ctx.sp_original_seq_len is None
|
|
|
|
|
|
def test_sp_attention_masks_padded_gen_and_joint_und_keys() -> None:
|
|
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
|
|
|
|
class RecordingAttention(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.metadata: Any | None = None
|
|
|
|
def forward(self, query, key, value, metadata):
|
|
del key, value
|
|
self.metadata = metadata
|
|
return query
|
|
|
|
module = object.__new__(Cosmos3CrossAttention)
|
|
nn.Module.__init__(module)
|
|
module.num_heads_local = 2
|
|
module.head_dim = 4
|
|
recording_attention = RecordingAttention()
|
|
module.attn = recording_attention
|
|
|
|
q = torch.zeros(1, 2, 2, 4)
|
|
k = torch.zeros_like(q)
|
|
v = torch.zeros_like(q)
|
|
k_und = torch.zeros(1, 3, 2, 4)
|
|
v_und = torch.zeros_like(k_und)
|
|
ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=3)
|
|
with override_forward_context(ctx):
|
|
output = module._forward_sp(q, k, v, k_und, v_und)
|
|
|
|
assert output.shape == (1, 2, 8)
|
|
assert recording_attention.metadata is not None
|
|
assert recording_attention.metadata.attn_mask.tolist() == [[True, True, True, False]]
|
|
assert recording_attention.metadata.joint_attn_mask.tolist() == [[True, True, True]]
|
|
|
|
|
|
def test_ulysses_2d_mask_uses_key_length_with_kv_only_joint_context(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.parallel import ulysses
|
|
|
|
def fake_all_to_all(
|
|
_group: object,
|
|
tensor: torch.Tensor,
|
|
_scatter_idx: int,
|
|
_gather_idx: int,
|
|
_use_sync: bool,
|
|
) -> torch.Tensor:
|
|
return torch.cat([tensor, tensor], dim=1)[:, :, :1]
|
|
|
|
monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", staticmethod(fake_all_to_all))
|
|
strategy = ulysses.UlyssesParallelAttention(
|
|
sp_group=SimpleNamespace(
|
|
ulysses_group=object(),
|
|
ulysses_world_size=2,
|
|
ulysses_rank=0,
|
|
ring_world_size=1,
|
|
),
|
|
scatter_idx=2,
|
|
gather_idx=1,
|
|
use_sync=False,
|
|
)
|
|
query = torch.zeros(1, 2, 2, 4)
|
|
key = torch.zeros_like(query)
|
|
value = torch.zeros_like(query)
|
|
joint_query = torch.zeros(1, 0, 2, 4)
|
|
joint_key = torch.zeros(1, 3, 2, 4)
|
|
metadata = AttentionMetadata(
|
|
attn_mask=torch.tensor([[True, True, True, False]]),
|
|
joint_attn_mask=torch.ones(1, 3, dtype=torch.bool),
|
|
joint_query=joint_query,
|
|
joint_key=joint_key,
|
|
joint_value=torch.zeros_like(joint_key),
|
|
joint_strategy="front",
|
|
)
|
|
|
|
query_out, key_out, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
|
|
|
|
assert query_out.shape[1] == 4
|
|
assert key_out.shape[1] == 7
|
|
assert metadata_out is not None
|
|
assert metadata_out.attn_mask.tolist() == [[True, True, True, True, True, True, False]]
|
|
|
|
|
|
def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
|
|
output = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)(
|
|
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
)
|
|
|
|
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
|
|
|
|
|
def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
class TrackingNorm(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.calls = 0
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
self.calls += 1
|
|
return hidden_states + 5.0
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
norm = TrackingNorm()
|
|
model.norm_moe_gen = norm
|
|
captured: dict[str, torch.Tensor] = {}
|
|
|
|
def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor:
|
|
del kwargs
|
|
captured["input"] = hidden_gen.detach().clone()
|
|
return hidden_gen + 2.0
|
|
|
|
monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers)
|
|
forward_kwargs = {
|
|
"hidden_states": torch.zeros(1, 2, 1, 2, 2),
|
|
"timestep": torch.tensor([1.0]),
|
|
"text_ids": torch.tensor([[1, 2]], dtype=torch.long),
|
|
"text_mask": torch.ones(1, 2, dtype=torch.long),
|
|
"video_shape": (1, 2, 2),
|
|
"fps": 24.0,
|
|
}
|
|
|
|
full_output = model(**forward_kwargs)
|
|
|
|
assert norm.calls == 1
|
|
norm.calls = 0
|
|
|
|
ctx = extract_cosmos3_context(model, **forward_kwargs)
|
|
execution_input = ctx.hidden_states.detach().clone()
|
|
execution_output = ctx.run_transformer_blocks()[0]
|
|
residual = execution_output - execution_input
|
|
extracted_output = ctx.postprocess(execution_output)
|
|
|
|
assert norm.calls == 1
|
|
expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"]
|
|
assert torch.equal(residual, expected_residual)
|
|
assert torch.equal(extracted_output, full_output)
|
|
|
|
norm.calls = 0
|
|
|
|
def fail_if_gen_layers_run(*args, **kwargs):
|
|
del args, kwargs
|
|
raise AssertionError("SeaCache hit unexpectedly executed GEN layers")
|
|
|
|
monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run)
|
|
|
|
cached_ctx = extract_cosmos3_context(model, **forward_kwargs)
|
|
cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual)
|
|
|
|
assert norm.calls == 0
|
|
assert torch.equal(cached_output, full_output)
|
|
for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
|
|
assert not hasattr(model, name)
|
|
|
|
|
|
@pytest.mark.parametrize("sequence_length", [16, 17])
|
|
@pytest.mark.parametrize("batch", [1, 2])
|
|
@pytest.mark.parametrize("rank", [0, 3])
|
|
@pytest.mark.parametrize("seacache", [False, True])
|
|
@pytest.mark.parametrize("control_count", [0, 1, 2])
|
|
@pytest.mark.parallel
|
|
@torch.inference_mode()
|
|
def test_sp_releases_full_gen_embedding_during_stack(
|
|
monkeypatch, sequence_length, batch, rank, seacache, control_count
|
|
) -> None:
|
|
from vllm_omni.diffusion.attention import selector
|
|
from vllm_omni.diffusion.cache.seacache import SeaCacheConfig, apply_sea_cache_hook
|
|
from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
|
|
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
|
|
from vllm_omni.diffusion.hooks import sequence_parallel
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
full_refs = []
|
|
full_pointers = []
|
|
layer_calls = 0
|
|
total_length = sequence_length * (control_count + 1)
|
|
uses_sp = control_count <= 1
|
|
|
|
class CheckingLayer(nn.Module):
|
|
def forward(self, hidden, **kwargs):
|
|
nonlocal layer_calls
|
|
layer_calls += 1
|
|
if full_refs:
|
|
# Inspect lifetime inside the production stack while the
|
|
# caller's prepared state is still live.
|
|
assert full_refs[-1]() is None
|
|
if layer_calls % 2 == 1:
|
|
assert hidden.untyped_storage().data_ptr() != full_pointers[-1]
|
|
assert hidden.shape == (batch, (total_length + 3) // 4, 8)
|
|
assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
|
|
assert kwargs["freqs_cos"].shape[1] == kwargs["freqs_sin"].shape[1] == hidden.shape[1]
|
|
elif not uses_sp:
|
|
assert hidden.shape == (batch, total_length, 8)
|
|
return hidden + 1
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
model.gen_layers = nn.ModuleList([CheckingLayer(), CheckingLayer()])
|
|
model.cached_kv = [(torch.empty(0), torch.empty(0)) for _ in model.gen_layers]
|
|
inputs = {
|
|
"hidden_states": torch.randn(batch, 2, 1, 1, sequence_length),
|
|
"timestep": torch.ones(batch),
|
|
"text_ids": torch.ones(batch, 2, dtype=torch.long),
|
|
"text_mask": torch.ones(batch, 2, dtype=torch.long),
|
|
"video_shape": (1, 1, sequence_length),
|
|
"control_latents": [torch.randn(batch, 2, 1, 1, sequence_length) for _ in range(control_count)],
|
|
}
|
|
gathered_reference = []
|
|
normalized_reference = []
|
|
handle = model.gen_sp_gather.register_forward_pre_hook(
|
|
lambda module, args: gathered_reference.append(args[0].clone())
|
|
)
|
|
norm_handle = model.norm_moe_gen.register_forward_hook(
|
|
lambda module, args, output: normalized_reference.append(output.clone())
|
|
)
|
|
expected = model(**inputs)
|
|
handle.remove()
|
|
norm_handle.remove()
|
|
|
|
class MaskBackend:
|
|
@staticmethod
|
|
def supports_attention_mask(spec):
|
|
return True
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (4, rank, None))
|
|
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 4)
|
|
monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
|
|
monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 4)
|
|
monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
|
|
monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
|
|
|
|
def gather(hidden, dim, validate):
|
|
shard_length = hidden.shape[1]
|
|
start = rank * shard_length
|
|
valid_length = min(shard_length, total_length - start)
|
|
reference = normalized_reference[0] if seacache else gathered_reference[0]
|
|
torch.testing.assert_close(hidden[:, :valid_length], reference[:, start : start + valid_length])
|
|
# Supply the other ranks' outputs without retaining the original input.
|
|
return torch.nn.functional.pad(reference, (0, 0, 0, -total_length % 4))
|
|
|
|
monkeypatch.setattr(sequence_parallel, "sp_gather", gather)
|
|
sequence_parallel.apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), model._sp_plan)
|
|
|
|
def record_full_embedding(module, args):
|
|
full_refs.append(weakref.ref(args[0]))
|
|
full_pointers.append(args[0].untyped_storage().data_ptr())
|
|
|
|
model.gen_sp_prepare.register_forward_pre_hook(record_full_embedding)
|
|
metadata = SimpleNamespace(step=0)
|
|
if seacache:
|
|
sea_hook = apply_sea_cache_hook(
|
|
model,
|
|
SeaCacheConfig(threshold=100.0, residual_order=0),
|
|
current_step_callback=lambda: metadata.step,
|
|
current_sigma_callback=lambda: 1.0,
|
|
num_inference_steps_callback=lambda: 4,
|
|
)
|
|
sea_hook.state_manager.set_context("cond")
|
|
layer_calls = 0
|
|
for step in range(4 if seacache else 2):
|
|
metadata.step = step
|
|
with set_forward_context():
|
|
output = model(**inputs)
|
|
torch.testing.assert_close(output, expected)
|
|
assert get_forward_context()._sp_shard_depth == 0
|
|
if seacache:
|
|
if uses_sp:
|
|
assert full_refs[-1]() is None
|
|
history = sea_hook.state_manager.get_state().history
|
|
assert history
|
|
residual = history[-1][1]
|
|
assert residual.shape == (batch, (total_length + 3) // 4 if uses_sp else total_length, 8)
|
|
assert residual.untyped_storage().nbytes() == residual.numel() * residual.element_size()
|
|
assert layer_calls == 4
|
|
assert len(full_refs) == ((4 if seacache else 2) if uses_sp else 0)
|
|
if seacache:
|
|
assert sea_hook.full_count == sea_hook.skip_count == 2
|
|
|
|
|
|
def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
class CountingNorm(nn.Module):
|
|
def __init__(self, original: nn.Module) -> None:
|
|
super().__init__()
|
|
self.original = original
|
|
self.calls = 0
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
self.calls += 1
|
|
return self.original(hidden_states)
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
norm = CountingNorm(model.norm_moe_gen)
|
|
model.norm_moe_gen = norm
|
|
|
|
output = model(
|
|
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
)
|
|
|
|
assert norm.calls == 1
|
|
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
|
|
|
|
|
def test_qwen_gen_mlp_remains_gated() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
|
|
|
|
model = Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32)
|
|
)
|
|
|
|
assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP)
|
|
assert hasattr(model.gen_layers[0].mlp, "gate_proj")
|
|
|
|
|
|
@pytest.mark.parametrize("edge", [False, True])
|
|
@pytest.mark.parametrize(
|
|
"use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")]
|
|
)
|
|
def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch):
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
|
|
from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
|
|
|
|
released = []
|
|
|
|
def release_and_check(block):
|
|
parameters = dict(block.named_parameters())
|
|
metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()}
|
|
release_module_parameters_to_meta(block)
|
|
for name, parameter in block.named_parameters():
|
|
assert parameter is parameters[name]
|
|
assert parameter.device.type == "meta"
|
|
assert parameter.__dict__ == metadata[name]
|
|
released.append(block)
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check)
|
|
monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check)
|
|
model_cls = (
|
|
transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer
|
|
)
|
|
config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
|
|
model = model_cls(
|
|
SimpleNamespace(
|
|
tf_model_config=config(num_hidden_layers=1),
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
|
|
hsdp_weight_load_strategy=strategy,
|
|
)
|
|
)
|
|
blocks = [*model.language_model.layers, *model.gen_layers]
|
|
expected_meta = use_hsdp and strategy == "pre_sharded"
|
|
assert len(released) == (len(blocks) if expected_meta else 0)
|
|
assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters())
|
|
assert model.time_embedder.linear_1.weight.device.type != "meta"
|
|
assert all(buffer.device.type != "meta" for buffer in model.buffers())
|
|
|
|
|
|
@pytest.mark.parametrize("edge", [False, True])
|
|
def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch):
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer
|
|
|
|
model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer
|
|
config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config
|
|
model = model_cls(
|
|
SimpleNamespace(
|
|
tf_model_config=config(num_hidden_layers=1),
|
|
dtype=torch.float32,
|
|
parallel_config=SimpleNamespace(use_hsdp=True),
|
|
hsdp_weight_load_strategy="pre_sharded",
|
|
)
|
|
)
|
|
groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)]
|
|
plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups)
|
|
processed = []
|
|
# Exercise model selection with real vLLM linears, without invoking the
|
|
# CPU-specific weight packing hook on meta tensors in CPU-only CI.
|
|
monkeypatch.setattr(
|
|
UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer)
|
|
)
|
|
DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
|
|
expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")]
|
|
assert expected
|
|
assert processed == expected
|
|
|
|
model._hsdp_pre_sharded_meta_post_load = False
|
|
with pytest.raises(ValueError, match="has not declared"):
|
|
DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
|
|
model._hsdp_pre_sharded_meta_post_load = True
|
|
|
|
class UnsupportedMethod(UnquantizedLinearMethod):
|
|
pass
|
|
|
|
expected[-1].quant_method = UnsupportedMethod()
|
|
processed.clear()
|
|
with pytest.raises(ValueError, match="UnsupportedMethod"):
|
|
DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan)
|
|
assert not processed
|
|
|
|
|
|
def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
|
|
class ToyReasonerLayer(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(1))
|
|
|
|
def forward(self, hidden: torch.Tensor, freqs):
|
|
del freqs
|
|
kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1)
|
|
return hidden + self.weight.to(hidden.dtype), kv, kv
|
|
|
|
class ToyGeneratorLayer(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(1))
|
|
|
|
def forward(self, hidden: torch.Tensor, **kwargs):
|
|
del kwargs
|
|
return hidden + self.weight.to(hidden.dtype)
|
|
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
model.language_model.layers = nn.ModuleList([ToyReasonerLayer()])
|
|
model.gen_layers = nn.ModuleList([ToyGeneratorLayer()])
|
|
model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False)
|
|
|
|
assert model._model_cpu_offload_enabled
|
|
assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"}
|
|
assert model.device == torch.device("cpu")
|
|
|
|
output = model(
|
|
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
)
|
|
|
|
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
|
# The reasoner runs once then the generator component stays resident for GEN.
|
|
assert model._active_model_cpu_offload_component == "generator"
|
|
|
|
model.disable_model_cpu_offload()
|
|
assert not model._model_cpu_offload_enabled
|
|
|
|
|
|
def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
|
|
monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device
|
|
) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)])
|
|
model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
|
|
model.to(accelerator_device)
|
|
|
|
reasoner_param = model.language_model.layers[0].weight
|
|
generator_param = model.gen_layers[0].weight
|
|
|
|
model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False)
|
|
|
|
assert model._model_cpu_offload_enabled
|
|
# On enable, every group is parked on CPU until a phase activates it.
|
|
assert reasoner_param.device.type == "cpu"
|
|
assert generator_param.device.type == "cpu"
|
|
|
|
model._activate_model_cpu_offload_component("reasoner")
|
|
assert reasoner_param.device == accelerator_device
|
|
assert generator_param.device.type == "cpu"
|
|
|
|
model._activate_model_cpu_offload_component("generator")
|
|
assert reasoner_param.device.type == "cpu"
|
|
assert generator_param.device == accelerator_device
|
|
|
|
model.disable_model_cpu_offload()
|
|
assert not model._model_cpu_offload_enabled
|
|
# Disable restores both components onto the device.
|
|
assert reasoner_param.device == accelerator_device
|
|
assert generator_param.device == accelerator_device
|
|
|
|
|
|
def test_multi_control_attention_weights_target_outputs() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
|
|
|
|
class ControlValueAttention:
|
|
def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
|
|
del k
|
|
# Each per-control K/V layout is [text, control_i, target].
|
|
return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1)
|
|
|
|
attention = SimpleNamespace(multi_control_attn=ControlValueAttention())
|
|
q = torch.zeros(1, 3, 1, 1)
|
|
k = torch.zeros_like(q)
|
|
v = torch.tensor([[[[10.0]], [[20.0]], [[30.0]]]])
|
|
text_kv = torch.zeros(1, 1, 1, 1)
|
|
|
|
output = Cosmos3CrossAttention._forward_multi_control(
|
|
attention,
|
|
q,
|
|
k,
|
|
v,
|
|
text_kv,
|
|
text_kv,
|
|
control_token_sizes=(1, 1),
|
|
control_weights=(0.25, 0.75),
|
|
)
|
|
|
|
torch.testing.assert_close(output, torch.tensor([[[10.0], [20.0], [17.5]]]))
|
|
|
|
|
|
def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
hidden_states = torch.zeros(1, 2, 1, 2, 2)
|
|
output = model(
|
|
hidden_states=hidden_states,
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
|
|
control_weights=[1.0, 3.0],
|
|
)
|
|
|
|
assert tuple(output.shape) == tuple(hidden_states.shape)
|
|
with pytest.raises(ValueError, match="control_weights length"):
|
|
model(
|
|
hidden_states=hidden_states,
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
|
|
control_weights=[1.0],
|
|
)
|
|
with pytest.raises(ValueError, match="control latent shape"):
|
|
model(
|
|
hidden_states=hidden_states,
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
control_latents=[torch.zeros(1, 2, 2, 2, 2)],
|
|
)
|
|
|
|
|
|
def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
class RecordingGather(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.calls = 0
|
|
|
|
def forward(self, hidden_gen: torch.Tensor) -> torch.Tensor:
|
|
self.calls += 1
|
|
return hidden_gen
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
gather = RecordingGather()
|
|
model.gen_sp_gather = gather
|
|
|
|
output = model(
|
|
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
)
|
|
|
|
assert gather.calls == 1
|
|
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
|
|
|
|
|
def test_sound_and_action_modules_follow_config() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
|
|
tiny = _tiny_cosmos3_config()
|
|
no_modal = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=tiny, dtype=torch.float32))
|
|
with_sound = Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=tiny, dtype=torch.float32),
|
|
sound_gen=True,
|
|
sound_dim=5,
|
|
sound_latent_fps=40.0,
|
|
)
|
|
with_action = Cosmos3VFMTransformer(
|
|
SimpleNamespace(
|
|
tf_model_config={**tiny, "action_gen": True, "max_action_dim": 6, "num_embodiment_domains": 9},
|
|
dtype=torch.float32,
|
|
)
|
|
)
|
|
|
|
assert no_modal.sound_gen is False
|
|
assert no_modal.action_gen is False
|
|
assert not hasattr(no_modal, "audio_proj_in")
|
|
assert not hasattr(no_modal, "action_proj_in")
|
|
assert with_sound.sound_dim == 5
|
|
assert with_sound.sound_latent_fps == 40.0
|
|
assert with_sound.audio_proj_in.in_features == 5
|
|
assert with_action.action_dim == 6
|
|
assert with_action.action_proj_in.num_domains == 9
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"kwargs",
|
|
[
|
|
{"sound_gen": True},
|
|
{"sound_gen": True, "sound_dim": 5},
|
|
{"sound_gen": True, "sound_latent_fps": 40.0},
|
|
],
|
|
)
|
|
def test_transformer_requires_sound_dim_and_fps_when_sound_gen_true(kwargs: dict[str, Any]) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
|
|
with pytest.raises(ValueError, match=r"requires an explicit sound_dim and sound_latent_fps"):
|
|
Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32),
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def test_sound_and_action_pack_unpack_validate_shapes() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
|
|
model = object.__new__(Cosmos3VFMTransformer)
|
|
nn.Module.__init__(model)
|
|
model.sound_dim = 3
|
|
model.action_dim = 3
|
|
|
|
sound = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4)
|
|
action = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
|
|
torch.testing.assert_close(model.unpack_sound(model.pack_sound(sound)), sound)
|
|
torch.testing.assert_close(model.unpack_action(model.pack_action(action)), action)
|
|
|
|
with pytest.raises(ValueError, match="channel mismatch"):
|
|
model.pack_sound(torch.zeros(1, 4, 2))
|
|
with pytest.raises(ValueError, match="dimension mismatch"):
|
|
model.pack_action(torch.zeros(1, 2, 4))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("config", "transformer_kwargs", "extra_kwargs", "expected_shapes"),
|
|
[
|
|
(
|
|
_tiny_cosmos3_config(),
|
|
{"sound_gen": True, "sound_dim": 3, "sound_latent_fps": 24.0},
|
|
{"sound_latents": torch.zeros(1, 3, 4)},
|
|
[(1, 2, 1, 2, 2), (1, 3, 4)],
|
|
),
|
|
(
|
|
_tiny_cosmos3_config(action_gen=True, max_action_dim=3, num_embodiment_domains=4),
|
|
{},
|
|
{"action_latents": torch.zeros(1, 5, 3), "action_domain_ids": torch.tensor([2])},
|
|
[(1, 2, 1, 2, 2), (1, 5, 3)],
|
|
),
|
|
],
|
|
)
|
|
def test_forward_returns_video_plus_optional_modality_predictions(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
config,
|
|
transformer_kwargs,
|
|
extra_kwargs,
|
|
expected_shapes,
|
|
) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
|
|
output = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=config, dtype=torch.float32),
|
|
**transformer_kwargs,
|
|
)(
|
|
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 2, 2),
|
|
fps=24.0,
|
|
action_noisy_mask=torch.ones(1, 5, 1),
|
|
**extra_kwargs,
|
|
)
|
|
|
|
assert isinstance(output, tuple)
|
|
assert [tuple(tensor.shape) for tensor in output] == expected_shapes
|
|
|
|
|
|
def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module
|
|
|
|
model = cosmos3_module.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32),
|
|
sound_gen=True,
|
|
sound_dim=3,
|
|
sound_latent_fps=40.0,
|
|
)
|
|
monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None))
|
|
|
|
output = model(
|
|
hidden_states=torch.zeros(1, 2, 1, 1, 2),
|
|
timestep=torch.tensor([1.0]),
|
|
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
|
text_mask=torch.ones(1, 2, dtype=torch.long),
|
|
video_shape=(1, 1, 2),
|
|
fps=24.0,
|
|
sound_latents=torch.zeros(1, 3, 1),
|
|
)
|
|
|
|
assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 1, 2), (1, 3, 1)]
|
|
|
|
|
|
def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.distributed import parallel_state
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
model = object.__new__(transformer_cosmos3.Cosmos3VFMTransformer)
|
|
model.latent_patch_size = 2
|
|
vs = (3, 16, 16)
|
|
|
|
monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 1)
|
|
assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 97
|
|
|
|
monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 2)
|
|
assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 98
|
|
assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=98) == 98
|
|
|
|
# ulysses=4, with a transfer control folded into the vision base.
|
|
monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 4)
|
|
padded = model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97, num_vision_items=2)
|
|
assert (2 * 192 + padded) % 4 == 0
|
|
|
|
|
|
def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
|
|
|
class FakeRotary:
|
|
def __init__(self) -> None:
|
|
self.position_ids: list[torch.Tensor] = []
|
|
|
|
def __call__(self, x, position_ids):
|
|
del x
|
|
self.position_ids.append(position_ids.detach().cpu())
|
|
batch, seq = position_ids.shape[1], position_ids.shape[2]
|
|
return torch.zeros(batch, seq, 4), torch.ones(batch, seq, 4)
|
|
|
|
rotary = FakeRotary()
|
|
model = object.__new__(Cosmos3VFMTransformer)
|
|
nn.Module.__init__(model)
|
|
model.language_model = SimpleNamespace(rotary_emb=rotary)
|
|
model.temporal_modality_margin = 100
|
|
model.base_fps = 24.0
|
|
model.temporal_compression_factor = 4
|
|
model.temporal_compression_factor_sound = 1
|
|
model.sound_latent_fps = 25.0
|
|
model.enable_fps_modulation = False
|
|
|
|
freqs_und, freqs_gen = model._compute_rope_freqs(
|
|
text_mask=torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long),
|
|
t=2,
|
|
hp=1,
|
|
wp=1,
|
|
fps=24.0,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.float32,
|
|
)
|
|
text_pos, vision_pos = rotary.position_ids
|
|
assert text_pos[:, 0, :].tolist() == [[0, 1, 0], [0, 1, 0], [0, 1, 0]]
|
|
assert vision_pos[0, 0].tolist() == [102, 103]
|
|
assert freqs_und[0].shape == (2, 3, 1, 4)
|
|
assert freqs_gen[0].shape == (2, 2, 1, 4)
|
|
|
|
rotary.position_ids.clear()
|
|
model._compute_rope_freqs(
|
|
text_mask=torch.tensor([[1, 1]], dtype=torch.long),
|
|
t=2,
|
|
hp=1,
|
|
wp=1,
|
|
fps=24.0,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.float32,
|
|
t_action=2,
|
|
action_start_frame_offset=1,
|
|
t_sound=1,
|
|
)
|
|
|
|
_, gen_pos = rotary.position_ids
|
|
assert gen_pos.shape == (3, 1, 5)
|
|
assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102]
|
|
|
|
rotary.position_ids.clear()
|
|
model._compute_rope_freqs(
|
|
text_mask=torch.tensor([[1, 1]], dtype=torch.long),
|
|
t=2,
|
|
hp=1,
|
|
wp=1,
|
|
fps=24.0,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.float32,
|
|
num_vision_items=3,
|
|
share_vision_temporal_positions=True,
|
|
)
|
|
_, shared_gen_pos = rotary.position_ids
|
|
assert shared_gen_pos[0, 0].tolist() == [102, 103, 102, 103, 102, 103]
|
|
|
|
rotary.position_ids.clear()
|
|
model._compute_rope_freqs(
|
|
text_mask=torch.tensor([[1, 1]], dtype=torch.long),
|
|
t=2,
|
|
hp=1,
|
|
wp=1,
|
|
fps=24.0,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.float32,
|
|
num_vision_items=3,
|
|
share_vision_temporal_positions=False,
|
|
)
|
|
_, offset_gen_pos = rotary.position_ids
|
|
assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107]
|
|
|
|
|
|
def test_shard_gen_prep_rejects_inputs_already_in_execution_layout(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
|
|
|
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
|
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
|
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
|
)
|
|
prep = model._gen_preprocess(
|
|
torch.zeros(1, 2, 1, 2, 2),
|
|
torch.tensor([1.0]),
|
|
torch.tensor([[1, 2]], dtype=torch.long),
|
|
torch.ones(1, 2, dtype=torch.long),
|
|
(1, 2, 2),
|
|
)
|
|
|
|
sharded = model._shard_gen_prep(prep, defer_gather=True)
|
|
assert sharded.freqs_gen is not None
|
|
assert sharded.defer_gen_gather
|
|
with pytest.raises(RuntimeError, match="already in the execution layout"):
|
|
model._shard_gen_prep(sharded)
|
|
# A deferred gather without sharding would gather a full-layout tensor.
|
|
with pytest.raises(RuntimeError, match="_shard_gen_prep"):
|
|
model._run_gen_stack(prep._replace(defer_gen_gather=True))
|