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

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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_world_group
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.

    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_size: int,
) -> DeviceMesh:
    """Create the FSDP2 DeviceMesh; it owns the HSDP process groups.

    Args:
        device_type: The device type (e.g., "cuda", "npu")
        replicate_size: Number of replica groups
        shard_size: Number of ranks in each FSDP shard group

    Returns:
        A 2D DeviceMesh with dimensions ("replicate", "shard")
    """
    world_size = replicate_size * shard_size
    actual_world_size = torch.distributed.get_world_size()
    if world_size != actual_world_size:
        raise ValueError(
            f"HSDP mesh dimensions ({replicate_size} x {shard_size} = {world_size}) "
            f"must equal WORLD size ({actual_world_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)

    device_mesh = init_device_mesh(
        device_type,
        mesh_shape=(replicate_size, shard_size),
        mesh_dim_names=("replicate", "shard"),
    )

    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,
    target_device: torch.device | None = None,
) -> 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
        target_device: Worker's execution device. When the model declares
            _hsdp_ignored_modules, those modules are excluded from FSDP's
            mesh-driven device placement, so the caller must specify where to
            put them. Optional only when there are no ignored modules.

    Returns:
        HSDP-wrapped model ready for inference
    """
    if not hsdp_config.enabled:
        raise ValueError("HSDP is not enabled in config")

    world_group = get_world_group()
    world_size = world_group.world_size
    rank = world_group.rank_in_group

    hsdp_replicate_size = hsdp_config.hsdp_replicate_size
    hsdp_shard_size = hsdp_config.hsdp_shard_size

    if hsdp_replicate_size * hsdp_shard_size != world_size:
        raise ValueError(
            f"HSDP mesh dimensions ({hsdp_replicate_size} x {hsdp_shard_size}) must equal WORLD size ({world_size})"
        )

    logger.info(
        "HSDP Inference (FSDP2): replicate_size=%d, shard_size=%d, world_size=%d, rank=%d",
        hsdp_replicate_size,
        hsdp_shard_size,
        world_size,
        rank,
    )

    # When the model contains FP8 parameters (online quantization), let FSDP
    # keep the original storage dtype on all-gather instead of casting to
    # hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect
    # FP8 inputs; an implicit FP8 -> bf16 cast would silently break them.
    has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters())
    preserve_parameter_dtypes = bool(getattr(model, "_hsdp_preserve_parameter_dtypes", False))
    mp_policy = MixedPrecisionPolicy(
        param_dtype=None if has_fp8_params or preserve_parameter_dtypes else 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

    # DeviceMesh is the single source of truth for FSDP2 shard/replicate groups
    # 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_size=hsdp_shard_size,
    )

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

    # Collect parameters of any modules the model wants excluded from FSDP sharding.
    # See _hsdp_ignored_modules on each model class. These params keep their
    # original dtype/storage and are not subject to MixedPrecisionPolicy on the root
    # FSDP wrap. Useful for small auxiliary modules (e.g., timestep embedders) that
    # need to stay in a higher-precision dtype than the bulk of the model.
    ignored_module_names = getattr(model, "_hsdp_ignored_modules", []) or []
    ignored_params: set[nn.Parameter] = set()
    # FSDP wraps move sharded params to mesh devices automatically, but
    # ignored modules stay wherever load_weights left them (CPU under HSDP).
    # The caller (diffusers_loader / similar) is responsible for telling us
    # which device the worker should run on, matching the convention used
    # for VAEs / encoders / resident_modules.
    if ignored_module_names and target_device is None:
        raise ValueError(
            f"Model {type(model).__name__} declares _hsdp_ignored_modules="
            f"{ignored_module_names} but apply_hsdp_to_model was called "
            "without target_device. The caller must pass target_device so "
            "the ignored modules can be placed on the worker's execution device."
        )
    for mod_name in ignored_module_names:
        try:
            sub_mod = model.get_submodule(mod_name)
        except AttributeError:
            logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name)
            continue
        sub_mod.to(target_device)
        ignored_params.update(sub_mod.parameters())
    if ignored_params:
        logger.info(
            "HSDP excluding %d parameter tensors from sharding (modules: %s, moved to %s)",
            len(ignored_params),
            ignored_module_names,
            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,
        reshard_after_forward=hsdp_config.reshard_after_forward,
        mp_policy=mp_policy,
        mesh=device_mesh,
        hsdp_shard_conditions=hsdp_shard_conditions,
        ignored_params=ignored_params if ignored_params else None,
    )

    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]],
    ignored_params: set[nn.Parameter] | None = None,
) -> None:
    """Apply HSDP sharding to model modules based on shard conditions.

    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.
    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,
        "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):
            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:
        raise ValueError("No modules were sharded")

    root_kwargs = dict(hsdp_kwargs)
    if ignored_params:
        root_kwargs["ignored_params"] = ignored_params
    fully_shard(model, **root_kwargs)
    logger.info(
        "Sharded %d modules + root (ignored_params=%d)",
        num_sharded,
        len(ignored_params) if ignored_params else 0,
    )