vllm-omni/vllm_omni/diffusion/distributed/hsdp.py

# 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)