diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index e6d782c0a..91c6784fc 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/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.""" diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index f884f3667..f370d5b49 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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( diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index 3a056d6d4..cae4356f4 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/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: diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index e1aa0fd1c..8ee172e6b 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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.