[Bugfix][Quantization] Support packed parameters with HSDP (#5088)

Signed-off-by: Alex Steiner <asteiner@nvidia.com>

tests/diffusion/distributed/test_hsdp.py

@@ -1,19 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.
+"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
-These tests verify HSDP configuration logic without requiring a distributed environment.
-"""
+import gc
+import os
+import socket
 
 import pytest
+import torch
+import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.tensor import DeviceMesh, DTensor
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    _unshardable_parameters,
+    shard_model,
+)
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
 
+class _PackedBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 2))
+        self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
+        self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)
+
+
+class _PackedModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block = _PackedBlock()
+        self.root_weight = nn.Parameter(torch.ones(2))
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield
+        return
+
+    master_port = _find_free_port()
+    os.environ.update(
+        {
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        dist.destroy_process_group()
+        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+            os.environ.pop(key, None)
+        gc.collect()
+
+
+def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
+    model = _PackedModel()
+    ignored_params = _unshardable_parameters(model)
+    expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
+    assert ignored_params == expected_ignored
+
+    packed_weight = model.block.packed_weight
+    input_global_scale = model.block.input_global_scale
+    shard_model(
+        model,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "block"],
+        ignored_params=ignored_params,
+    )
+
+    assert isinstance(model.block.weight, DTensor)
+    assert isinstance(model.root_weight, DTensor)
+    assert model.block.packed_weight is packed_weight
+    assert model.block.input_global_scale is input_global_scale
+    assert not isinstance(model.block.packed_weight, DTensor)
+    assert not isinstance(model.block.input_global_scale, DTensor)
+
+
 class TestHSDPInferenceConfig:
     """Tests for HSDPInferenceConfig dataclass."""
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -179,6 +179,59 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert get_current_diffusion_config_or_none() is None
 
 
+def test_hsdp_processes_quantized_weights_before_sharding(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    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)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    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: events.append("prepare"),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    assert events == ["load", "process", "prepare", "shard"]
+
+
 def test_get_all_weights(prefetch_helios_model, mock_tp_group):
     """Ensure that get all weights on a tiny model resolves to nonempty weights."""
     od_config = OmniDiffusionConfig(

vllm_omni/diffusion/distributed/hsdp.py

@@ -26,6 +26,15 @@ from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 
 
+def _unshardable_parameters(model: nn.Module) -> set[nn.Parameter]:
+    """Return packed/integer or scalar parameters unsupported by FSDP2."""
+    return {
+        param
+        for param in model.parameters()
+        if param.ndim == 0 or not (param.is_floating_point() or param.is_complex())
+    }
+
+
 @dataclass
 class HSDPInferenceConfig:
     """Configuration for HSDP inference.
@@ -202,6 +211,32 @@ def apply_hsdp_to_model(
             target_device,
         )
 
+    # Serialized low-bit checkpoints may store packed weights in integer
+    # Parameters (for example, ModelOpt NVFP4 uses uint8) and global scales as
+    # scalar Parameters. FSDP2 cannot represent non-floating or zero-dimensional
+    # sharded parameters. Keep those tensors resident and replicated; eligible
+    # scales, biases, and other parameters remain HSDP-sharded.
+    unshardable_params = _unshardable_parameters(model)
+    if unshardable_params:
+        if target_device is None:
+            raise ValueError(
+                f"Model {type(model).__name__} has parameters that HSDP must ignore, "
+                "but apply_hsdp_to_model was called without target_device."
+            )
+        for param in unshardable_params:
+            if param.device != target_device:
+                param.data = param.data.to(target_device)
+        ignored_params.update(unshardable_params)
+        logger.info(
+            "HSDP excluding %d unshardable parameter tensors from sharding "
+            "(non_floating=%d, scalar=%d, dtypes=%s, moved to %s)",
+            len(unshardable_params),
+            sum(not (param.is_floating_point() or param.is_complex()) for param in unshardable_params),
+            sum(param.ndim == 0 for param in unshardable_params),
+            sorted({str(param.dtype) for param in unshardable_params}),
+            target_device,
+        )
+
     # Apply HSDP sharding, this will automatically handle weight distribution
     shard_model(
         model,
@@ -233,9 +268,9 @@ def shard_model(
     ignored_params (if provided) are excluded from the root fully_shard
     wrap, so they are not collected into the root flat-parameter, are not
     subject to MixedPrecisionPolicy, and retain their original dtype.
-    Per-submodule shard wraps do not receive ignored_params because the
-    ignored modules are expected to live at the root level, not inside any
-    block matched by hsdp_shard_conditions.
+    Each per-submodule wrap receives the subset of ignored_params that it owns.
+    This is required for packed integer parameters inside sharded transformer
+    blocks; the root wrap receives the full set for all remaining parameters.
     """
     hsdp_kwargs: dict[str, Any] = {
         "reshard_after_forward": reshard_after_forward,
@@ -246,7 +281,12 @@ def shard_model(
     num_sharded = 0
     for name, module in reversed(list(model.named_modules())):
         if any(cond(name, module) for cond in hsdp_shard_conditions):
-            fully_shard(module, **hsdp_kwargs)
+            module_kwargs = dict(hsdp_kwargs)
+            if ignored_params:
+                module_ignored_params = ignored_params.intersection(module.parameters())
+                if module_ignored_params:
+                    module_kwargs["ignored_params"] = module_ignored_params
+            fully_shard(module, **module_kwargs)
             num_sharded += 1
 
     if num_sharded == 0:

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -381,10 +381,9 @@ class DiffusersPipelineLoader:
                     cast(DiffusersAdapterPipeline, model).load_weights()
                 else:
                     self.load_weights(model)
-
-            # Process weights after loading for quantization (e.g., FP8 online quantization)
-            # This is needed for vLLM's quantization methods that need to transform weights
-            self._process_weights_after_loading(model, target_device)
+                # HSDP processes quantized weights before wrapping parameters as
+                # DTensors. The non-HSDP path can process them here as usual.
+                self._process_weights_after_loading(model, target_device)
 
             if offload_after_quant:
                 model.to("cpu")
@@ -629,6 +628,11 @@ class DiffusersPipelineLoader:
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
+        # Quantization methods must finish while parameters are ordinary local
+        # tensors. Some post-load transforms use operations (for example,
+        # torch.unique in ModelOpt NVFP4) that do not support DTensor inputs.
+        self._process_weights_after_loading(model, target_device)
+
         # Discover pipeline components (DiT, encoders, VAEs) via
         # ModuleDiscovery, which consults SupportsComponentDiscovery
         # when available and falls back to well-known attribute names.