[Feature] Support HSDP for diffusion models (#1339)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/distributed/hsdp.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed
+from torch import nn
+from torch.distributed import DeviceMesh, init_device_mesh
+from torch.distributed.fsdp import (
+    MixedPrecisionPolicy,
+    fully_shard,
+)
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_fs_group,
+    get_fully_shard_rank,
+    get_fully_shard_world_size,
+    get_world_group,
+)
+from vllm_omni.platforms import current_omni_platform
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class HSDPInferenceConfig:
+    """Configuration for HSDP inference.
+
+    This is a runtime config created from DiffusionParallelConfig's HSDP settings.
+    """
+
+    enabled: bool = False
+    hsdp_replicate_size: int = 1
+    hsdp_shard_size: int = -1  # -1 = auto (shard across entire world)
+    param_dtype: torch.dtype = torch.bfloat16
+    reduce_dtype: torch.dtype = torch.float32
+    output_dtype: torch.dtype | None = None
+    reshard_after_forward: bool = True
+
+
+def _create_hsdp_mesh(
+    device_type: str,
+    replicate_size: int,
+    shard_pg: torch.distributed.ProcessGroup,
+) -> DeviceMesh:
+    """Create a 2D DeviceMesh for HSDP using an existing ProcessGroup for the shard dimension.
+
+    Args:
+        device_type: The device type (e.g., "cuda", "npu")
+        replicate_size: Number of replica groups
+        shard_pg: The ProcessGroup for the shard dimension (from FS GroupCoordinator)
+
+    Returns:
+        A 2D DeviceMesh with dimensions ("replicate", "shard")
+    """
+    shard_size = torch.distributed.get_world_size(shard_pg)
+    world_size = replicate_size * shard_size
+
+    # Build 2D mesh tensor: shape (replicate_size, shard_size)
+    # Ranks are arranged so that each row is a shard group
+    mesh_tensor = torch.arange(world_size).reshape(replicate_size, shard_size)
+
+    # Create DeviceMesh with the shard ProcessGroup
+    # For the shard dimension, we reuse the existing FS ProcessGroup
+    device_mesh = init_device_mesh(
+        device_type,
+        mesh_shape=(replicate_size, shard_size),
+        mesh_dim_names=("replicate", "shard"),
+    )
+
+    # Note: init_device_mesh creates new ProcessGroups internally.
+    # For consistency, we verify the mesh structure matches our FS group.
+    # In a future optimization, we could pass the existing ProcessGroups directly.
+    logger.debug(
+        "Created HSDP mesh: replicate_size=%d, shard_size=%d, mesh=%s",
+        replicate_size,
+        shard_size,
+        mesh_tensor.tolist(),
+    )
+
+    return device_mesh
+
+
+def apply_hsdp_to_model(
+    model: nn.Module,
+    hsdp_config: HSDPInferenceConfig,
+) -> nn.Module:
+    """
+    Apply HSDP sharding to a model that already has weights loaded.
+
+    This function redistributes the model's parameters across GPUs using HSDP.
+    The model should already have its weights loaded via the standard load_weights method.
+
+    Args:
+        model: Model instance with weights already loaded
+        hsdp_config: HSDP configuration with HSDP mesh dimensions
+
+    Returns:
+        HSDP-wrapped model ready for inference
+    """
+    if not hsdp_config.enabled:
+        raise ValueError("HSDP is not enabled in config")
+
+    # Use GroupCoordinator for distributed info
+    world_group = get_world_group()
+    fs_group = get_fs_group()
+
+    world_size = world_group.world_size
+    rank = world_group.rank_in_group
+    fs_world_size = get_fully_shard_world_size()
+    fs_rank = get_fully_shard_rank()
+
+    hsdp_replicate_size = hsdp_config.hsdp_replicate_size
+    hsdp_shard_size = hsdp_config.hsdp_shard_size
+
+    # Validate that the FS group matches the HSDP shard size
+    if fs_world_size != hsdp_shard_size:
+        raise ValueError(
+            f"FS group world_size ({fs_world_size}) does not match "
+            f"HSDP shard_size ({hsdp_shard_size}). "
+            "Ensure fully_shard_degree is set correctly in initialize_model_parallel."
+        )
+
+    logger.info(
+        "HSDP Inference: replicate_size=%d, shard_size=%d, world_size=%d, rank=%d, fs_world_size=%d, fs_rank=%d",
+        hsdp_replicate_size,
+        hsdp_shard_size,
+        world_size,
+        rank,
+        fs_world_size,
+        fs_rank,
+    )
+
+    mp_policy = MixedPrecisionPolicy(
+        param_dtype=hsdp_config.param_dtype,
+        reduce_dtype=hsdp_config.reduce_dtype,
+        output_dtype=hsdp_config.output_dtype,
+        cast_forward_inputs=False,
+    )
+
+    device_type = current_omni_platform.device_type
+
+    # Create 2D DeviceMesh for HSDP using the FS group's ProcessGroup for shard dimension
+    # The mesh shape is (replicate, shard) where:
+    # - replicate: groups of ranks that hold the same shard (for gradient all-reduce in training)
+    # - shard: groups of ranks that each hold different shards (for parameter all-gather)
+    device_mesh = _create_hsdp_mesh(
+        device_type=device_type,
+        replicate_size=hsdp_replicate_size,
+        shard_pg=fs_group.device_group,
+    )
+
+    hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
+    if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0:
+        raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
+
+    # Apply HSDP sharding, this will automatically handle weight distribution
+    shard_model(
+        model,
+        reshard_after_forward=hsdp_config.reshard_after_forward,
+        mp_policy=mp_policy,
+        mesh=device_mesh,
+        hsdp_shard_conditions=hsdp_shard_conditions,
+    )
+
+    for param in model.parameters():
+        param.requires_grad = False
+
+    logger.info("HSDP applied to model: %s", type(model).__name__)
+    return model
+
+
+def shard_model(
+    model: nn.Module,
+    *,
+    reshard_after_forward: bool = True,
+    mp_policy: MixedPrecisionPolicy | None = None,
+    mesh: DeviceMesh | None = None,
+    hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]],
+) -> None:
+    """Apply HSDP sharding to model modules based on shard conditions."""
+    hsdp_kwargs: dict[str, Any] = {
+        "reshard_after_forward": reshard_after_forward,
+        "mesh": mesh,
+        "mp_policy": mp_policy,
+    }
+
+    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)
+            num_sharded += 1
+
+    if num_sharded == 0:
+        raise ValueError("No modules were sharded")
+
+    fully_shard(model, **hsdp_kwargs)
+    logger.info("Sharded %d modules + root", num_sharded)