diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md index d8da4e1a9..07d05730e 100644 --- a/docs/user_guide/quantization/modelopt.md +++ b/docs/user_guide/quantization/modelopt.md @@ -204,6 +204,91 @@ omni = Omni( | `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages | | `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints | +## Cosmos3 Mixed-Precision Schedule + +The experimental Cosmos3 schedule can use the checkpoint-native quantized GEMM for middle denoising +steps and dense 16-bit activation execution for selected first and last steps. +The schedule supports serialized ModelOpt FP8 and NVFP4 checkpoints: + +| Detected weight format | Native middle steps | A16 boundary steps | +|------------------------|---------------------|--------------------| +| FP8 | W8A8 | W8A16 | +| NVFP4 | W4A4 | W4A16 | + +The native and A16 paths share one quantized weight representation. The A16 +path dequantizes the live backend weight and calls `F.linear`; it does not keep +a second checkpoint-weight snapshot. Scheduled FP8 currently requires +serialized tensorwise scales and a backend that retains canonical FP8 weights. +Scheduled NVFP4 requires a supported CUTLASS-compatible native or FlashInfer +layout. SmoothQuant, Marlin-repacked FP8, and per-channel/per-token FP8 +schedules fail closed. + +The reasoner uses dense A16 execution by default when its weights are FP8 or +NVFP4. Set the nested `reasoner` field to `native` to retain the checkpoint-native +W8A8/W4A4 path. A BF16 reasoner remains BF16 in either mode. The schedule +currently requires tensor parallel size 1 and one active request per worker, +does not support HSDP, and does not support block-scaled FP8. Its live-weight +design is compatible with model-level and standard layer-wise offload, but +those paths still require GPU end-to-end validation. Distributed layer-wise +offload is rejected because its direct loader bypasses ModelOpt post-load +transformations. + +A checkpoint can opt in by adding a versioned policy to its authoritative +`transformer/config.json`. The same policy schema is used for FP8 and NVFP4: + +```json +{ + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": { + "diffusion_step_policy": { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": ["transformer"], + "default_mode": "native", + "first_steps": {"count": 3, "mode": "a16"}, + "last_steps": {"count": 3, "mode": "a16"}, + "overlap": "a16", + "reasoner": "a16" + } + } + } +} +``` + +The example preserves the checkpoint's existing `quant_method` and +`quant_algo`; only `runtime` is added. + +When present and valid, the checkpoint policy is the default. An explicit +runtime object overrides it. To retain ordinary checkpoint-native execution +for every step, disable only the schedule: + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config '{"cosmos3_mixed_precision":{"enabled":false}}' +``` + +This does not disable checkpoint quantization; it prevents installation of the +mixed-step runtime. + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config \ + '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}' +``` + +The presence of `cosmos3_mixed_precision` supplies an explicit runtime +override; an empty object uses the defaults shown above. +The schedule uses vLLM's selected native backend and fails during loading if +its live weight layout cannot support A16 dequantization. +The runtime selects FP8 or NVFP4 from the checkpoint's ModelOpt configuration. +BF16 linears are untouched. Mixed FP8/NVFP4 checkpoints are not supported by +this schedule. + ## Validation and Notes 1. Compare the ModelOpt checkpoint against the BF16 baseline with the same diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 490f5b1f4..f81e5e887 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -115,6 +115,38 @@ shift at the same seed). The pipeline auto-resolves from `model_index.json`; pass `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly. +For a serialized ModelOpt FP8 or NVFP4 checkpoint, an experimental +mixed-precision schedule can use native W8A8/W4A4 in middle denoising steps +and dense W8A16/W4A16 in the first and last steps: + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config \ + '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}' +``` + +The nested object's presence supplies an explicit runtime override; use an +empty object for these defaults. A compatible checkpoint can instead declare +the versioned `runtime.diffusion_step_policy` in its authoritative +`transformer/config.json`. The checkpoint policy is used only when no runtime +override is present. Use +`--additional-config '{"cosmos3_mixed_precision":{"enabled":false}}'` to +disable the checkpoint schedule without disabling checkpoint quantization. +This path currently requires tensor parallel size 1 and does not support HSDP +or block-scaled FP8. It keeps one live quantized weight representation: +scheduled FP8 requires serialized tensorwise scales and a canonical backend, +while scheduled NVFP4 requires a supported CUTLASS-compatible native or +FlashInfer layout. Its dequantization is a correctness baseline; no speedup is +implied. Quantized reasoner weights use A16 by default; set +`"reasoner":"native"` in the nested object to keep W8A8/W4A4. +The checkpoint's ModelOpt configuration selects FP8 or NVFP4; mixed FP8/NVFP4 +checkpoints are not supported by this schedule. BF16 linears are left unchanged. +The schedule currently supports one active request per worker. Model-level and +layer-wise offload are designed to work with this live-weight path but still +require GPU validation. Distributed layer-wise offload and incompatible native +weight layouts fail during loading. + #### Verification Best quality uses the JSON-upsampled prompts from `assets/` (download with diff --git a/tests/diffusion/quantization/test_cosmos3_mixed_precision.py b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py new file mode 100644 index 000000000..026090fe9 --- /dev/null +++ b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.cosmos3.mixed_precision import ( + Cosmos3MixedPrecisionConfig, + Cosmos3MixedPrecisionRuntime, + resolve_mixed_precision_config, +) +from vllm_omni.diffusion.models.cosmos3.mixed_precision import runtime as runtime_impl +from vllm_omni.diffusion.models.cosmos3.mixed_precision.checkpoint import read_checkpoint_policy +from vllm_omni.diffusion.models.cosmos3.mixed_precision.runtime import ( + Cosmos3MixedPrecisionLinearMethod, +) +from vllm_omni.diffusion.models.cosmos3.mixed_precision.strategy import ( + Fp8W8A8W8A16Strategy, + Nvfp4W4A4W4A16Strategy, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _policy( + *, + first_steps: int = 3, + last_steps: int = 3, + reasoner: str = "a16", + scope: list[str] | None = None, +) -> dict[str, object]: + return { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": scope or ["transformer"], + "default_mode": "native", + "first_steps": {"count": first_steps, "mode": "a16"}, + "last_steps": {"count": last_steps, "mode": "a16"}, + "overlap": "a16", + "reasoner": reasoner, + } + + +def _fake_quant_config(name: str, *, serialized: bool = True): + if name == "modelopt": + return SimpleNamespace( + get_name=lambda: name, + is_checkpoint_fp8_serialized=serialized, + ) + if name == "modelopt_fp4": + return SimpleNamespace( + get_name=lambda: name, + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=serialized, + ) + return SimpleNamespace(get_name=lambda: name) + + +def _checkpoint_od_config( + quant_config, + *, + policy: object | None = None, + additional_config: dict[str, object] | None = None, +): + quant_config_name = quant_config.get_name() + quant_algo = { + "modelopt": "FP8", + "modelopt_fp4": "NVFP4", + }.get(quant_config_name, "FP8") + disk_quant_config: dict[str, object] = { + "quant_method": quant_config_name, + "quant_algo": quant_algo, + } + if policy is not None: + disk_quant_config["runtime"] = {"diffusion_step_policy": policy} + return SimpleNamespace( + quantization_config=quant_config, + tf_model_config=SimpleNamespace( + params={"quantization_config": disk_quant_config}, + quant_config=quant_config, + ), + additional_config=additional_config, + ) + + +def _swizzle_blockscale_cpu(scale: torch.Tensor) -> torch.Tensor: + rows, cols = scale.shape + padded_rows = ((rows + 127) // 128) * 128 + padded_cols = ((cols + 3) // 4) * 4 + padded = torch.zeros( + (1, padded_rows, padded_cols), + dtype=scale.dtype, + device=scale.device, + ) + padded[0, :rows, :cols] = scale + padded = padded.reshape(1, padded_rows // 128, 4, 32, padded_cols // 4, 4) + return ( + padded.permute(0, 1, 4, 3, 2, 5) + .contiguous() + .reshape( + padded_rows, + padded_cols, + ) + ) + + +def test_config_parses_nested_schedule_and_reasoner_policy() -> None: + config = Cosmos3MixedPrecisionConfig.from_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 2, + "last_steps": 4, + "reasoner": "native", + } + } + ) + + assert config is not None + assert config.reasoner == "native" + assert [index for index in range(10) if config.use_high_precision(index, 10)] == [0, 1, 6, 7, 8, 9] + + +@pytest.mark.parametrize( + ("first_steps", "last_steps", "selected"), + [ + (0, 2, [5, 6]), + (2, 0, [0, 1]), + (0, 0, []), + (4, 4, list(range(7))), + ], +) +def test_schedule_boundaries_and_overlap( + first_steps: int, + last_steps: int, + selected: list[int], +) -> None: + config = Cosmos3MixedPrecisionConfig( + first_steps=first_steps, + last_steps=last_steps, + ) + assert [index for index in range(7) if config.use_high_precision(index, 7)] == selected + + +def test_one_step_request_honors_boundary_precision() -> None: + config = Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1) + assert config.use_high_precision(0, 1) + + +@pytest.mark.parametrize( + ("values", "message"), + [ + ({"cosmos3_mixed_precision": True}, "must be a mapping"), + ({"cosmos3_mixed_precision": {"first_steps": -1}}, "non-negative"), + ({"cosmos3_mixed_precision": {"reasoner": "fp16"}}, "must be one of"), + ({"cosmos3_mixed_precision": {"enabled": 0}}, "must be a boolean"), + ( + {"cosmos3_mixed_precision": {"enabled": False, "first_steps": 1}}, + "cannot be combined", + ), + ({"cosmos3_mixed_precision": {"unknown": 1}}, "Unknown"), + ], +) +def test_config_rejects_invalid_values(values: dict, message: str) -> None: + with pytest.raises((TypeError, ValueError), match=message): + Cosmos3MixedPrecisionConfig.from_additional_config(values) + + +def test_config_presence_enables_defaults() -> None: + assert Cosmos3MixedPrecisionConfig.from_additional_config({}) is None + config = Cosmos3MixedPrecisionConfig.from_additional_config({"cosmos3_mixed_precision": {}}) + assert config == Cosmos3MixedPrecisionConfig() + + +def test_noop_config_does_not_install_runtime() -> None: + assert ( + Cosmos3MixedPrecisionConfig.from_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 0, + "last_steps": 0, + "reasoner": "native", + } + } + ) + is None + ) + + +def test_additional_config_distinguishes_absent_and_disabled() -> None: + assert Cosmos3MixedPrecisionConfig.resolve_additional_config({}) == (False, None) + assert Cosmos3MixedPrecisionConfig.resolve_additional_config({"cosmos3_mixed_precision": {"enabled": False}}) == ( + True, + None, + ) + assert Cosmos3MixedPrecisionConfig.resolve_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 0, + "last_steps": 0, + "reasoner": "native", + } + } + ) == (True, None) + + +def test_checkpoint_policy_round_trip_through_transformer_config() -> None: + from vllm_omni.diffusion.data import TransformerConfig + + disk = { + "_class_name": "Cosmos3VFMTransformer", + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": _policy(first_steps=2, last_steps=4)}, + }, + } + + tf_config = TransformerConfig.from_dict(disk) + assert tf_config.quant_config is not None + assert tf_config.quant_config.get_name() == "modelopt" + assert tf_config.to_dict()["quantization_config"]["runtime"] == disk["quantization_config"]["runtime"] + assert read_checkpoint_policy(SimpleNamespace(tf_model_config=tf_config)) == Cosmos3MixedPrecisionConfig( + first_steps=2, + last_steps=4, + reasoner="a16", + ) + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_checkpoint_policy_supports_serialized_modelopt_formats(name: str) -> None: + od_config = _checkpoint_od_config(_fake_quant_config(name), policy=_policy()) + assert read_checkpoint_policy(od_config) == Cosmos3MixedPrecisionConfig() + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_checkpoint_policy_rejects_nonserialized_modelopt_formats(name: str) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name, serialized=False), + policy=_policy(), + ) + with pytest.raises(ValueError, match="serialized ModelOpt"): + read_checkpoint_policy(od_config) + + +def test_checkpoint_policy_rejects_mixed_modelopt_format() -> None: + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt_mixed"), policy=_policy())) + + +def test_checkpoint_policy_missing_metadata_preserves_ordinary_path() -> None: + assert read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"))) is None + + +def test_checkpoint_policy_rejects_malformed_runtime_container() -> None: + od_config = _checkpoint_od_config(_fake_quant_config("modelopt")) + od_config.tf_model_config.params["quantization_config"]["runtime"] = True + with pytest.raises(TypeError, match="runtime must be a mapping"): + read_checkpoint_policy(od_config) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("schema_version", 2, "schema_version"), + ("type", "sigma_ranges", "type"), + ("index_space", "scheduler_timestep", "index_space"), + ("default_mode", "a16", "default_mode"), + ("overlap", "native", "overlap"), + ("reasoner", "fp16", "reasoner"), + ], +) +def test_checkpoint_policy_rejects_invalid_schema(field: str, value: object, message: str) -> None: + policy = _policy() + policy[field] = value + with pytest.raises(ValueError, match=message): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy)) + + +@pytest.mark.parametrize( + ("step_range", "message"), + [ + ({"count": -1, "mode": "a16"}, "non-negative"), + ({"count": True, "mode": "a16"}, "non-negative"), + ({"count": 1, "mode": "native"}, "mode"), + ({"mode": "a16"}, "Missing"), + ({"count": 1, "mode": "a16", "sigma": 1.0}, "Unknown"), + ], +) +def test_checkpoint_policy_rejects_invalid_step_range(step_range: object, message: str) -> None: + policy = _policy() + policy["first_steps"] = step_range + with pytest.raises((TypeError, ValueError), match=message): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy)) + + +def test_checkpoint_policy_rejects_unsupported_quantization_config() -> None: + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("compressed-tensors"), policy=_policy())) + + +def test_checkpoint_policy_for_other_scope_is_not_applied() -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(scope=["vae"]), + ) + assert read_checkpoint_policy(od_config) is None + + +def test_additional_config_overrides_and_disables_checkpoint_policy() -> None: + checkpoint = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(first_steps=1, last_steps=1), + ) + assert resolve_mixed_precision_config(checkpoint) == ( + Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1), + "checkpoint", + ) + + checkpoint.additional_config = { + "cosmos3_mixed_precision": { + "first_steps": 2, + "last_steps": 0, + "reasoner": "native", + } + } + assert resolve_mixed_precision_config(checkpoint) == ( + Cosmos3MixedPrecisionConfig(first_steps=2, last_steps=0, reasoner="native"), + "additional_config", + ) + + checkpoint.additional_config = {"cosmos3_mixed_precision": {"enabled": False}} + assert resolve_mixed_precision_config(checkpoint) == (None, "additional_config_disabled") + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_additional_config_supports_serialized_modelopt_formats_without_checkpoint_policy(name: str) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name), + additional_config={"cosmos3_mixed_precision": {}}, + ) + assert resolve_mixed_precision_config(od_config) == (Cosmos3MixedPrecisionConfig(), "additional_config") + + +@pytest.mark.parametrize("has_checkpoint_policy", [False, True]) +@pytest.mark.parametrize( + ("name", "serialized"), + [ + ("modelopt_mixed", True), + ("mxfp8", True), + ("compressed-tensors", True), + ("modelopt", False), + ("modelopt_fp4", False), + ], +) +def test_additional_config_rejects_incompatible_quantization( + name: str, serialized: bool, has_checkpoint_policy: bool +) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name, serialized=serialized), + policy=_policy() if has_checkpoint_policy else None, + additional_config={"cosmos3_mixed_precision": {}}, + ) + with pytest.raises(ValueError, match="serialized ModelOpt"): + resolve_mixed_precision_config(od_config) + + +@pytest.mark.parametrize("additional_config", [None, {"cosmos3_mixed_precision": {}}]) +def test_policy_validates_effective_quantization_config(additional_config: dict[str, object] | None) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(), + additional_config=additional_config, + ) + od_config.quantization_config = _fake_quant_config("modelopt_mixed") + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + resolve_mixed_precision_config(od_config) + + +def test_additional_config_rejects_missing_quantization_config() -> None: + od_config = SimpleNamespace(additional_config={"cosmos3_mixed_precision": {}}) + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + resolve_mixed_precision_config(od_config) + + +@pytest.mark.parametrize("policy", [_policy(), True]) +def test_additional_config_disable_bypasses_incompatible_checkpoint(policy: object) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt_mixed"), + policy=policy, + additional_config={"cosmos3_mixed_precision": {"enabled": False}}, + ) + od_config.quantization_config = _fake_quant_config("modelopt_mixed") + assert resolve_mixed_precision_config(od_config) == (None, "additional_config_disabled") + + +def test_runtime_allows_standard_offload_and_rejects_distributed(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr( + transformer_cosmos3, + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=True, + enable_distributed_layerwise_offload=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + od_config.enable_distributed_layerwise_offload = True + with pytest.raises(ValueError, match="distributed layer-wise offload"): + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + od_config.enable_distributed_layerwise_offload = False + od_config.max_num_seqs = 2 + with pytest.raises(ValueError, match="one active request"): + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + + +class _Layer(torch.nn.Module): + def __init__(self, output_size: int, input_size: int) -> None: + super().__init__() + self.output_size_per_partition = output_size + self.input_size_per_partition = input_size + + +class _BaseMethod: + def __init__(self, *, nvfp4: bool = False) -> None: + self.processed = False + self.apply_calls = 0 + self.fp8_linear = object() + if nvfp4: + from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, + ) + + self.kernel = object.__new__(CutlassNvFp4LinearKernel) + else: + self.kernel = None + + def create_weights(self, *args, **kwargs) -> None: + pass + + def process_weights_after_loading(self, layer) -> None: + self.processed = True + if layer.weight.dtype == torch.float8_e4m3fn: + layer.weight = layer.weight.t().contiguous() + layer.weight_scale = layer.weight_scale.max().reshape(1) + elif layer.weight.dtype == torch.uint8: + layer.weight_scale = _swizzle_blockscale_cpu(layer.weight_scale) + layer.weight_global_scale = layer.weight_scale_2.max().float().reshape(1) + del layer.weight_scale_2 + + def apply(self, layer, x, bias=None): + del bias + self.apply_calls += 1 + return torch.full( + (*x.shape[:-1], layer.output_size_per_partition), + 17, + dtype=x.dtype, + device=x.device, + ) + + +def _fp8_layer(output_size: int = 2, input_size: int = 4) -> _Layer: + layer = _Layer(output_size, input_size) + layer.weight = torch.ones(output_size, input_size, dtype=torch.float8_e4m3fn) + layer.weight_scale = torch.tensor([0.5], dtype=torch.float32) + return layer + + +def _nvfp4_layer(output_size: int = 128, input_size: int = 64) -> _Layer: + layer = _Layer(output_size, input_size) + layer.weight = torch.full((output_size, input_size // 2), 0x22, dtype=torch.uint8) + layer.weight_scale = torch.ones(output_size, input_size // 16, dtype=torch.float8_e4m3fn) + layer.weight_scale_2 = torch.ones(1) + return layer + + +def _runtime_and_method( + strategy, + layer: _Layer, + *, + path: str = "generation", + reasoner: str = "native", +): + config = Cosmos3MixedPrecisionConfig( + first_steps=1, + last_steps=1, + reasoner=reasoner, + ) + runtime = Cosmos3MixedPrecisionRuntime(config) + base = _BaseMethod(nvfp4=isinstance(strategy, Nvfp4W4A4W4A16Strategy)) + method = Cosmos3MixedPrecisionLinearMethod( + base, + strategy, + runtime, + f"{path}.linear", + path, + ) + method.process_weights_after_loading(layer) + return runtime, base, method + + +@pytest.mark.parametrize( + ("strategy", "layer"), + [ + (Fp8W8A8W8A16Strategy(), _fp8_layer()), + (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()), + ], +) +def test_generation_dispatches_native_middle_and_a16_edges(strategy, layer) -> None: + runtime, base, method = _runtime_and_method(strategy, layer) + x = torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16) + + runtime.set_step(1, 3) + assert torch.equal( + method.apply(layer, x), + torch.full((1, layer.output_size_per_partition), 17, dtype=x.dtype), + ) + runtime.set_step(0, 3) + assert torch.equal( + method.apply(layer, x), + torch.nn.functional.linear(x, strategy.materialize(layer)), + ) + assert base.apply_calls == 1 + + +@pytest.mark.parametrize( + ("strategy", "layer"), + [ + (Fp8W8A8W8A16Strategy(), _fp8_layer()), + (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()), + ], +) +def test_a16_uses_only_live_native_weights(strategy, layer) -> None: + runtime, base, method = _runtime_and_method(strategy, layer) + assert base.processed + assert not hasattr(layer, "_cosmos3_precision_weight") + assert not hasattr(layer, "_cosmos3_precision_weight_scale") + + runtime.set_step(0, 3) + output = method.apply( + layer, + torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16), + ) + assert output.shape == (1, layer.output_size_per_partition) + + +def test_a16_follows_live_weight_storage_after_offload_rebinding() -> None: + layer = _fp8_layer() + strategy = Fp8W8A8W8A16Strategy() + runtime, _, method = _runtime_and_method(strategy, layer) + runtime.set_step(0, 3) + x = torch.ones(1, 4, dtype=torch.bfloat16) + before = method.apply(layer, x) + + layer.weight = torch.zeros_like(layer.weight) + after = method.apply(layer, x) + assert not torch.equal(before, after) + assert torch.count_nonzero(after) == 0 + + +def test_fp8_rejects_smoothquant() -> None: + layer = _fp8_layer() + layer.pre_quant_scale = torch.ones(1) + with pytest.raises(ValueError, match="SmoothQuant"): + Fp8W8A8W8A16Strategy().validate_before_processing( + _BaseMethod(), + layer, + "gen.linear", + ) + + +def test_fp8_rejects_non_tensorwise_scales() -> None: + layer = _fp8_layer(output_size=3, input_size=2) + layer.weight_scale = torch.tensor([[1.0], [2.0], [3.0]]) + with pytest.raises(ValueError, match="tensorwise FP8"): + Fp8W8A8W8A16Strategy().validate_before_processing( + _BaseMethod(), + layer, + "gen.linear", + ) + + +def test_nvfp4_reference_materializes_live_cutlass_weights() -> None: + layer = _nvfp4_layer() + strategy = Nvfp4W4A4W4A16Strategy() + _runtime_and_method(strategy, layer) + output = strategy.materialize(layer) + assert output.shape == (128, 64) + assert torch.equal(output, torch.ones_like(output)) + + +def test_nvfp4_rejects_fused_global_scales() -> None: + layer = _nvfp4_layer() + layer.weight_scale_2 = torch.tensor([1.0, 2.0]) + strategy = Nvfp4W4A4W4A16Strategy() + with pytest.raises(ValueError, match="NVFP4 scales"): + strategy.validate_before_processing( + _BaseMethod(nvfp4=True), + layer, + "gen.qkv", + ) + + +@pytest.mark.parametrize( + ("policy", "uses_native"), + [("a16", False), ("native", True)], +) +def test_reasoner_policy_is_independent_of_generation_step( + policy: str, + uses_native: bool, +) -> None: + layer = _fp8_layer() + runtime, base, method = _runtime_and_method( + Fp8W8A8W8A16Strategy(), + layer, + path="reasoner", + reasoner=policy, + ) + runtime.set_step(1, 3) + method.apply(layer, torch.ones(1, 4, dtype=torch.bfloat16)) + assert (base.apply_calls == 1) is uses_native + + +def test_reset_clears_generation_state() -> None: + runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1)) + runtime.set_step(0, 5) + assert runtime.use_high_precision("generation") + runtime.reset() + assert not runtime.use_high_precision("generation") + + +def test_install_discovers_generation_and_opt_in_reasoner(monkeypatch) -> None: + class _FakeLinear(torch.nn.Module): + def __init__(self, prefix: str) -> None: + super().__init__() + self.prefix = prefix + self.input_size_per_partition = 4 + self.output_size_per_partition = 2 + self.quant_method = SimpleNamespace(name="fp8") + + monkeypatch.setattr(runtime_impl, "LinearBase", _FakeLinear) + strategy = Fp8W8A8W8A16Strategy() + monkeypatch.setattr(strategy, "accepts", lambda method: getattr(method, "name", None) == "fp8") + monkeypatch.setattr(runtime_impl, "_STRATEGIES", (strategy,)) + runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(reasoner="a16")) + reasoner = _FakeLinear("reasoner.q_proj") + generation = [ + _FakeLinear("generation.q_proj"), + _FakeLinear("generation.out_proj"), + ] + transformer = SimpleNamespace( + language_model=SimpleNamespace(layers=torch.nn.Sequential(reasoner)), + gen_layers=torch.nn.Sequential(*generation), + ) + runtime.install(transformer) + assert isinstance(reasoner.quant_method, Cosmos3MixedPrecisionLinearMethod) + assert all(isinstance(layer.quant_method, Cosmos3MixedPrecisionLinearMethod) for layer in generation) + + +def test_pipeline_helpers_forward_and_reset_schedule() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + calls = [] + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + pipeline.transformer = SimpleNamespace( + set_mixed_precision_step=lambda step, count: calls.append((step, count)), + reset_mixed_precision=lambda: calls.append("reset"), + ) + pipeline._set_mixed_precision_step(2, 7) + pipeline._reset_mixed_precision() + assert calls == [(2, 7), "reset"] diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py new file mode 100644 index 000000000..1d5ee972b --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Scheduled A16 execution for quantized Cosmos3 checkpoints.""" + +from .checkpoint import resolve_mixed_precision_config +from .config import Cosmos3MixedPrecisionConfig +from .runtime import Cosmos3MixedPrecisionRuntime + +__all__ = [ + "Cosmos3MixedPrecisionConfig", + "Cosmos3MixedPrecisionRuntime", + "resolve_mixed_precision_config", +] diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py new file mode 100644 index 000000000..c8a62f142 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Checkpoint discovery for the Cosmos3 diffusion-step precision policy.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Literal, cast + +from .config import Cosmos3MixedPrecisionConfig, ReasonerPolicy, _non_negative_int + +PolicySource = Literal[ + "none", + "checkpoint", + "additional_config", + "additional_config_disabled", +] + +_COMPONENT = "transformer" +_POLICY_FIELDS = frozenset( + { + "schema_version", + "type", + "index_space", + "scope", + "default_mode", + "first_steps", + "last_steps", + "overlap", + "reasoner", + } +) +_STEP_RANGE_FIELDS = frozenset({"count", "mode"}) +_SUPPORTED_MODEL_OPT_CONFIGS = frozenset( + { + "modelopt", + "modelopt_fp4", + } +) + + +def resolve_mixed_precision_config( + od_config: object, +) -> tuple[Cosmos3MixedPrecisionConfig | None, PolicySource]: + """Resolve an explicit runtime override before the checkpoint default.""" + override_present, override = Cosmos3MixedPrecisionConfig.resolve_additional_config( + getattr(od_config, "additional_config", None) + ) + if override_present: + if override is not None: + _validate_checkpoint_quantization(od_config) + source: PolicySource = "additional_config" if override is not None else "additional_config_disabled" + return override, source + + checkpoint = read_checkpoint_policy(od_config) + if checkpoint is not None: + return checkpoint, "checkpoint" + return None, "none" + + +def read_checkpoint_policy(od_config: object) -> Cosmos3MixedPrecisionConfig | None: + """Read and validate a policy from ``transformer/config.json``. + + Missing metadata preserves ordinary checkpoint behavior. Metadata that is + present but malformed or incompatible fails closed. + """ + tf_config = getattr(od_config, "tf_model_config", None) + params = getattr(tf_config, "params", None) + if not isinstance(params, Mapping): + return None + + disk_quant_config = params.get("quantization_config") + if not isinstance(disk_quant_config, Mapping): + return None + + if "runtime" not in disk_quant_config: + return None + runtime_config = disk_quant_config["runtime"] + if not isinstance(runtime_config, Mapping): + raise TypeError("quantization_config.runtime must be a mapping") + + if "diffusion_step_policy" not in runtime_config: + return None + raw_policy = runtime_config["diffusion_step_policy"] + if not isinstance(raw_policy, Mapping): + raise TypeError("quantization_config.runtime.diffusion_step_policy must be a mapping") + + policy = _parse_policy(raw_policy) + if policy is None: + return None + _validate_checkpoint_quantization(od_config) + return policy + + +def _parse_policy(policy: Mapping[str, object]) -> Cosmos3MixedPrecisionConfig | None: + fields = set(policy) + unknown = fields - _POLICY_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy fields: {sorted(unknown)}") + missing = _POLICY_FIELDS - fields + if missing: + raise ValueError(f"Missing diffusion_step_policy fields: {sorted(missing)}") + + schema_version = policy["schema_version"] + if not isinstance(schema_version, int) or isinstance(schema_version, bool) or schema_version != 1: + raise ValueError("diffusion_step_policy.schema_version must be the integer 1") + if policy["type"] != "first_last_n": + raise ValueError("diffusion_step_policy.type must be 'first_last_n'") + if policy["index_space"] != "denoising_loop_iteration": + raise ValueError("diffusion_step_policy.index_space must be 'denoising_loop_iteration'") + if policy["default_mode"] != "native": + raise ValueError("diffusion_step_policy.default_mode must be 'native'") + if policy["overlap"] != "a16": + raise ValueError("diffusion_step_policy.overlap must be 'a16'") + + scope = policy["scope"] + if not isinstance(scope, list) or not scope or not all(isinstance(item, str) for item in scope): + raise TypeError("diffusion_step_policy.scope must be a non-empty list of strings") + + first_steps = _parse_step_range(policy["first_steps"], "first_steps") + last_steps = _parse_step_range(policy["last_steps"], "last_steps") + + reasoner = policy["reasoner"] + if reasoner not in {"native", "a16"}: + raise ValueError("diffusion_step_policy.reasoner must be 'native' or 'a16'") + + if _COMPONENT not in scope: + return None + return Cosmos3MixedPrecisionConfig( + first_steps=first_steps, + last_steps=last_steps, + reasoner=cast(ReasonerPolicy, reasoner), + ) + + +def _parse_step_range(value: object, name: str) -> int: + if not isinstance(value, Mapping): + raise TypeError(f"diffusion_step_policy.{name} must be a mapping") + fields = set(value) + unknown = fields - _STEP_RANGE_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy.{name} fields: {sorted(unknown)}") + missing = _STEP_RANGE_FIELDS - fields + if missing: + raise ValueError(f"Missing diffusion_step_policy.{name} fields: {sorted(missing)}") + if value["mode"] != "a16": + raise ValueError(f"diffusion_step_policy.{name}.mode must be 'a16'") + return _non_negative_int(value["count"], f"diffusion_step_policy.{name}.count") + + +def _validate_checkpoint_quantization(od_config: object) -> None: + # Match the effective quantization used to construct the transformer linears. + quant_config = getattr(od_config, "quantization_config", None) + if quant_config is None: + tf_config = getattr(od_config, "tf_model_config", None) + quant_config = getattr(tf_config, "quant_config", None) + + get_name = getattr(quant_config, "get_name", None) + name = get_name() if callable(get_name) else None + if name not in _SUPPORTED_MODEL_OPT_CONFIGS: + raise ValueError(f"diffusion_step_policy requires a serialized ModelOpt FP8 or NVFP4 checkpoint, got {name!r}") + + if name == "modelopt": + _require_serialized(quant_config, "is_checkpoint_fp8_serialized", "FP8") + return + if name == "modelopt_fp4": + if getattr(quant_config, "quant_method", "NVFP4") != "NVFP4": + raise ValueError("diffusion_step_policy requires checkpoint-native NVFP4 W4A4 linears") + _require_serialized(quant_config, "is_checkpoint_nvfp4_serialized", "NVFP4") + return + + +def _require_serialized(quant_config: object | None, flag: str, format_name: str) -> None: + if not bool(getattr(quant_config, flag, False)): + raise ValueError(f"diffusion_step_policy requires serialized ModelOpt {format_name} weights") diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py new file mode 100644 index 000000000..13748c8fd --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration parsing and denoising-step policy for mixed precision.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Literal, cast + +ReasonerPolicy = Literal["native", "a16"] + +_REASONER_POLICIES = frozenset({"native", "a16"}) +_CONFIG_FIELDS = frozenset({"enabled", "first_steps", "last_steps", "reasoner"}) + + +@dataclass(frozen=True) +class Cosmos3MixedPrecisionConfig: + """Validated precision policy shared by every quantization strategy.""" + + first_steps: int = 3 + last_steps: int = 3 + reasoner: ReasonerPolicy = "a16" + + @classmethod + def from_additional_config( + cls, + additional_config: Mapping[str, object] | None, + ) -> Cosmos3MixedPrecisionConfig | None: + """Parse Cosmos3 fields from vLLM-Omni's additional configuration.""" + _, config = cls.resolve_additional_config(additional_config) + return config + + @classmethod + def resolve_additional_config( + cls, + additional_config: Mapping[str, object] | None, + ) -> tuple[bool, Cosmos3MixedPrecisionConfig | None]: + """Return whether an override was supplied and its effective policy. + + The separate presence bit distinguishes an absent runtime override from + an explicit disable. That distinction matters when a checkpoint owns a + default policy. + """ + if additional_config is not None and not isinstance(additional_config, Mapping): + raise TypeError("additional_config must be a mapping") + values = additional_config or {} + if "cosmos3_mixed_precision" not in values: + return False, None + raw_config = values["cosmos3_mixed_precision"] + if not isinstance(raw_config, Mapping): + raise TypeError("cosmos3_mixed_precision must be a mapping") + unknown = set(raw_config) - _CONFIG_FIELDS + if unknown: + raise ValueError(f"Unknown cosmos3_mixed_precision fields: {sorted(unknown)}") + + enabled = raw_config.get("enabled", True) + if not isinstance(enabled, bool): + raise TypeError("cosmos3_mixed_precision.enabled must be a boolean") + if not enabled: + conflicting = set(raw_config) - {"enabled"} + if conflicting: + raise ValueError( + "cosmos3_mixed_precision.enabled=false cannot be combined with " + f"schedule fields: {sorted(conflicting)}" + ) + return True, None + + first_steps = _non_negative_int( + raw_config.get("first_steps", 3), + "cosmos3_mixed_precision.first_steps", + ) + last_steps = _non_negative_int( + raw_config.get("last_steps", 3), + "cosmos3_mixed_precision.last_steps", + ) + + reasoner = str(raw_config.get("reasoner", "a16")).lower() + if reasoner not in _REASONER_POLICIES: + raise ValueError( + f"cosmos3_mixed_precision.reasoner must be one of {sorted(_REASONER_POLICIES)}, got {reasoner!r}" + ) + config = cls( + first_steps=first_steps, + last_steps=last_steps, + reasoner=cast(ReasonerPolicy, reasoner), + ) + if first_steps == 0 and last_steps == 0 and reasoner == "native": + return True, None + return True, config + + def use_high_precision(self, step_index: int, num_steps: int) -> bool: + """Select the high path for an actual scheduler-step index.""" + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}") + if step_index < 0 or step_index >= num_steps: + raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}") + return step_index < self.first_steps or step_index >= num_steps - self.last_steps + + +def _non_negative_int(value: object, name: str) -> int: + """Validate an integer configuration field without accepting booleans.""" + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise TypeError(f"{name} must be a non-negative integer") + return value diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py new file mode 100644 index 000000000..f528262e5 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Format-agnostic schedule state, linear discovery, and dispatch.""" + +from __future__ import annotations + +from typing import Literal + +import torch +from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase + +from .config import Cosmos3MixedPrecisionConfig +from .strategy import ( + Cosmos3PrecisionStrategy, + Fp8W8A8W8A16Strategy, + Nvfp4W4A4W4A16Strategy, +) + +PrecisionPath = Literal["reasoner", "generation"] +_STRATEGIES: tuple[Cosmos3PrecisionStrategy, ...] = ( + Fp8W8A8W8A16Strategy(), + Nvfp4W4A4W4A16Strategy(), +) + + +def _strategy_for(method: object | None) -> Cosmos3PrecisionStrategy | None: + return next((strategy for strategy in _STRATEGIES if strategy.accepts(method)), None) + + +class Cosmos3MixedPrecisionLinearMethod(LinearMethodBase): + """Dispatch between a checkpoint-native method and dense A16.""" + + def __init__( + self, + base_method: LinearMethodBase, + strategy: Cosmos3PrecisionStrategy, + runtime: Cosmos3MixedPrecisionRuntime, + module_name: str, + path: PrecisionPath, + ) -> None: + self.base_method = base_method + self.strategy = strategy + self.runtime = runtime + self.module_name = module_name + self.path = path + + def create_weights(self, *args, **kwargs) -> None: + self.base_method.create_weights(*args, **kwargs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.strategy.validate_before_processing( + self.base_method, + layer, + self.module_name, + ) + self.base_method.process_weights_after_loading(layer) + self.strategy.validate_after_processing(layer, self.module_name) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if not self.runtime.use_high_precision(self.path): + return self.base_method.apply(layer, x, bias) + return self.strategy.apply_high(layer, x, bias) + + +class Cosmos3MixedPrecisionRuntime: + """Own one transformer's schedule state and wrapped linear inventory.""" + + def __init__(self, config: Cosmos3MixedPrecisionConfig) -> None: + self.config = config + # TODO: Make this request-local before Cosmos3 supports interleaved + # denoising requests on one transformer instance. + self._generation_high_precision = False + + def install(self, transformer: torch.nn.Module) -> None: + components: dict[PrecisionPath, torch.nn.Module] = { + "generation": transformer.gen_layers, + } + if self.config.reasoner == "a16": + components["reasoner"] = transformer.language_model.layers + + wrapped_count = 0 + for path, component in components.items(): + for local_name, layer in component.named_modules(): + if not isinstance(layer, LinearBase): + continue + base_method = getattr(layer, "quant_method", None) + strategy = _strategy_for(base_method) + if strategy is None: + continue + module_name = getattr(layer, "prefix", None) or f"{path}.{local_name}" + layer.quant_method = Cosmos3MixedPrecisionLinearMethod( + base_method, + strategy, + self, + module_name, + path, + ) + wrapped_count += 1 + + if not wrapped_count: + raise ValueError("Cosmos3 mixed precision found no compatible FP8 or NVFP4 ModelOpt linears") + + def use_high_precision(self, path: PrecisionPath) -> bool: + return self.config.reasoner == "a16" if path == "reasoner" else self._generation_high_precision + + def set_step(self, step_index: int, num_steps: int) -> None: + self._generation_high_precision = self.config.use_high_precision(step_index, num_steps) + + def reset(self) -> None: + self._generation_high_precision = False diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py new file mode 100644 index 000000000..8208b7622 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Format-specific live-weight validation and reference A16 materialization.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptFp8LinearMethod, + ModelOptNvFp4LinearMethod, +) +from vllm.model_executor.layers.quantization.utils import nvfp4_emulation_utils + +_NVFP4_BLOCK_SIZE = 16 + + +class Cosmos3PrecisionStrategy: + """Quantization-format boundary behind the shared precision schedule.""" + + def accepts(self, method: object | None) -> bool: + """Return whether this strategy owns a native linear method.""" + raise NotImplementedError + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + """Reject native transforms that cannot support live dequantization.""" + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + """Validate the live native weight representation.""" + raise NotImplementedError + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + """Materialize the live quantized weight as a dense BF16 matrix.""" + raise NotImplementedError + + def apply_high( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ) -> torch.Tensor: + """Apply the reference dense A16 path.""" + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + weight = self.materialize(layer) + expected_shape = (output_size, input_size) + if tuple(weight.shape) != expected_shape: + raise RuntimeError(f"Dense weight shape {tuple(weight.shape)} does not match {expected_shape}") + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.shape[1] != input_size: + raise ValueError(f"Expected activation width {input_size}, got {x_2d.shape[1]}") + weight = weight.to(dtype=x.dtype) + if bias is not None and bias.dtype != x.dtype: + bias = bias.to(dtype=x.dtype) + output = F.linear(x_2d, weight, bias) + return output.view(*x.shape[:-1], output_size) + + +class Fp8W8A8W8A16Strategy(Cosmos3PrecisionStrategy): + """Use native ModelOpt W8A8 or reference dense W8A16.""" + + def accepts(self, method: object | None) -> bool: + return isinstance(method, ModelOptFp8LinearMethod) + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + if hasattr(layer, "pre_quant_scale"): + raise ValueError(f"{module_name} uses unsupported SmoothQuant pre_quant_scale") + weight_scale = getattr(layer, "weight_scale", None) + if not isinstance(weight_scale, torch.Tensor) or weight_scale.numel() != 1: + raise ValueError(f"{module_name} requires serialized tensorwise FP8 weights") + fp8_kernel = getattr(method, "fp8_linear", None) + from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( + MarlinFP8ScaledMMLinearKernel, + ) + + if isinstance(fp8_kernel, MarlinFP8ScaledMMLinearKernel): + raise ValueError( + f"{module_name} selected Marlin FP8, which repacks the only weight copy; " + "use a backend that retains canonical FP8 weights" + ) + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + weight = getattr(layer, "weight", None) + scale = getattr(layer, "weight_scale", None) + if not isinstance(weight, torch.Tensor) or weight.dtype != torch.float8_e4m3fn: + dtype = getattr(weight, "dtype", None) + raise TypeError(f"{module_name} requires a canonical FP8 weight, got {dtype}") + if weight.ndim != 2: + raise ValueError(f"{module_name} requires a rank-2 FP8 weight, got {tuple(weight.shape)}") + if not isinstance(scale, torch.Tensor): + raise ValueError(f"{module_name} has an FP8 weight but no weight_scale") + if scale.numel() != 1: + raise ValueError(f"{module_name} requires one live FP8 weight scale") + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + if weight.shape[0] < input_size or weight.shape[1] < output_size: + raise ValueError( + f"{module_name} live FP8 weight shape {tuple(weight.shape)} does not cover " + f"(K, N)=({input_size}, {output_size})" + ) + _validate_positive_finite_scale(scale, module_name) + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + weight = layer.weight[:input_size, :output_size] + scale = layer.weight_scale.reshape(1).to(device=weight.device, dtype=torch.float32) + return (weight.to(torch.float32) * scale).t().to(torch.bfloat16) + + +class Nvfp4W4A4W4A16Strategy(Cosmos3PrecisionStrategy): + """Use native ModelOpt W4A4 or reference dense W4A16.""" + + def accepts(self, method: object | None) -> bool: + return isinstance(method, ModelOptNvFp4LinearMethod) + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + kernel = getattr(method, "kernel", None) + from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, + ) + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ) + + compatible_kernels = ( + CutlassNvFp4LinearKernel, + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ) + if not isinstance(kernel, compatible_kernels): + kernel_name = type(kernel).__name__ if kernel is not None else "none" + raise ValueError( + f"{module_name} selected unsupported NVFP4 backend {kernel_name}; " + "live W4A16 dequantization requires a CUTLASS-compatible layout" + ) + global_scale = _nvfp4_global_scale(layer) + if global_scale is None or global_scale.numel() != 1: + raise ValueError(f"{module_name} is missing ModelOpt NVFP4 scales") + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + packed = getattr(layer, "weight", None) + scale = getattr(layer, "weight_scale", None) + global_scale = getattr(layer, "weight_global_scale", None) + if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8: + dtype = getattr(packed, "dtype", None) + raise TypeError(f"{module_name} requires a live packed NVFP4 weight, got {dtype}") + if not isinstance(scale, torch.Tensor) or not isinstance(global_scale, torch.Tensor): + raise ValueError(f"{module_name} is missing live NVFP4 scales") + output_size = int(layer.output_size_per_partition) + input_size = int(layer.input_size_per_partition) + if input_size % _NVFP4_BLOCK_SIZE != 0: + raise ValueError(f"{module_name} input size {input_size} is not divisible by {_NVFP4_BLOCK_SIZE}") + expected_weight = (output_size, input_size // 2) + if packed.shape[0] < expected_weight[0] or packed.shape[1] < expected_weight[1]: + raise ValueError( + f"{module_name} live NVFP4 weight shape {tuple(packed.shape)} does not cover {expected_weight}" + ) + _validate_positive_finite_scale(global_scale, module_name) + lut_handle = nvfp4_emulation_utils.kE2M1ToFloat_handle + if lut_handle.val.device != packed.device: + lut_handle.val = lut_handle.val.to(packed.device) + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + output_size = int(layer.output_size_per_partition) + input_size = int(layer.input_size_per_partition) + packed = layer.weight[:output_size, : input_size // 2] + return nvfp4_emulation_utils.dequantize_to_dtype( + packed, + layer.weight_scale, + layer.weight_global_scale, + torch.bfloat16, + _NVFP4_BLOCK_SIZE, + True, + ) + + +def _validate_positive_finite_scale(scale: torch.Tensor, module_name: str) -> None: + values = scale.detach().float() + if not torch.isfinite(values).all() or not (values > 0).all(): + raise ValueError(f"{module_name} has a non-finite or non-positive weight scale") + + +def _nvfp4_global_scale(layer: torch.nn.Module) -> torch.Tensor | None: + for name in ("weight_scale_2", "weight_global_scale"): + value = getattr(layer, name, None) + if isinstance(value, torch.Tensor): + return value + return None diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 49a4e4027..c60817afb 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1822,6 +1822,16 @@ class Cosmos3OmniDiffusersPipeline( def num_timesteps(self): return self._num_timesteps + def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: + setter = getattr(self.transformer, "set_mixed_precision_step", None) + if setter is not None: + setter(step_index, num_steps) + + def _reset_mixed_precision(self) -> None: + resetter = getattr(self.transformer, "reset_mixed_precision", None) + if resetter is not None: + resetter() + @staticmethod def _distilled_unsupported_error(detail: str) -> ValueError: return ValueError( @@ -2819,7 +2829,8 @@ class Cosmos3OmniDiffusersPipeline( # Each CFG-parallel rank runs exactly one branch (rank 0 -> cond, # else uncond), so session keying loads/stores only this rank's branch. cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0 - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) # Outside the interval, scale=1 makes the combined output equal # the cond branch. Every rank remains on the same iteration and @@ -2858,7 +2869,8 @@ class Cosmos3OmniDiffusersPipeline( uncond_cache: tuple = (None, None) keep_uncond_for_cache = self._cache_requires_paired_cfg() - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) cfg_active = _cfg_active_at(t) @@ -2912,7 +2924,8 @@ class Cosmos3OmniDiffusersPipeline( else: # No CFG: a single cond branch per step. Bespoke (state None) keeps # using the transformer-instance cache exactly as before. - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) self._kv_load_und(kv_state, is_negative=False) noise_pred = self.transformer( @@ -2928,6 +2941,7 @@ class Cosmos3OmniDiffusersPipeline( self._kv_capture_und(kv_state, is_negative=False) _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) finally: + self._reset_mixed_precision() # Cosmos3 currently receives a unique request_id rather than a # reusable rollout session id. Retaining its state would only pin # K/V buffers on device after this generation finishes. @@ -3112,7 +3126,8 @@ class Cosmos3OmniDiffusersPipeline( self.transformer.reset_cache() self._cosmos3_branch_caches = {} try: - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0 @@ -3224,6 +3239,7 @@ class Cosmos3OmniDiffusersPipeline( )[0] latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents finally: + self._reset_mixed_precision() self._cosmos3_branch_caches = None self.transformer.reset_cache() return latents diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index fe7a34c8b..1ded3db2f 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -40,6 +40,12 @@ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_ from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm from vllm_omni.platforms import current_omni_platform +from .mixed_precision import ( + Cosmos3MixedPrecisionConfig, + Cosmos3MixedPrecisionRuntime, + resolve_mixed_precision_config, +) + if TYPE_CHECKING: from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook @@ -136,6 +142,26 @@ def _od_config_get(od_config: Any, key: str, default: Any = None) -> Any: return default if value is None else value +def _validate_mixed_precision_runtime( + config: Cosmos3MixedPrecisionConfig | None, + od_config: OmniDiffusionConfig, +) -> None: + if config is None: + return + if get_tensor_model_parallel_world_size() != 1: + raise ValueError("Cosmos3 mixed precision currently supports tensor parallel size 1 only") + if int(getattr(od_config, "max_num_seqs", 1)) != 1: + raise ValueError("Cosmos3 mixed precision currently supports one active request per worker") + if bool(getattr(od_config, "enable_distributed_layerwise_offload", False)): + raise ValueError( + "Cosmos3 mixed precision does not support distributed layer-wise offload " + "because its direct loader bypasses ModelOpt post-load transformations" + ) + parallel_config = getattr(od_config, "parallel_config", None) + if bool(getattr(parallel_config, "use_hsdp", False)): + raise ValueError("Cosmos3 mixed precision has not validated live backend weights under HSDP") + + def _as_bool(value: Any) -> bool: if isinstance(value, str): return value.strip().lower() in {"1", "true", "yes", "on"} @@ -1241,6 +1267,19 @@ class Cosmos3VFMTransformer(nn.Module): dtype = od_config.dtype quant_config = getattr(od_config, "quantization_config", None) if od_config else 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": + logger.info("Cosmos3 checkpoint mixed-precision policy disabled by additional_config") + else: + logger.info( + "Cosmos3 mixed precision active (source=%s): first_steps=%d last_steps=%d reasoner=%s", + mixed_precision_source, + mixed_precision_config.first_steps, + mixed_precision_config.last_steps, + mixed_precision_config.reasoner, + ) + _validate_mixed_precision_runtime(mixed_precision_config, od_config) self.language_model = self._language_model_cls( hidden_size=self.hidden_size, @@ -1300,6 +1339,11 @@ class Cosmos3VFMTransformer(nn.Module): ] ) + self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None + if mixed_precision_config is not None: + self.mixed_precision_runtime = Cosmos3MixedPrecisionRuntime(mixed_precision_config) + self.mixed_precision_runtime.install(self) + self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) self.gen_sp_prepare = Cosmos3GenSPPrepare() self.gen_sp_gather = nn.Identity() @@ -1958,3 +2002,11 @@ class Cosmos3VFMTransformer(nn.Module): def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" self.time_embedder.to(torch.float32) + + def set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: + if self.mixed_precision_runtime is not None: + self.mixed_precision_runtime.set_step(step_index, num_steps) + + def reset_mixed_precision(self) -> None: + if self.mixed_precision_runtime is not None: + self.mixed_precision_runtime.reset()