diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index b6e67c5ca..d36afb79e 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -353,9 +353,14 @@ modes can resolve the same `None` value to different schedules. Cache-DiT on any selected component. Prompt-embedding cache is also incompatible when the text encoder uses `allgather`. Rank-local cache hits could otherwise make ranks enter different weight collectives. -- Per-tensor online FP8 linears use the ordinary loader and can run with either - DiT transfer path. With DiT `allgather`, every rank temporarily materializes - the complete FP8 model in host memory before DLO retains only its shard. +- Per-tensor online FP8, INT8, and MXFP8 linears use the ordinary loader and + can run with either DiT transfer path. With DiT `allgather`, every rank + temporarily materializes the complete quantized model in host memory before + DLO retains only its shard. Layers too wide for the NPU quantization kernel + (`npu_quant_matmul` rejects an output dimension past 65535) stay + unquantized and load straight into host memory under offload-after-quant — + their runtime layout is a plain contiguous bf16 weight, the same layout DLO + shards on the ordinary path, so they are allowed under `allgather` too. Other online quantization methods require rank-local transfer for the affected component until their runtime layouts are validated. - Resident leading layers require DiT `rank-local` transfer and a model diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index c7a87ac34..8a2ee0900 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -1241,6 +1241,169 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch): loader.load_model(load_device="cpu") +def test_dlo_allgather_allows_unquantized_host_fallback(): + """Host-loaded unquantized fallback layers are plain contiguous bf16 — + the same runtime layout DLO already shards on the ordinary path — so the + allowlist must not reject them.""" + from vllm_omni.quantization.int8_config import UnquantizedHostLinearMethod + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = object.__new__(UnquantizedHostLinearMethod) + + assert DiffusersPipelineLoader._unsupported_dlo_allgather_online_quant_methods(model) == () + + +def test_dlo_load_model_keeps_host_fallback_on_cpu_through_post_load_sweep(monkeypatch, mocker): + """Through load_model(): DLO + online quant must build the model inside + load_unquantizable_fallback_on_cpu(), and the over-wide fallback weight + must survive the post-load sweep (the _process_weights_after_loading pass + plus model.to("cpu")) as the same host tensor it was loaded into — never + bounced to the accelerator.""" + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.quantization import int8_config + from vllm_omni.quantization.int8_config import ( + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + DiffusionInt8Config, + NPUInt8OnlineLinearMethod, + UnquantizedHostLinearMethod, + ) + + # Same stand-in as TestHostFallbackLoading: the eager fallback delegates to + # UnquantizedLinearMethod, which reads the TP group. + mock_group = mocker.Mock() + mock_group.rank_in_group = 0 + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + out_features = NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1 + built_with_ctx: list[bool] = [] + + def copy_loader(param, loaded_weight, *args, **kwargs): + param.data.copy_(loaded_weight) + + def build_model(*_args, **_kwargs): + built_with_ctx.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get()) + quant_config = DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic") + method = NPUInt8OnlineLinearMethod(quant_config) + layer = nn.Module() + layer.quant_method = method + method.create_weights( + layer, + input_size_per_partition=8, + output_partition_sizes=[out_features], + input_size=8, + output_size=out_features, + params_dtype=torch.bfloat16, + weight_loader=copy_loader, + ) + model = nn.Module() + model.transformer = layer + return model + + def load_weights(_model, **_kwargs): + layer = _model.transformer + loaded = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16) + layer.weight.weight_loader(layer.weight, loaded) + + loader._init_from_load_format = build_model # type: ignore[method-assign] + loader.load_weights = load_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + # Diffusion DiT models have no vLLM Attention layers, so the upstream + # layerwise finalize is a no-op here; stub it to keep the real + # _process_weights_after_loading sweep CPU-only. + monkeypatch.setattr( + "vllm.model_executor.model_loader.reload.layerwise.finalize_layerwise_processing", + lambda *_args, **_kwargs: None, + ) + + model = loader.load_model(load_device="cpu", device=torch.device("cpu")) + + # Construction ran inside load_unquantizable_fallback_on_cpu(), so the + # over-wide layer took the host-loading fallback... + assert built_with_ctx == [True] + assert type(model.transformer.quant_method) is UnquantizedHostLinearMethod + # ...and the real post-load sweep skipped it via the fully-loaded flag and + # model.to("cpu") left the loaded host tensor untouched. + assert model.transformer._already_called_process_weights_after_loading + assert model.transformer.weight.device.type == "cpu" + expected = torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16) + assert torch.equal(model.transformer.weight, expected) + + +@pytest.mark.parametrize(("offload_after_quant", "ctx_expected"), [(True, True), (False, False)]) +def test_hsdp_enters_host_fallback_context_only_when_offloading_after_quant(mocker, offload_after_quant, ctx_expected): + """The HSDP path must apply the same host-fallback bound as the ordinary + path: with a quant config it initializes on the accelerator, so the + load_unquantizable_fallback_on_cpu() context is what keeps over-wide + fallback weights off the device before sharding.""" + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + from vllm_omni.quantization import int8_config + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=2, + ), + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.quant_config = object() + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + ctx_seen: list[bool] = [] + + def build_model(*_args, **_kwargs): + ctx_seen.append(int8_config._LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get()) + return model + + loader._init_from_load_format = build_model # type: ignore[method-assign] + loader.load_weights = lambda _model: None # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign] + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch( + "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp", + side_effect=lambda _model: None, + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: None, + ) + + loader._load_model_with_hsdp(torch.device("cpu"), offload_after_quant=offload_after_quant) + + assert ctx_seen == [ctx_expected] + + def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch): mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config") NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod diff --git a/tests/diffusion/quantization/test_int8_config.py b/tests/diffusion/quantization/test_int8_config.py index f706cc68b..6eba09c21 100644 --- a/tests/diffusion/quantization/test_int8_config.py +++ b/tests/diffusion/quantization/test_int8_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for Int8 quantization config.""" import pytest @@ -12,7 +12,7 @@ from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization import build_quant_config from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.") @@ -248,6 +248,124 @@ class TestNPUQuantMatmulShapeFallback: assert layer.quant_method is method +class TestHostFallbackLoading: + """Over-wide fallback layers load straight into host memory when the loader + will offload the model back to CPU after online quantization (DLO). + + Without the context the eager path is preserved: the fallback weight is a + real tensor created at construction time. + """ + + @pytest.fixture(autouse=True) + def _mock_tp(self, mocker): + # Same stand-in as TestNPUQuantMatmulShapeFallback: the eager fallback + # delegates to UnquantizedLinearMethod, which reads the TP group. + mock_group = mocker.Mock() + mock_group.rank_in_group = 0 + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + + @pytest.fixture + def quant_config(self): + from vllm_omni.quantization.int8_config import DiffusionInt8Config + + return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic") + + def _make_over_wide_layer(self, quant_config, weight_loader): + from vllm_omni.quantization.int8_config import ( + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + NPUInt8OnlineLinearMethod, + ) + + out_features = NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1 + method = NPUInt8OnlineLinearMethod(quant_config) + layer = Module() + layer.quant_method = method + method.create_weights( + layer, + input_size_per_partition=8, + output_partition_sizes=[out_features], + input_size=8, + output_size=out_features, + params_dtype=torch.bfloat16, + weight_loader=weight_loader, + ) + return layer, out_features + + @staticmethod + def _loaded_values(out_features): + return torch.arange(out_features * 8, dtype=torch.float32).reshape(out_features, 8).to(torch.bfloat16) + + def test_fallback_without_context_keeps_eager_weights(self, quant_config): + layer, _ = self._make_over_wide_layer(quant_config, lambda param, loaded_weight, *a, **k: None) + + assert type(layer.quant_method) is UnquantizedLinearMethod + assert layer.weight.device.type != "meta" + + def test_fallback_under_context_defers_weight_to_meta(self, quant_config): + from vllm_omni.quantization.int8_config import ( + UnquantizedHostLinearMethod, + load_unquantizable_fallback_on_cpu, + ) + + with load_unquantizable_fallback_on_cpu(): + layer, _ = self._make_over_wide_layer(quant_config, lambda param, loaded_weight, *a, **k: None) + + assert isinstance(layer.quant_method, UnquantizedHostLinearMethod) + assert layer.quant_method.uses_meta_device + assert layer.weight.device.type == "meta" + + def test_fallback_weight_loads_straight_to_cpu(self, quant_config): + from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu + + def copy_loader(param, loaded_weight, *args, **kwargs): + param.data.copy_(loaded_weight) + + with load_unquantizable_fallback_on_cpu(): + layer, out_features = self._make_over_wide_layer(quant_config, copy_loader) + + loaded = self._loaded_values(out_features) + layer.weight.weight_loader(layer.weight, loaded) + + assert layer.weight.device.type == "cpu" + assert torch.equal(layer.weight, loaded) + # Marks the post-load sweep to skip the layer instead of bouncing the + # host tensor to the accelerator and back for a no-op process call. + assert layer._already_called_process_weights_after_loading + + def test_chunked_load_marks_completion_only_when_full(self, quant_config): + from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu + + state = {"offset": 0} + + def half_loader(param, loaded_weight, *args, **kwargs): + start = state["offset"] + param.data[start : start + loaded_weight.shape[0]].copy_(loaded_weight) + state["offset"] += loaded_weight.shape[0] + + with load_unquantizable_fallback_on_cpu(): + layer, out_features = self._make_over_wide_layer(quant_config, half_loader) + + loaded = self._loaded_values(out_features) + halves = loaded.chunk(2, dim=0) + layer.weight.weight_loader(layer.weight, halves[0]) + assert not getattr(layer, "_already_called_process_weights_after_loading", False) + + layer.weight.weight_loader(layer.weight, halves[1]) + assert layer._already_called_process_weights_after_loading + assert torch.equal(layer.weight, loaded) + + def test_context_does_not_leak(self, quant_config): + from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu + + with load_unquantizable_fallback_on_cpu(): + pass + layer, _ = self._make_over_wide_layer(quant_config, lambda param, loaded_weight, *a, **k: None) + + assert type(layer.quant_method) is UnquantizedLinearMethod + + class TestOffloadAfterQuant: def test_only_methods_advertising_the_capability_are_asked(self): from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader diff --git a/tests/diffusion/quantization/test_mxfp8_config.py b/tests/diffusion/quantization/test_mxfp8_config.py index cc34fb232..7c3745d4b 100644 --- a/tests/diffusion/quantization/test_mxfp8_config.py +++ b/tests/diffusion/quantization/test_mxfp8_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for MXFP8 quantization config and linear method dispatch. Coverage: @@ -9,11 +9,13 @@ Coverage: - Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU) - build_quant_config integration - MXFP8_QUANT_CONFIG structure as the auto-detection contract +- Offload-after-quant capability of the shared online-method mixin (CPU) """ import pytest import torch from pytest_mock import MockerFixture +from torch.nn import Module from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm_omni.platforms import current_omni_platform @@ -335,3 +337,99 @@ def test_supported_methods_include_mxfp8(): from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS + + +# --------------------------------------------------------------------------- +# Offload-after-quant capability (CPU, no NPU ops) +# --------------------------------------------------------------------------- + + +class TestOffloadAfterQuant: + """Online MXFP methods advertise ``supports_offload_after_quant`` so the + diffusion loader can stream each layer back to host memory right after it + is quantized, instead of accumulating the whole model on the accelerator + until the final ``model.to("cpu")``. + + The capability lives on ``_LazyWeightMixin``, which is shared by all + online MXFP8/MXFP4 methods, so pinning it here covers mxfp4 too. + """ + + @pytest.fixture(autouse=True) + def _mock_tp(self, mocker: MockerFixture): + # create_weights resolves the TP group while computing per-partition + # shapes; these tests never initialize torch.distributed. + mock_group = mocker.Mock() + mock_group.rank_in_group = 0 + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + + def _make_method(self): + from vllm_omni.quantization.mxfp8_config import ( + DiffusionMXFP8Config, + NPUMxfp8OnlineLinearMethod, + ) + + return NPUMxfp8OnlineLinearMethod(DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)) + + def test_only_methods_advertising_the_capability_are_asked(self): + from vllm_omni.diffusion.model_loader.diffusers_loader import ( + DiffusersPipelineLoader, + ) + + class MetaDeviceOnlyMethod: + """Stands in for upstream online FP8: lazy, but no offload hook.""" + + uses_meta_device = True + + model = torch.nn.Sequential(Module(), Module()) + lazy_mxfp8 = self._make_method() + model[0].quant_method = lazy_mxfp8 + model[1].quant_method = MetaDeviceOnlyMethod() + + marked = DiffusersPipelineLoader._request_offload_after_quant(model) + + assert marked == 1 + assert lazy_mxfp8._offload_after_quant + + def _load_layer_in_one_chunk(self, method, layer): + method.create_weights( + layer, + input_size_per_partition=8, + output_partition_sizes=[4], + input_size=8, + output_size=4, + params_dtype=torch.bfloat16, + weight_loader=lambda param, loaded_weight, *args, **kwargs: param.copy_(loaded_weight), + ) + full_weight = torch.randn(4, 8, dtype=torch.bfloat16) + layer.weight.weight_loader(layer.weight, full_weight) + + def test_finished_layer_returns_to_cpu_when_enabled(self, mocker: MockerFixture): + method = self._make_method() + method.enable_offload_after_quant() + mocker.patch.object(method, "process_weights_after_loading") + + layer = Module() + moved_to = [] + layer.to = lambda device, *args, **kwargs: moved_to.append(device) # type: ignore[method-assign] + + self._load_layer_in_one_chunk(method, layer) + + method.process_weights_after_loading.assert_called_once_with(layer) + assert layer._already_called_process_weights_after_loading + assert moved_to == ["cpu"] + + def test_finished_layer_stays_put_by_default(self, mocker: MockerFixture): + method = self._make_method() + mocker.patch.object(method, "process_weights_after_loading") + + layer = Module() + moved_to = [] + layer.to = lambda device, *args, **kwargs: moved_to.append(device) # type: ignore[method-assign] + + self._load_layer_in_one_chunk(method, layer) + + method.process_weights_after_loading.assert_called_once_with(layer) + assert layer._already_called_process_weights_after_loading + assert moved_to == [] diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 834887f41..c47eab7f7 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -679,10 +679,21 @@ class DiffusersPipelineLoader(HWRLoaderMixin): with set_default_torch_dtype(self.od_config.dtype): if self.parallel_config.use_hsdp: model = self._load_model_with_hsdp( - target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name + target_device=device, + load_format=load_format, + custom_pipeline_name=custom_pipeline_name, + offload_after_quant=offload_after_quant, ) else: - model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) + # The model is headed back to host memory right after online + # quantization, so over-wide NPU-unquantizable fallback weights + # load straight into host memory instead of round-tripping + # through the accelerator (~24 GiB startup peak on MiniMax H3). + from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu + + fallback_ctx = load_unquantizable_fallback_on_cpu() if offload_after_quant else contextlib.nullcontext() + with fallback_ctx: + model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) resolved_offload = resolve_offload(self.od_config) distributed_offload = resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE @@ -721,8 +732,9 @@ class DiffusersPipelineLoader(HWRLoaderMixin): if unsupported_methods: raise ValueError( "DLO+AllGather supports online quantization only for " - "per-tensor FP8, INT8, and MXFP8 linears; unsupported " - f"online methods: {', '.join(sorted(unsupported_methods))}. " + "per-tensor FP8, INT8, and MXFP8 linears " + "(host-loaded unquantized fallback layers are also allowed); " + f"unsupported online methods: {', '.join(sorted(unsupported_methods))}. " "Use rank-local transfer for the affected component or " "disable online quantization." ) @@ -930,6 +942,11 @@ class DiffusersPipelineLoader(HWRLoaderMixin): scale, packing, or aliasing layouts (e.g. dual-scale fp4 pairs, swizzled or NZ hardware formats) and remain fail-closed until validated. + + ``UnquantizedHostLinearMethod`` is also allowed: it backs layers too + wide for npu_quant_matmul, loads their weights straight into host + memory, and its runtime layout is a plain contiguous bf16 weight — + identical to the ordinary unquantized path DLO already shards. """ from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, @@ -938,6 +955,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): from vllm_omni.quantization.int8_config import ( Int8OnlineLinearMethod, NPUInt8OnlineLinearMethod, + UnquantizedHostLinearMethod, ) try: @@ -959,6 +977,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): Fp8PerTensorOnlineLinearMethod, Int8OnlineLinearMethod, NPUInt8OnlineLinearMethod, + UnquantizedHostLinearMethod, *mxfp8_online_methods, ) @@ -1208,6 +1227,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): target_device: torch.device, load_format: str = "default", custom_pipeline_name: str | type[nn.Module] | None = None, + offload_after_quant: bool = False, ) -> nn.Module: """Load model with HSDP sharding for inference. @@ -1231,40 +1251,54 @@ class DiffusersPipelineLoader(HWRLoaderMixin): # mapping (QKV fusion, etc.). if load_format == "diffusers": raise ValueError("HSDP is not supported with the diffusers adapter load format") - model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) - world_size = 1 - rank = 0 - if torch.distributed.is_initialized(): - world_size = torch.distributed.get_world_size() - rank = torch.distributed.get_rank() + # Same host-fallback bound as the ordinary path: with a quant config the + # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so + # over-wide NPU-unquantizable fallback weights would otherwise round-trip + # through the device before apply_hsdp_to_model shards them. Loading them + # straight into host memory keeps the load-time device peak at the + # quantizable layers alone; sharding then distributes the fallback like + # any other parameter. Broadcast loading excludes online quantization + # already, so the context below only ever matters on the ordinary + # per-rank branch -- but spanning both is harmless. + from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu - has_online_quant = self._has_online_quant(model) or ( - self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False) - ) - enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1 + fallback_ctx = load_unquantizable_fallback_on_cpu() if offload_after_quant else contextlib.nullcontext() + with fallback_ctx: + model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) + world_size = 1 + rank = 0 + if torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() - if enable_broadcast and has_online_quant: - logger.info( - "Worker %d: Online quantization detected; falling back to ordinary per-rank weight loading for HSDP", - rank, + has_online_quant = self._has_online_quant(model) or ( + self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False) ) - enable_broadcast = False + enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1 - if enable_broadcast: - if rank == 0: + if enable_broadcast and has_online_quant: + logger.info( + "Worker %d: Online quantization detected; falling back to ordinary " + "per-rank weight loading for HSDP", + rank, + ) + enable_broadcast = False + + if enable_broadcast: + if rank == 0: + self.load_weights(model) + self._maybe_fuse_distilled_lora(model) + self._broadcast_model_weights(model, target_device=target_device, src_rank=0) + if ( + rank != 0 + and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill") + and getattr(self.od_config, "lora_path", None) + and hasattr(model, "load_lora_weights") + ): + setattr(model, "lora_is_fused", True) + else: self.load_weights(model) self._maybe_fuse_distilled_lora(model) - self._broadcast_model_weights(model, target_device=target_device, src_rank=0) - if ( - rank != 0 - and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill") - and getattr(self.od_config, "lora_path", None) - and hasattr(model, "load_lora_weights") - ): - setattr(model, "lora_is_fused", True) - else: - self.load_weights(model) - self._maybe_fuse_distilled_lora(model) # Quantization methods must finish while parameters are ordinary local # tensors. Some post-load transforms use operations (for example, diff --git a/vllm_omni/quantization/int8_config.py b/vllm_omni/quantization/int8_config.py index de31fbc3a..c7abe209f 100644 --- a/vllm_omni/quantization/int8_config.py +++ b/vllm_omni/quantization/int8_config.py @@ -7,6 +7,8 @@ on CUDA and NPU platforms. """ from collections.abc import Callable +from contextlib import contextmanager +from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Optional import torch @@ -62,6 +64,32 @@ NPU_QUANT_MATMUL_MAX_OUT_FEATURES = 65535 logger = init_logger(__name__) +# Set by the diffusion loader while it constructs a model whose weights will be +# offloaded back to host memory after online quantization (DLO). Over-wide +# layers that npu_quant_matmul cannot run then load straight into host memory +# instead of being built on the accelerator and moved off at the end of +# loading — on MiniMax H3 the 50 fallback adaln layers are ~24 GiB of bf16, +# which is the dominant startup memory peak. +_LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU: ContextVar[bool] = ContextVar( + "int8_load_unquantizable_fallback_on_cpu", default=False +) + + +@contextmanager +def load_unquantizable_fallback_on_cpu(): + """Create npu_quant_matmul-incompatible fallback weights on meta and load + them straight into host memory. + + Only valid when the whole model returns to the host after loading (the + loader's offload-after-quant path); otherwise the fallback weights would + stay on CPU while the rest of the model runs on the accelerator. + """ + token = _LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.set(True) + try: + yield + finally: + _LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.reset(token) + def _fell_back_to_unquantized_npu( layer: torch.nn.Module, @@ -90,7 +118,15 @@ def _fell_back_to_unquantized_npu( output_size_per_partition, NPU_QUANT_MATMUL_MAX_OUT_FEATURES, ) - fallback = UnquantizedLinearMethod() + if _LOAD_UNQUANTIZABLE_FALLBACK_ON_CPU.get(): + logger.info_once( + "Loading over-wide unquantized fallback weights straight into host memory " + "(offload-after-quant is active); they only reach the accelerator via " + "the offload backend's runtime prefetch." + ) + fallback: LinearMethodBase = UnquantizedHostLinearMethod() + else: + fallback = UnquantizedLinearMethod() layer.quant_method = fallback fallback.create_weights( layer, @@ -297,6 +333,14 @@ class LazyWeightMixin: """ self._offload_after_quant = True + def _lazy_load_device(self) -> torch.device: + """Device the meta weight materializes on at first load. + + Subclasses that never want the accelerator involved (e.g. the + over-wide host fallback) override this to return host memory. + """ + return torch.get_default_device() + def create_weights( self, layer: torch.nn.Module, @@ -347,7 +391,7 @@ class LazyWeightMixin: # process_weights_after_loading target_loaded_numel = layer.weight.numel() if layer._loaded_numel == target_loaded_numel: - self.process_weights_after_loading(layer) + self.process_weights_after_loading(layer) # type: ignore[attr-defined] # Prevent the usual `process_weights_after_loading` call from doing # anything @@ -378,10 +422,46 @@ class LazyWeightMixin: weight_loader=patched_weight_loader, ) # stash the correct device for `patched_weight_loader` - layer._load_device = torch.get_default_device() + layer._load_device = self._lazy_load_device() layer.register_parameter("weight", weight) +class UnquantizedHostLinearMethod(LazyWeightMixin, UnquantizedLinearMethod): + """Unquantized linear method whose weight loads straight into host memory. + + Layers wider than ``NPU_QUANT_MATMUL_MAX_OUT_FEATURES`` cannot be quantized, + so under the ordinary path their bf16 weights materialize on the accelerator + at construction time and stay there until the whole model is moved off after + loading. When the loader has entered ``load_unquantizable_fallback_on_cpu()`` + (offload-after-quant, i.e. DLO), that round trip is pure startup peak — the + weights end up pinned on the host either way and only visit the accelerator + through the offload backend's runtime prefetch. + + This method reuses ``LazyWeightMixin``'s meta + just-in-time materialization + wholesale (the weight really is deferred on meta until loading, like the + online quant methods, so the loader's meta-aware bookkeeping treats the + layer correctly) and only overrides the materialization target to host + memory: checkpoint chunks are CPU tensors, so loading stays CPU->CPU with + no accelerator round trip. + """ + + # Offload-after-quant marking is deliberately not advertised even though + # the mixin provides the hook: the weight never visits the accelerator, + # so there is nothing to return to host. + supports_offload_after_quant: bool = False + + def _lazy_load_device(self) -> torch.device: + return torch.device("cpu") + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Unquantized weights need no post-load transform. + + The mixin calls this when the weight is fully loaded, right before it + flags the layer — and the flag is what keeps the loader's post-load + sweep from bouncing the host tensor to the accelerator and back. + """ + + class Int8LinearMethod(BaseInt8LinearMethod): """ Linear method for Int8 diff --git a/vllm_omni/quantization/mxfp8_config.py b/vllm_omni/quantization/mxfp8_config.py index 69e51dc6b..28a91574b 100644 --- a/vllm_omni/quantization/mxfp8_config.py +++ b/vllm_omni/quantization/mxfp8_config.py @@ -167,6 +167,21 @@ class _LazyWeightMixin: uses_meta_device: bool = True + # This mixin knows when a layer's weight is final, so it can hand the layer + # back to the host right there. Loaders that intend to offload the whole + # model after loading opt in per layer via ``enable_offload_after_quant``. + supports_offload_after_quant: bool = True + _offload_after_quant: bool = False + + def enable_offload_after_quant(self) -> None: + """Return each layer to host memory as soon as it has been quantized. + + Caps the load-time device footprint at one layer instead of the whole + model. A quant method instance belongs to a single layer, so this is not + a global switch. + """ + self._offload_after_quant = True + def create_weights( self, layer: torch.nn.Module, @@ -202,12 +217,17 @@ class _LazyWeightMixin: param = layer.weight counter = CopyNumelCounter() with counter: - res = weight_loader(param, loaded_weight, *args, **kwargs) + res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc] layer._loaded_numel += counter.copied_numel if layer._loaded_numel == layer.weight.numel(): - self.process_weights_after_loading(layer) + self.process_weights_after_loading(layer) # type: ignore[attr-defined] layer._already_called_process_weights_after_loading = True + + # This layer's weight is final, so nothing needs it on the + # accelerator until inference. + if self._offload_after_quant: + layer.to("cpu") return res weight = ModelWeightParameter(