diff --git a/docs/design/feature/quantization.md b/docs/design/feature/quantization.md index 7ad612d12..8d6cd09c3 100644 --- a/docs/design/feature/quantization.md +++ b/docs/design/feature/quantization.md @@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter. ModelOpt checkpoint metadata is detected from its method, producer metadata, or `quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration. +Unless the input is a per-component dictionary, the resolved configuration is +passed unchanged to every quantization-aware component constructed by the +pipeline. It affects only layers supported by that component's quantization +method; it does not rewrite arbitrary `torch.nn` modules. A per-component +dictionary narrows this global scope through runtime layer-prefix routing. + The factory uses Omni-specific overrides for methods whose diffusion behavior or platform dispatch is not provided by the vLLM registry. The current override set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/ @@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for: Model code must keep precision-sensitive or unsupported modules unquantized. The shared `safe_quant_config()` helper is used by model integrations for normalization and modulation layers where a generic FP8 configuration would be -unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers -are not implicitly quantized just because the main transformer is quantized; -they require explicit component routing and method support. +unsafe. Encoders exposed through the quantization factory receive a global +configuration just like the main transformer, but only their supported vLLM +quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs, +schedulers, and tokenizers are not rewritten automatically. A component map can +select a narrower scope, but it cannot add quantization support to a layer. Weight quantization and quantized KV cache are separate paths. The former is resolved through `quantization_config`; runtime attention/KV-cache quantization diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 05473b15e..397a412e8 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -95,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | | | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | | | Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | | -| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | | +| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) @@ -138,6 +138,39 @@ outputs = omni.generate( ) ``` +### Global and per-component scope + +A plain, global quantization configuration is passed unchanged to every +quantization-aware component constructed by a pipeline. This includes an +eligible encoder when that encoder is implemented with vLLM quantizable +layers; it does not rewrite arbitrary `torch.nn` modules. A structured +component map is the only way to narrow that scope. Pipeline integrations that +do not yet expose an encoder through the quantization factory remain DiT-only. + +For a pipeline that exposes both a transformer and a quantization-aware text +encoder, the scope is: + +| Configuration | Transformer | Text encoder | Components without supported quantizable layers | +|---------------|-------------|--------------|-------------------------------------------------| +| `quantization="fp8"` | FP8 | FP8 | checkpoint precision | +| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision | +| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision | + +Use `quantization_config` for component-selective Python configuration. For +example, quantize only the text decoder: + +```python +omni = Omni( + model="", + quantization_config={"text_encoder": {"method": "fp8"}}, +) +``` + +Component keys are runtime prefixes exposed by the pipeline integration; common +keys include `transformer` and `text_encoder`. Entries may be combined, and +`ignored_layers` can keep named eligible layers in checkpoint precision. Check +the model recipe for supported components and their runtime prefixes. + CLI: ```bash diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 114352343..f6e5dd43b 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -357,12 +357,19 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while ### Online FP8 quantization -MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint -remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and -uses dynamic activation scaling during inference. By default, attention and -MLP linears in the token refiner and main DiT blocks, the condition -projection, and all AdaLN projections use FP8. Patch, timestep, and final -projections remain FP32; the text encoder and VAEs are unchanged. +MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL +text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8 +weights at runtime and uses dynamic activation scaling during inference. By +default, `--quantization fp8` quantizes eligible attention and MLP linears in +the text decoder, token refiner, and main DiT blocks, as well as the DiT +condition and AdaLN projections. The Qwen vision tower, embeddings, norms, +RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections +keep checkpoint precision. + +To select a component, use `--diffusion-quantization-config` with +`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or +`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two +entries can be combined. The shorthand below enables both components. Add this option to an existing H3 server command: @@ -370,6 +377,28 @@ Add this option to an existing H3 server command: --quantization fp8 ``` +#### Single 96 GB GPU, no-offload capacity check + +Use the FL2VA-only partition for this capacity test. Loading the combined +service would also load the Ref2VA DiT and would test a different memory +budget. A no-offload capacity check should contain none of +`--enable-cpu-offload`, `--enable-layerwise-offload`, or +`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement +but does not offload model weights to the CPU. + +The run passes the capacity check when the server initializes, the request +finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the +card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus +32 kHz stereo AAC audio. Report the measured headroom rather than assuming +that every nominal 96 GB SKU exposes the same MiB total. + +As a capacity proxy only, the same five-second case on one B300 measured a +92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its +encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s, +6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but +they are not an RTX PRO 6000 validation: kernels, allocator behavior, and +reported device capacity differ across GPUs. + Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3 resolves the `transformer` component before constructing the DiT, so names do not start with `transformer.`. Entries are exact runtime linear prefixes; a @@ -386,10 +415,10 @@ For example, keep the first main block's attention projections in BF16 with: '{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}' ``` -The structured option replaces `--quantization fp8`. Online FP8 is currently -incompatible with H3 layerwise offload because the offload path produces a -weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor -parallelism and VAE tiling instead. +The structured option replaces `--quantization fp8`. Online FP8 can be used +with H3 layerwise offload and with distributed layerwise offload's full-weight +per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not +supported for runtime-created FP8 weights. ## AMD ROCm (gfx942 / gfx950) @@ -750,6 +779,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22% reduction. +For direct human inspection, see the external +[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/), +which includes matched five-second T2VA, I2VA, and Ref2VA videos. The +[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16) +also record per-task fidelity metrics and provenance without storing generated +media in this repository. + ## TeaCache acceleration TeaCache reuses DiT block residuals across denoising steps when consecutive @@ -834,7 +870,8 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ H3 native `tile` mode only. - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use pure Ulysses. -- Online FP8 is not compatible with layerwise offload. +- Runtime-created FP8 weights do not support the sharded DLO AllGather path; + use `--dlo-no-use-allgather` when combining online FP8 with DLO. - TeaCache and Cache-DiT cannot be enabled on the same server. - Image+audio Ref2VA accepts exactly one image and one audio reference. - Video Ref2VA accepts one or more video files, but not an additional standalone diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index f370d5b49..69600d811 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -110,6 +110,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check(): loader.load_weights(model) +def test_stream_online_quant_weights_offloads_layers_after_processing(): + from vllm.model_executor.model_loader.reload.layerwise import ( + get_layerwise_info, + ) + + events: list[str] = [] + + class _OnlineQuantMethod: + uses_meta_device = True + + class _TrackedLayer(nn.Linear): + def __init__(self, name: str): + super().__init__(2, 2, bias=False) + self.name = name + self.quant_method = _OnlineQuantMethod() + get_layerwise_info(self).load_numel_total = self.weight.numel() + + def to(self, *args, **kwargs): + events.append(self.name) + return super().to(*args, **kwargs) + + class _StreamingModel(nn.Module): + def __init__(self): + super().__init__() + self.first = _TrackedLayer("first") + self.second = _TrackedLayer("second") + + model = _StreamingModel() + weights = iter( + [ + ("first.weight", torch.zeros((2, 2))), + ("second.weight", torch.zeros((2, 2))), + ] + ) + streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights) + + assert next(streamed)[0] == "first.weight" + get_layerwise_info(model.first).reset() + assert next(streamed)[0] == "second.weight" + assert events == ["first"] + + get_layerwise_info(model.second).reset() + with pytest.raises(StopIteration): + next(streamed) + assert events == ["first", "second"] + + +def test_process_weights_skips_completed_online_quant_layer(monkeypatch): + from unittest.mock import Mock + + from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase + from vllm.model_executor.model_loader.reload import layerwise + + class _TrackedLayer(nn.Linear): + def __init__(self): + super().__init__(2, 2, bias=False) + self.to_calls: list[object] = [] + + def to(self, *args, **kwargs): + self.to_calls.append(args[0] if args else kwargs.get("device")) + return super().to(*args, **kwargs) + + model = nn.Module() + model.layer = _TrackedLayer() + quant_method = Mock(spec=QuantizeMethodBase) + quant_method.uses_meta_device = True + model.layer.quant_method = quant_method + model.layer._already_called_process_weights_after_loading = True + finalize = Mock() + monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize) + + loader = _make_loader_with_weights([]) + loader._process_weights_after_loading(model, torch.device("cuda")) + + finalize.assert_called_once_with(model, model_config=None) + quant_method.process_weights_after_loading.assert_not_called() + assert model.layer.to_calls == [] + + class _ConfigAwareModel(nn.Module): def __init__(self, *, od_config): super().__init__() diff --git a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py index 3900c0c94..ff9e0e110 100644 --- a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py +++ b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py @@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import ( Flux2Pipeline, _resolve_component_quant_config, ) +from vllm_omni.quantization import ComponentQuantizationConfig pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_resolve_component_quant_config_routes_components(): text_encoder_config = object() - component_config = Mock() - component_config.resolve.side_effect = { - "text_encoder": text_encoder_config, - "transformer": None, - }.get + component_config = ComponentQuantizationConfig( + component_configs={"text_encoder": text_encoder_config}, + default_config=None, + ) assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config assert _resolve_component_quant_config(component_config, "transformer") is None diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index f872e381f..34250abf2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -5,7 +5,7 @@ import json import sys from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import numpy as np import pytest @@ -363,8 +363,18 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once( if expected_partition == "ref2va": assert pipeline._transformer_for_task("ref2va") is pipeline.transformer assert [source.model_or_path for source in pipeline.weights_sources] == [ - str(tmp_path / partition_name) for partition_name in source_partitions + *(str(tmp_path / partition_name) for partition_name in source_partitions), + str(component_path), ] + assert [source.subfolder for source in pipeline.weights_sources] == [ + *("transformer" for _ in source_partitions), + "text_encoder", + ] + expected_prefixes = ["transformer."] + if expected_dits == 2: + expected_prefixes.append("transformers_ref.") + expected_prefixes.append("text_encoder.") + assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes expected_patterns = ( pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS if expected_partition == "combined" @@ -936,6 +946,66 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): assert list(encoder.named_parameters()) == [] +def test_global_quant_config_is_shared_by_dit_and_encoder(): + from vllm_omni.quantization import build_quant_config, resolve_component_quant_config + + global_config = build_quant_config("fp8") + + assert resolve_component_quant_config(global_config, "transformer") is global_config + assert resolve_component_quant_config(global_config, "text_encoder") is global_config + + +def test_minimax_h3_fp8_component_selection(): + from vllm_omni.quantization import build_quant_config, resolve_component_quant_config + + dit_only = build_quant_config({"transformer": {"method": "fp8"}}) + encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}}) + + assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8" + assert resolve_component_quant_config(dit_only, "text_encoder") is None + assert resolve_component_quant_config(encoder_only, "transformer") is None + assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8" + + +def test_text_encoder_rejects_serialized_fp8(): + from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config + + config = SimpleNamespace( + get_name=lambda: "fp8", + is_checkpoint_fp8_serialized=True, + ) + with pytest.raises(ValueError, match="online FP8 only"): + _validate_encoder_quant_config(config) + + +def test_text_encoder_linear_delegates_quantization_to_vllm_factory(): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module + + group = SimpleNamespace(rank_in_group=0, world_size=1) + method = UnquantizedLinearMethod() + quant_config = Mock() + quant_config.get_quant_method.return_value = method + prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj" + + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear( + group, + input_size=4, + intermediate_size=8, + dtype=torch.bfloat16, + quant_config=quant_config, + prefix=prefix, + ) + + assert linear.quant_method is method + quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix) + + def test_no_offload_keeps_text_encoder_resident(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -978,15 +1048,21 @@ def test_model_offload_uses_hooked_text_encoder_call(): pipeline.text_encoder.load_to_device.assert_not_called() -def test_layerwise_offload_releases_text_encoder(): +@pytest.mark.parametrize( + "offload_flag", + ["enable_layerwise_offload", "enable_distributed_layerwise_offload"], +) +def test_layerwise_offload_releases_text_encoder(offload_flag): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = SimpleNamespace( enable_cpu_offload=False, - enable_layerwise_offload=True, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, ) + setattr(pipeline.od_config, offload_flag, True) pipeline.text_encoder = Mock() expected = torch.ones(2, 3) pipeline.text_encoder.encode_ids.return_value = expected @@ -1988,21 +2064,25 @@ def _one_layer_text_encoder(): group = _SingleRankEncoderGroup(0) layer = nn.Module() layer.self_attn = nn.Module() - layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear( - group, - hidden_size=_ENCODER_HIDDEN, - num_heads=_ENCODER_NUM_HEADS, - num_kv_heads=_ENCODER_NUM_KV_HEADS, - head_dim=_ENCODER_HEAD_DIM, - dtype=torch.float32, - ) layer.mlp = nn.Module() - layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear( - group, - input_size=_ENCODER_HIDDEN, - intermediate_size=_ENCODER_INTERMEDIATE, - dtype=torch.float32, - ) + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear( + group, + hidden_size=_ENCODER_HIDDEN, + num_heads=_ENCODER_NUM_HEADS, + num_kv_heads=_ENCODER_NUM_KV_HEADS, + head_dim=_ENCODER_HEAD_DIM, + dtype=torch.float32, + ) + layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear( + group, + input_size=_ENCODER_HIDDEN, + intermediate_size=_ENCODER_INTERMEDIATE, + dtype=torch.float32, + ) layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN) encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) @@ -2035,6 +2115,7 @@ def _write_text_encoder_checkpoint(path, *, omit=()): safetensors.torch.save_file(tensors, str(path / "model.safetensors")) index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")} (path / "model.safetensors.index.json").write_text(json.dumps(index)) + return tensors @pytest.mark.parametrize( @@ -2056,10 +2137,10 @@ def test_encoder_load_reports_missing_fused_source_shards( tmp_path, omitted_sources, expected_detail, unreported_target ): encoder = _one_layer_text_encoder() - _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources) + weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources) with pytest.raises(RuntimeError) as excinfo: - encoder._load_weights(str(tmp_path)) + encoder.load_weights(weights.items()) message = str(excinfo.value) assert expected_detail in message @@ -2068,17 +2149,17 @@ def test_encoder_load_reports_missing_fused_source_shards( def test_encoder_load_rejects_unloaded_plain_param(tmp_path): encoder = _one_layer_text_encoder() - _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",)) + weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",)) with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"): - encoder._load_weights(str(tmp_path)) + encoder.load_weights(weights.items()) def test_encoder_load_places_every_source_in_its_own_rows(tmp_path): encoder = _one_layer_text_encoder() - _write_text_encoder_checkpoint(tmp_path) + weights = _write_text_encoder_checkpoint(tmp_path) - encoder._load_weights(str(tmp_path)) + encoder.load_weights(weights.items()) params = dict(encoder.named_parameters()) q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index b9ae3335c..5f220d6ab 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -211,3 +211,18 @@ def test_pipeline_resolves_transformer_component_quant_config(): assert transformer_config.ignored_layers == ignored_layers assert _resolve_component_quant_config(component_config, "transformer") is transformer_config assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config + + +def test_pipeline_strips_prequantized_text_encoder_config(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config + from vllm_omni.quantization import ComponentQuantizationConfig + + prequantized = Mock() + prequantized.get_name.return_value = "modelopt" + online_fp8 = Mock() + online_fp8.get_name.return_value = "fp8" + component_config = ComponentQuantizationConfig({"text_encoder": prequantized}) + + assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8 + assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None + assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index fd004d056..4c0cceb68 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -139,6 +139,41 @@ class TestDistributedLayerwiseOffloadHook: assert shard.numel() == 4 assert torch.equal(shard, _make_values(10.0)) + def test_prefetch_preserves_transposed_weight_stride( + self, + dist_group, + patched_offload_runtime, + ): + """Online-FP8 Cutlass weights must retain their transposed layout.""" + + class StridedBlock(nn.Module): + def __init__(self, start: float): + super().__init__() + base = torch.arange(start, start + 12).reshape(3, 4) + self.weight = nn.Parameter(base.t(), requires_grad=False) + + current_block = StridedBlock(1.0) + next_block = StridedBlock(20.0) + expected = next_block.weight.detach().clone() + expected_stride = next_block.weight.stride() + assert expected_stride == (1, 4) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook.initialize_hook(current_block) + hook.prefetch_layer(slot=0, non_blocking=False) + + assert next_block.weight.stride() == expected_stride + assert torch.equal(next_block.weight, expected) + def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime): """DTensor wrapper should be preserved through prefetch/offload cycle.""" current_block = TinyBlock(_make_values(1.0)) @@ -386,6 +421,24 @@ class TestPinnedResidentLayerGroup: group.offload() + def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime): + block = nn.Module() + weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t() + expected = weight.detach().clone() + expected_stride = weight.stride() + block.weight = nn.Parameter(weight) + group = PinnedResidentLayerGroup( + [block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + ) + + group.load() + assert torch.equal(block.weight, expected) + assert block.weight.stride() == expected_stride + group.offload() + def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self): module = nn.Linear(2, 2, bias=False) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 5fd80e75c..f62bd6472 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -79,11 +79,11 @@ def dist_group(): @pytest.fixture -def patched_offload_runtime(mocker): - mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream) - mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent) - mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) - mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream) +def patched_offload_runtime(monkeypatch): + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream) class TinyBlock(nn.Module): @@ -128,6 +128,33 @@ class TestLayerwiseOffloadHook: assert current_block.weight.to_local().shape == torch.Size([4]) assert not hook.is_materialized + def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime): + """Online-FP8 Cutlass weights must retain their transposed layout.""" + + class StridedBlock(nn.Module): + def __init__(self, start: float): + super().__init__() + base = torch.arange(start, start + 12).reshape(3, 4) + self.weight = nn.Parameter(base.t(), requires_grad=False) + + current_block = StridedBlock(1.0) + next_block = StridedBlock(20.0) + expected = next_block.weight.detach().clone() + expected_stride = next_block.weight.stride() + assert expected_stride == (1, 4) + + hook = LayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + stream=DummyStream(), + pin_memory=False, + ) + hook.initialize_hook(current_block) + hook.prefetch_layer(non_blocking=False) + + assert next_block.weight.stride() == expected_stride + assert torch.equal(next_block.weight, expected) + class _DummyBlock(nn.Module): def __init__(self): diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py index c8b3837e2..9a306b6e4 100644 --- a/tests/diffusion/quantization/test_component_routing.py +++ b/tests/diffusion/quantization/test_component_routing.py @@ -16,10 +16,11 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( ) from vllm_omni.quantization.component_config import ( ComponentQuantizationConfig, + resolve_component_quant_config, ) from vllm_omni.quantization.inc_config import OmniINCConfig -pytestmark = [pytest.mark.core_model] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # --------------------------------------------------------------------------- @@ -347,6 +348,24 @@ class TestTalkerVisualRouting: # =================================================================== +class TestResolveComponentQuantConfig: + def test_global_config_applies_to_each_quantization_aware_component(self): + config = _MockQuantConfig("fp8") + + assert resolve_component_quant_config(config, "transformer") is config + assert resolve_component_quant_config(config, "text_encoder") is config + + def test_component_map_is_the_only_scope_override(self): + transformer = _MockQuantConfig("fp8") + config = ComponentQuantizationConfig( + component_configs={"transformer": transformer}, + default_config=None, + ) + + assert resolve_component_quant_config(config, "transformer") is transformer + assert resolve_component_quant_config(config, "text_encoder") is None + + class TestComponentResolve: def test_longest_prefix_match(self): a = _MockQuantConfig("a") diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d1a1ed15e..bde9a1543 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -298,6 +298,60 @@ class DiffusersPipelineLoader: for source in sources: yield from self._get_weights_iterator(source, model=model) + @staticmethod + def _stream_online_quant_weights_to_cpu( + model: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Generator[tuple[str, torch.Tensor], None, None]: + """Offload each online-quantized layer as soon as it is complete. + + Upstream vLLM's online layerwise loader materializes and quantizes a + layer synchronously while consuming ``weights``. Generator execution + resumes after that weight has been consumed, which gives us a safe + point to move completed layers to CPU before the next layer is loaded. + This bounds accelerator residency during CPU-offloaded model startup + instead of retaining the entire quantized model until loading ends. + """ + from vllm.model_executor.model_loader.reload.layerwise import ( + get_layerwise_info, + ) + + pending = { + module + for module in model.modules() + if getattr(getattr(module, "quant_method", None), "uses_meta_device", False) + and get_layerwise_info(module).can_load() + } + offloaded = 0 + + def offload_completed() -> None: + nonlocal offloaded + for module in tuple(pending): + if get_layerwise_info(module).can_load(): + continue + module.to("cpu") + pending.remove(module) + offloaded += 1 + + for weight in weights: + # This runs after the consumer has handled the previous yield. + offload_completed() + yield weight + offload_completed() + + # Quantization workspaces and the old accelerator-side parameter + # storages are now reusable cache blocks. Release them before the + # remaining (unquantized) model tensors are copied to CPU; otherwise + # the cached quantization footprint and final offload overlap in the + # process-level startup peak. + if offloaded: + torch.accelerator.empty_cache() + + logger.info( + "Stream-offloaded %d online-quantized layers to CPU during weight loading", + offloaded, + ) + def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]: return tuple( cast( @@ -403,7 +457,10 @@ class DiffusersPipelineLoader: if load_format == "diffusers": cast(DiffusersAdapterPipeline, model).load_weights() else: - self.load_weights(model) + if offload_after_quant: + self.load_weights(model, stream_online_quant_to_cpu=True) + else: + self.load_weights(model) self._process_weights_after_loading(model, target_device) if offload_after_quant: @@ -498,6 +555,14 @@ class DiffusersPipelineLoader: continue if has_online_quant: + # finalize_layerwise_processing() and the synchronous online + # loader already processed these layers. Avoid moving their + # quantized CPU weights back to the accelerator merely to call + # an idempotent no-op; doing so rebuilds a large CUDA allocator + # cache and defeats streaming CPU offload's startup-memory bound. + if getattr(module, "_already_called_process_weights_after_loading", False): + continue + # Online quant may leave straggler params on the ``meta`` device. # Move only real (non-meta) params onto the target device for # processing and restore them afterward, mirroring upstream vLLM's @@ -535,9 +600,17 @@ class DiffusersPipelineLoader: if needs_device_move: module.to(module_device) - def load_weights(self, model: nn.Module) -> None: + def load_weights( + self, + model: nn.Module, + *, + stream_online_quant_to_cpu: bool = False, + ) -> None: weights_to_load = self._get_expected_parameter_names(model) - loaded_weights = model.load_weights(self.get_all_weights(model)) + weights = self.get_all_weights(model) + if stream_online_quant_to_cpu: + weights = self._stream_online_quant_weights_to_cpu(model, weights) + loaded_weights = model.load_weights(weights) self.counter_after_loading_weights = time.perf_counter() logger.info_once( diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 0dbfe0f8c..973f8b7ff 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -39,16 +39,11 @@ from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import resolve_component_quant_config as _resolve_component_quant_config logger = logging.getLogger(__name__) -def _resolve_component_quant_config(quant_config, component: str): - if hasattr(quant_config, "resolve"): - return quant_config.resolve(component) - return quant_config - - class Flux2ImageProcessor(VaeImageProcessor): """Image processor to preprocess the reference image for Flux2.""" diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 2e07ace96..88f0ab193 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -30,14 +30,16 @@ complete ``[seq, 5120]`` hidden state. from __future__ import annotations import itertools -import json import re +from collections.abc import Iterable from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm @@ -47,6 +49,19 @@ MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 logger = init_logger(__name__) +def _validate_encoder_quant_config(quant_config: QuantizationConfig | None) -> None: + if ( + quant_config is not None + and quant_config.get_name() == "fp8" + and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False)) + ): + raise ValueError("MiniMax-H3 text encoder supports online FP8 only, not serialized FP8 checkpoints") + + +def _is_quantized_linear(module: LinearBase) -> bool: + return not isinstance(module.quant_method, UnquantizedLinearMethod) + + def _default_weight_loader( param: nn.Parameter, loaded_weight: torch.Tensor, @@ -114,7 +129,7 @@ class MiniMaxH3Qwen3VLVocabParallelEmbedding(nn.Module): param.data.copy_(loaded_weight.narrow(0, start_idx, shard_size)) -class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): +class MiniMaxH3Qwen3VLMergedColumnParallelLinear(LinearBase): """Packed gate/up projection sharded along the output dimension.""" def __init__( @@ -123,23 +138,43 @@ class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): input_size: int, intermediate_size: int, dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: - super().__init__() self.group = group - self.input_size = input_size self.intermediate_size = intermediate_size tp_rank, tp_size = _tp_range(group) assert intermediate_size % tp_size == 0, ( f"intermediate_size {intermediate_size} must be divisible by text_encoder_tp_size {tp_size}" ) self.intermediate_size_per_partition = intermediate_size // tp_size - self.weight = nn.Parameter(torch.empty(2 * self.intermediate_size_per_partition, input_size, dtype=dtype)) - self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + self.output_dtype = dtype + super().__init__( + input_size=input_size, + output_size=2 * intermediate_size, + bias=False, + params_dtype=dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=True, + ) + self.quant_method.create_weights( + self, + input_size_per_partition=input_size, + output_partition_sizes=[ + self.intermediate_size_per_partition, + self.intermediate_size_per_partition, + ], + input_size=input_size, + output_size=2 * intermediate_size, + params_dtype=dtype, + weight_loader=self.weight_loader, + ) self._tp_rank = tp_rank self._tp_size = tp_size def forward(self, input_: torch.Tensor) -> torch.Tensor: - return F.linear(input_, self.weight) + return self.quant_method.apply(self, input_) def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None @@ -152,7 +187,7 @@ class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): param.data[0:shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size)) -class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): +class MiniMaxH3Qwen3VLQKVParallelLinear(LinearBase): """QKV projection with GQA head sharding along the output dimension.""" def __init__( @@ -163,8 +198,9 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): num_kv_heads: int, head_dim: int, dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: - super().__init__() self.group = group self.hidden_size = hidden_size self.num_heads = num_heads @@ -179,15 +215,33 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): ) self.local_num_heads = num_heads // tp_size self.local_num_kv_heads = num_kv_heads // tp_size + self.output_dtype = dtype q_local = self.local_num_heads * head_dim kv_local = self.local_num_kv_heads * head_dim - self.weight = nn.Parameter(torch.empty(q_local + 2 * kv_local, hidden_size, dtype=dtype)) - self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + output_size = (num_heads + 2 * num_kv_heads) * head_dim + super().__init__( + input_size=hidden_size, + output_size=output_size, + bias=False, + params_dtype=dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=True, + ) + self.quant_method.create_weights( + self, + input_size_per_partition=hidden_size, + output_partition_sizes=[q_local, kv_local, kv_local], + input_size=hidden_size, + output_size=output_size, + params_dtype=dtype, + weight_loader=self.weight_loader, + ) self._tp_rank = tp_rank self._tp_size = tp_size def forward(self, input_: torch.Tensor) -> torch.Tensor: - return F.linear(input_, self.weight) + return self.quant_method.apply(self, input_) def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None @@ -240,7 +294,7 @@ def _verify_fused_shards_complete( raise RuntimeError(f"MiniMax H3 text encoder fused weights are missing source shards: {details}") -class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): +class MiniMaxH3Qwen3VLRowParallelLinear(LinearBase): def __init__( self, group: Any, @@ -248,17 +302,33 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): output_size: int, dtype: torch.dtype, input_is_parallel: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: - super().__init__() self.group = group - self.input_size = input_size - self.output_size = output_size self.input_is_parallel = input_is_parallel + self.output_dtype = dtype tp_rank, tp_size = _tp_range(group) assert input_size % tp_size == 0, f"input_size {input_size} must be divisible by text_encoder_tp_size {tp_size}" self.input_size_per_partition = input_size // tp_size - self.weight = nn.Parameter(torch.empty(output_size, self.input_size_per_partition, dtype=dtype)) - self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + super().__init__( + input_size=input_size, + output_size=output_size, + bias=False, + params_dtype=dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=True, + ) + self.quant_method.create_weights( + self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=dtype, + weight_loader=self.weight_loader, + ) self._tp_rank = tp_rank self._tp_size = tp_size @@ -268,7 +338,7 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): else: split_input = input_.split(self.input_size_per_partition, dim=-1) input_parallel = split_input[self._tp_rank].contiguous() - output_parallel = F.linear(input_parallel, self.weight) + output_parallel = self.quant_method.apply(self, input_parallel) if self._tp_size > 1: # Reduce in fp32: the reference path accumulates the full (K=8192) # dot product inside a single cuBLAS GEMM before rounding to bf16. @@ -276,7 +346,7 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): # and amplify error on this model's large-magnitude activations. output_parallel = output_parallel.float() self.group.all_reduce(output_parallel) - output_parallel = output_parallel.to(self.weight.dtype) + output_parallel = output_parallel.to(self.output_dtype) return output_parallel def weight_loader( @@ -665,7 +735,14 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module): class MiniMaxH3Qwen3VLTextAttention(nn.Module): - def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.group = group self.config = config @@ -683,6 +760,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): num_kv_heads=self.num_kv_heads, head_dim=self.head_dim, dtype=dtype, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", ) self.o_proj = MiniMaxH3Qwen3VLRowParallelLinear( group, @@ -690,6 +769,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): output_size=self.hidden_size, dtype=dtype, input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", ) self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) @@ -703,14 +784,23 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): hidden_shape = (*input_shape, -1, self.head_dim) q_local = self.qkv_proj.local_num_heads * self.head_dim kv_local = self.qkv_proj.local_num_kv_heads * self.head_dim - q_weight = self.qkv_proj.weight[0:q_local] - k_weight = self.qkv_proj.weight[q_local : q_local + kv_local] - v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local] - # Three separate GEMMs keep the numerics identical to the reference - # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding). - query_states = self.q_norm(F.linear(hidden_states, q_weight).view(hidden_shape)).transpose(1, 2) - key_states = self.k_norm(F.linear(hidden_states, k_weight).view(hidden_shape)).transpose(1, 2) - value_states = F.linear(hidden_states, v_weight).view(hidden_shape).transpose(1, 2) + if _is_quantized_linear(self.qkv_proj): + query, key, value = self.qkv_proj(hidden_states).split( + (q_local, kv_local, kv_local), + dim=-1, + ) + else: + q_weight = self.qkv_proj.weight[0:q_local] + k_weight = self.qkv_proj.weight[q_local : q_local + kv_local] + v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local] + # Three separate GEMMs keep the numerics identical to the reference + # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding). + query = F.linear(hidden_states, q_weight) + key = F.linear(hidden_states, k_weight) + value = F.linear(hidden_states, v_weight) + query_states = self.q_norm(query.view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(key.view(hidden_shape)).transpose(1, 2) + value_states = value.view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin) @@ -734,7 +824,14 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): class MiniMaxH3Qwen3VLTextMLP(nn.Module): - def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size @@ -743,6 +840,8 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module): input_size=self.hidden_size, intermediate_size=self.intermediate_size, dtype=dtype, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", ) self.down_proj = MiniMaxH3Qwen3VLRowParallelLinear( group, @@ -750,22 +849,46 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module): output_size=self.hidden_size, dtype=dtype, input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", ) def forward(self, x: torch.Tensor) -> torch.Tensor: ip_local = self.gate_up_proj.intermediate_size_per_partition - gate = F.linear(x, self.gate_up_proj.weight[0:ip_local]) - up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local]) + if _is_quantized_linear(self.gate_up_proj): + gate, up = self.gate_up_proj(x).split(ip_local, dim=-1) + else: + gate = F.linear(x, self.gate_up_proj.weight[0:ip_local]) + up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local]) x = F.silu(gate) * up return self.down_proj(x) class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): - def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.hidden_size = config.hidden_size - self.self_attn = MiniMaxH3Qwen3VLTextAttention(group, config, dtype) - self.mlp = MiniMaxH3Qwen3VLTextMLP(group, config, dtype) + self.self_attn = MiniMaxH3Qwen3VLTextAttention( + group, + config, + dtype, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = MiniMaxH3Qwen3VLTextMLP( + group, + config, + dtype, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype) self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm( config.hidden_size, eps=config.rms_norm_eps, dtype=dtype @@ -791,7 +914,15 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): class MiniMaxH3Qwen3VLTextModel(nn.Module): """TP-sharded Qwen3-VL text model returning unnormalized layer-50 states.""" - def __init__(self, group: Any, config: Any, selected_layer: int, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + selected_layer: int, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "text_encoder.text_model", + ) -> None: super().__init__() self.group = group self.config = config @@ -804,7 +935,16 @@ class MiniMaxH3Qwen3VLTextModel(nn.Module): ) self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config) self.layers = nn.ModuleList( - [MiniMaxH3Qwen3VLTextDecoderLayer(group, config, dtype) for _ in range(self.num_layers)] + [ + MiniMaxH3Qwen3VLTextDecoderLayer( + group, + config, + dtype, + quant_config=quant_config, + prefix=f"{prefix}.layers.{layer_idx}", + ) + for layer_idx in range(self.num_layers) + ] ) def _deepstack_process( @@ -951,10 +1091,13 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): device: torch.device, load_model: bool, encoder_group: Any | None = None, + quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() self.device_target = device self.encoder_group = encoder_group + self.quant_config = quant_config + _validate_encoder_quant_config(quant_config) self.image_token_id = 151655 self.video_token_id = 151656 self._tp_size = 1 @@ -975,9 +1118,9 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): config.text_config, MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, dtype, + quant_config=quant_config, ) self.text_model.to(dtype=dtype) - self._load_weights(model_path) logger.info( "MiniMax H3 Qwen3-VL encoder: %d retained decoder layers, text_encoder_tp_size=%d, vision replicated", MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, @@ -1032,48 +1175,37 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): return ("text_model." + rest, None) return None - def _load_weights(self, model_path: str) -> None: - import os - - import safetensors.torch - - index_path = os.path.join(model_path, "model.safetensors.index.json") - with open(index_path, encoding="utf-8") as handle: - index = json.load(handle) - weight_map: dict[str, str] = index["weight_map"] - files = sorted(set(weight_map.values())) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params = dict(self.named_parameters()) loaded: set[str] = set() - # Seeded from the modules that own a multi-shard loader, so a fused weight - # that receives no source shard at all is caught too, not only a partially - # filled one. + # Seeded from modules that own a multi-shard loader, so a fused weight + # receiving no source shard is caught as well as a partially filled one. expected_fused_shards: dict[str, tuple[tuple[str, str | int], ...]] = { f"{module_name}.weight": sources for module_name, module in self.named_modules() if (sources := _fused_source_shards(module)) is not None } loaded_fused_shards: dict[str, set[str | int]] = {name: set() for name in expected_fused_shards} - for shard in files: - tensors = safetensors.torch.load_file(os.path.join(model_path, shard), device="cpu") - for name, tensor in tensors.items(): - mapped = self._map_weight_name(name) - if mapped is None: - continue - param_name, shard_id = mapped - param = params.get(param_name) - if param is None: - logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name) - continue - weight_loader = getattr(param, "weight_loader", _default_weight_loader) - weight_loader(param, tensor, shard_id) - loaded.add(param_name) - if shard_id is not None and param_name in loaded_fused_shards: - loaded_fused_shards[param_name].add(shard_id) + for name, tensor in weights: + mapped = self._map_weight_name(name) + if mapped is None: + continue + param_name, shard_id = mapped + param = params.get(param_name) + if param is None: + logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name) + continue + weight_loader = getattr(param, "weight_loader", _default_weight_loader) + weight_loader(param, tensor, shard_id) + loaded.add(param_name) + if shard_id is not None and param_name in loaded_fused_shards: + loaded_fused_shards[param_name].add(shard_id) _verify_fused_shards_complete(expected_fused_shards, loaded_fused_shards) missing = sorted(set(params) - loaded) if missing: # Listed in full: this aborts startup, so the message is the only diagnostic. raise RuntimeError(f"MiniMax H3 text encoder weights not loaded: {len(missing)} params: {missing}") + return loaded def load_to_device(self) -> None: if not self.is_loaded: diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 3b96b8ab7..5de203721 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -20,6 +20,7 @@ import torch.nn as nn from PIL import Image from transformers import Qwen2TokenizerFast, Qwen3VLProcessor from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion import envs from vllm_omni.diffusion.cache.cachedit import ( @@ -54,6 +55,12 @@ from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import ( + resolve_component_quant_config as _resolve_component_quant_config, +) +from vllm_omni.quantization.component_config import ( + resolve_encoder_quant_config as _resolve_encoder_quant_config, +) from .condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, @@ -128,6 +135,13 @@ MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = { } +def _resolve_minimax_h3_text_encoder_quant_config( + quant_config: QuantizationConfig | None, +) -> QuantizationConfig | None: + resolved = _resolve_component_quant_config(quant_config, "text_encoder") + return _resolve_encoder_quant_config(resolved) + + def _minimax_h3_partition_for_task( task_type: str | None, model: str | None = None, @@ -175,12 +189,6 @@ def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None: return DMD2SigmaSchedule.from_metadata(release) -def _resolve_component_quant_config(quant_config, component: str): - if hasattr(quant_config, "resolve"): - return quant_config.resolve(component) - return quant_config - - def _minimax_h3_post_process(output, output_type: str = "np"): """Convert the joint video/audio output without capturing worker state. @@ -688,12 +696,24 @@ class MiniMaxH3Pipeline( ) self.text_encoder_tp_size = text_encoder_tp_size self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size) + load_text_encoder = rank < text_encoder_tp_size self.text_encoder = MiniMaxH3Qwen3VLEncoder( os.path.join(model_path, "text_encoder"), device=self.device, - load_model=rank < text_encoder_tp_size, + load_model=load_text_encoder, encoder_group=self.text_encoder_group, + quant_config=_resolve_minimax_h3_text_encoder_quant_config(od_config.quantization_config), ) + if load_text_encoder: + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=str(model_path), + subfolder="text_encoder", + revision=od_config.revision, + prefix="text_encoder.", + fall_back_to_pt=False, + ) + ) stage_components = bool( od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) ) @@ -725,7 +745,7 @@ class MiniMaxH3Pipeline( def source_prefix(item: tuple[str, torch.Tensor]) -> str: name, _ = item prefix = name.partition(".")[0] + "." - if prefix in {"transformer.", "transformers_ref."}: + if prefix in {"transformer.", "transformers_ref.", "text_encoder."}: return prefix raise ValueError(f"unexpected MiniMax-H3 weight {name!r}") @@ -735,17 +755,16 @@ class MiniMaxH3Pipeline( if prefix in loaded_prefixes: raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous") loaded_prefixes.add(prefix) - transformer = getattr(self, prefix.removesuffix(".")) - loaded = transformer.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights) - transformer.post_load_weights() + component = getattr(self, prefix.removesuffix(".")) + loaded = component.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights) + if prefix != "text_encoder.": + component.post_load_weights() loaded_with_prefix.update(prefix + name for name in loaded) - # The text encoder and both VAEs load eagerly in ``__init__`` rather - # than through ``weights_sources``. Record them for the runner's strict - # missing-parameter check. For the text encoder that report is exact - # because its loader raises on any unloaded parameter or partially filled - # fused parameter; weaken that and gaps here become invisible. The two - # VAEs carry no equivalent guarantee. - for component_name in ("text_encoder", "video_vae", "audio_vae"): + # Both VAEs load eagerly in ``__init__`` rather than through + # ``weights_sources``. The text encoder uses the shared component + # loader so online quantization and offload processing follow the same + # path as the DiT. + for component_name in ("video_vae", "audio_vae"): component = getattr(self, component_name) loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) return loaded_with_prefix diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 7da7c8d37..2d6a34978 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -178,7 +178,16 @@ class DistributedLayerwiseOffloadHook(ModelHook): if m["name"] in self.next_block_parameters else self.next_block_buffers[m["name"]] ) - repoint.append((target, dtype, m["offset"], m["numel"], m["shape"])) + repoint.append( + ( + target, + dtype, + m["offset"], + m["numel"], + m["shape"], + m["stride"], + ) + ) self._cached_repoint.append(repoint) # Pre-compute AG output sizes (avoid sum() per layer). @@ -228,9 +237,15 @@ class DistributedLayerwiseOffloadHook(ModelHook): # parameter as an mmap view avoids a private full-model # copy in every worker. local_t = mmap_transform(local_t) - weights_with_local.append((name, t, local_t)) + stride = local_t.stride() + storage_numel = ( + 0 + if local_t.numel() == 0 + else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride)) + ) + weights_with_local.append((name, t, local_t, storage_numel, stride)) - total_numel = sum(local.numel() for _, _, local in weights_with_local) + total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local) # Equal-sized shards (ceil division) for all_gather_into_tensor shard_size = (total_numel + dp_size - 1) // dp_size # ceil @@ -242,8 +257,13 @@ class DistributedLayerwiseOffloadHook(ModelHook): shard = torch.zeros(shard_size, dtype=dtype, device="cpu") current_offset = 0 - for name, original_tensor, local_tensor in weights_with_local: - numel = local_tensor.numel() + for ( + name, + original_tensor, + local_tensor, + storage_numel, + stride, + ) in weights_with_local: if dtype not in dtype_metadata: dtype_metadata[dtype] = [] # Offsets remain relative to the FULL flattened buffer @@ -252,27 +272,48 @@ class DistributedLayerwiseOffloadHook(ModelHook): { "name": name, "offset": current_offset, - "numel": numel, + "numel": storage_numel, "shape": local_tensor.shape, + "stride": stride, } ) # Copy ONLY the portion within [shard_start, shard_end) overlap_start = max(current_offset, shard_start) - overlap_end = min(current_offset + numel, shard_end) + overlap_end = min(current_offset + storage_numel, shard_end) if overlap_start < overlap_end: + if local_tensor.is_contiguous(): + flat_storage = local_tensor.flatten() + else: + # Online FP8 stores Cutlass weights as transposed views + # (e.g. stride=(1, K)). Flattening such a tensor in + # logical order and later rebuilding it with .view() + # changes its layout and makes scaled_mm reject it. + # Pack the physical storage order and preserve the + # original stride for zero-copy reconstruction. + flat_storage = torch.zeros( + storage_numel, + dtype=dtype, + device=local_tensor.device, + ) + physical_view = torch.as_strided( + flat_storage, + size=local_tensor.shape, + stride=stride, + ) + physical_view.copy_(local_tensor) src_start = overlap_start - current_offset src_end = overlap_end - current_offset dst_start = overlap_start - shard_start dst_end = overlap_end - shard_start - shard[dst_start:dst_end].copy_(local_tensor.flatten()[src_start:src_end]) + shard[dst_start:dst_end].copy_(flat_storage[src_start:src_end]) # Replace original tensor with placeholder (frees CPU storage) set_tensor_storage( original_tensor, make_offload_placeholder(original_tensor), ) - current_offset += numel + current_offset += storage_numel if pin_memory: shard = shard.pin_memory() @@ -365,10 +406,14 @@ class DistributedLayerwiseOffloadHook(ModelHook): self._shared_slot_group[slot] = self._group_id # Re-point using cached metadata (avoids per-layer dict lookups). - for target, dtype, offset, numel, shape in self._cached_repoint[slot]: + for target, dtype, offset, numel, shape, stride in self._cached_repoint[slot]: set_tensor_storage( target, - gpu_weights[dtype][offset : offset + numel].view(shape), + torch.as_strided( + gpu_weights[dtype][offset : offset + numel], + size=shape, + stride=stride, + ), ) def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None: @@ -577,7 +622,11 @@ class PinnedResidentLayerGroup: for meta in metas: set_tensor_storage( targets[meta["name"]], - gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]].view(meta["shape"]), + torch.as_strided( + gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]], + size=meta["shape"], + stride=meta["stride"], + ), ) current_omni_platform.current_stream().wait_event(ready) diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 9979d01b1..662e1304f 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -139,29 +139,57 @@ class LayerwiseOffloadHook(ModelHook): weights_with_local = [] for name, t in name2weights.items(): local_t = t.to_local() if hasattr(t, "to_local") else t - weights_with_local.append((name, t, local_t)) - total_numel = sum(local.numel() for _, _, local in weights_with_local) + stride = local_t.stride() + storage_numel = ( + 0 + if local_t.numel() == 0 + else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride)) + ) + weights_with_local.append((name, t, local_t, storage_numel, stride)) + total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local) cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory) current_offset = 0 - for name, original_tensor, local_tensor in weights_with_local: - numel = local_tensor.numel() - cpu_tensor[current_offset : current_offset + numel].copy_(local_tensor.flatten()) + for ( + name, + original_tensor, + local_tensor, + storage_numel, + stride, + ) in weights_with_local: + if local_tensor.is_contiguous(): + flat_storage = local_tensor.flatten() + else: + # Cutlass FP8 weights use a transposed physical layout. + # Preserve it across the flattened CPU staging buffer. + flat_storage = torch.zeros( + storage_numel, + dtype=dtype, + device=local_tensor.device, + ) + physical_view = torch.as_strided( + flat_storage, + size=local_tensor.shape, + stride=stride, + ) + physical_view.copy_(local_tensor) + cpu_tensor[current_offset : current_offset + storage_numel].copy_(flat_storage) if dtype not in dtype_metadata: dtype_metadata[dtype] = [] dtype_metadata[dtype].append( { "name": name, "offset": current_offset, - "numel": numel, + "numel": storage_numel, "shape": local_tensor.shape, + "stride": stride, } ) LayerwiseOffloadHook._set_tensor_storage( original_tensor, LayerwiseOffloadHook._make_offload_placeholder(original_tensor) ) - current_offset += numel + current_offset += storage_numel dtype_cpu_flattened_weights[dtype] = cpu_tensor @@ -210,7 +238,11 @@ class LayerwiseOffloadHook(ModelHook): LayerwiseOffloadHook._set_tensor_storage( target_param_or_buf, - gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view(metadata["shape"]), + torch.as_strided( + gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]], + size=metadata["shape"], + stride=metadata["stride"], + ), ) self._prefetch_done = evt diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py index 8896903e8..89d445d61 100644 --- a/vllm_omni/quantization/__init__.py +++ b/vllm_omni/quantization/__init__.py @@ -11,7 +11,7 @@ Adds per-component quantization for multi-stage models. config = build_quant_config({"transformer": {"method": "fp8"}, "vae": None}) """ -from .component_config import ComponentQuantizationConfig +from .component_config import ComponentQuantizationConfig, resolve_component_quant_config from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override from .inc_config import OmniINCConfig @@ -23,6 +23,7 @@ from .inc_config import OmniINCConfig __all__ = [ "build_quant_config", "ComponentQuantizationConfig", + "resolve_component_quant_config", "OmniINCConfig", "SUPPORTED_QUANTIZATION_METHODS", "register_quantization_override", diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index 28c33d637..638419689 100644 --- a/vllm_omni/quantization/component_config.py +++ b/vllm_omni/quantization/component_config.py @@ -29,6 +29,21 @@ if TYPE_CHECKING: PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}) +def resolve_component_quant_config( + quant_config: QuantizationConfig | None, + component: str, +) -> QuantizationConfig | None: + """Resolve one pipeline component from a global or component config. + + A plain config is global and therefore applies unchanged to every + quantization-aware component. Only ``ComponentQuantizationConfig`` narrows + the scope through its explicit prefix map. + """ + if isinstance(quant_config, ComponentQuantizationConfig): + return quant_config.resolve(component) + return quant_config + + def resolve_encoder_quant_config( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: