diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index de14ee233..58fe4b10b 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -75,6 +75,248 @@ def _tiny_cosmos3_edge_config(**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, diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index fc5a6f7cc..8510e0da9 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -39,6 +39,10 @@ from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, Seque from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.component_config import ( + ComponentQuantizationConfig, + resolve_component_quant_config, +) from .mixed_precision import ( Cosmos3MixedPrecisionConfig, @@ -52,6 +56,53 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _pathway_quant_config( + components: dict[str, QuantizationConfig | None], + pathway: str, + fallback: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Leaf overlay, or a scoped router when nested keys exist under ``pathway``. + + Linear layers announce prefixes such as ``language_model.layers.0.mlp``. + Passing a single pathway leaf would drop more-specific keys the factory + still accepts (for example ``language_model.layers.0.mlp``). Nested keys + keep a ``ComponentQuantizationConfig`` so longest-prefix matching still + selects them; unmatched prefixes use the pathway root or ``fallback``. + """ + nested = {key: value for key, value in components.items() if key.startswith(f"{pathway}.")} + has_root = pathway in components + root = components[pathway] if has_root else fallback + if not nested: + return root + scoped = dict(nested) + if has_root: + scoped[pathway] = root + return ComponentQuantizationConfig(scoped, default_config=root) + + +def _resolve_cosmos3_quant_configs( + quant_config: QuantizationConfig | None, +) -> tuple[QuantizationConfig | None, QuantizationConfig | None]: + """Resolve the Cosmos3 reasoner and generator quantization configs. + + A pipeline-level ``transformer`` entry is the default for both internal + pathways (same leaf-unwrapping pattern as Flux2 / MiniMax / Boogu). The + historical ``language_model`` and ``gen_layers`` scopes remain supported as + exact-key overlays, including explicit ``None`` entries that leave one + pathway unquantized. Nested keys under those roots keep longest-prefix + routing via a pathway-scoped ``ComponentQuantizationConfig``. + """ + if not isinstance(quant_config, ComponentQuantizationConfig): + return quant_config, quant_config + + transformer_config = resolve_component_quant_config(quant_config, "transformer") + components = quant_config.component_configs + return ( + _pathway_quant_config(components, "language_model", transformer_config), + _pathway_quant_config(components, "gen_layers", transformer_config), + ) + + class RMSNorm(_VllmRMSNorm): """Cosmos3-local RMSNorm that uses the FP32 native implementation.""" @@ -1288,7 +1339,9 @@ class Cosmos3VFMTransformer(nn.Module): self.use_und_k_norm_for_gen = _tf_config_get(model_config, "use_und_k_norm_for_gen", None) dtype = od_config.dtype - quant_config = getattr(od_config, "quantization_config", None) if od_config else None + language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs( + getattr(od_config, "quantization_config", None) + ) mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config) if mixed_precision_config is None: if mixed_precision_source == "additional_config_disabled": @@ -1314,7 +1367,7 @@ class Cosmos3VFMTransformer(nn.Module): rms_norm_eps=self.rms_norm_eps, rope_theta=self.rope_theta, mrope_section=self.mrope_section, - quant_config=quant_config, + quant_config=language_model_quant_config, prefix="language_model", **self._language_model_kwargs(), ) @@ -1352,7 +1405,7 @@ class Cosmos3VFMTransformer(nn.Module): num_key_value_heads=self.num_key_value_heads, head_dim=self.head_dim, rms_norm_eps=self.rms_norm_eps, - quant_config=quant_config, + quant_config=gen_layers_quant_config, mlp_cls=self._gen_mlp_cls, qk_norm=self.qk_norm_for_diffusion, prefix=f"gen_layers.{i}",