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

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

docs/design/feature/hsdp.md

@@ -0,0 +1,146 @@
+# HSDP
+
+This section describes how to add HSDP (Hybrid Sharded Data Parallel) support to a diffusion transformer model. We use the Wan2.2 transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+
+---
+
+## Overview
+
+### What is HSDP?
+
+HSDP (Hybrid Sharded Data Parallel) is a memory optimization technique that **shards model weights** across multiple GPUs using PyTorch's FSDP2. Unlike Tensor Parallelism which splits computation, HSDP:
+
+- Shards weights across GPUs to reduce per-GPU memory usage
+- Gathers weights on-demand during forward passes
+- Can work standalone or combined with other parallelism (e.g., Sequence Parallel)
+
+This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+**Important constraints:**
+- HSDP cannot be used with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+
+### Architecture
+
+HSDP implementation relies on:
+
+1. **`_hsdp_shard_conditions`**: Model attribute specifying which modules to shard
+2. **`apply_hsdp_to_model`**: Function that applies FSDP2 sharding
+3. **`HSDPInferenceConfig`**: Runtime configuration for HSDP
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Identify Modules to Shard
+
+Determine which modules in your transformer should be sharded. Typically, these are:
+
+- Transformer blocks (e.g., `blocks.0`, `blocks.1`, ...)
+- Large submodules with significant weight memory
+
+**Key questions:**
+- Which modules have the largest weights?
+- Which modules are repeated (like transformer blocks)?
+
+### Step 2: Define Shard Conditions
+
+Add `_hsdp_shard_conditions` to your model class. This is a list of functions that return `True` for modules that should be sharded.
+
+**Example (Transformer Blocks):**
+
+```python
+class MyTransformerModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        """Match transformer blocks for HSDP sharding.
+
+        Args:
+            name: Module name from named_modules() (e.g., "blocks.0", "blocks.0.attn")
+            module: The module instance
+
+        Returns:
+            True if this module should be sharded
+        """
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Multiple Conditions Example:**
+
+```python
+class MyModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name: str, module) -> bool:
+        # Also shard MoE expert layers
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    # Module is sharded if ANY condition returns True
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+---
+
+## Testing
+
+After adding HSDP support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+parallel_config = DiffusionParallelConfig(
+    use_hsdp=True,
+    hsdp_shard_size=8,  # Shard across 8 GPUs
+)
+omni = Omni(model="your-model-name", parallel_config=parallel_config)
+
+output = omni.generate(
+    "a cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Or via command line:**
+
+```bash
+vllm serve Your-org/your-model --omni --port 8091 --use-hsdp
+```
+
+**Verify:**
+
+1. Check logs for "HSDP Inference: replicate_size=..., shard_size=..."
+2. Check logs for "Sharded N modules + root"
+3. Verify memory usage is reduced proportionally
+4. Compare generated output quality with HSDP disabled
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference implementation |
+| **HSDP Core** | `vllm_omni/diffusion/distributed/hsdp.py` | `apply_hsdp_to_model`, `shard_model` |
+| **HSDP Tests** | `tests/diffusion/distributed/test_hsdp.py` | Unit tests |
+
+---

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -16,6 +16,8 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
 
+6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
+
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
@@ -49,9 +51,9 @@ The following table shows which models are currently supported by parallelism me
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel |
-|-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
+| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
+|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
 
 ### Tensor Parallelism
 
@@ -330,3 +332,84 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa
 ```bash
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
 ```
+
+### HSDP
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
+
+#### Configuration
+
+HSDP is configured via `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `use_hsdp` | bool | False | Enable HSDP |
+| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
+| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
+
+**Constraints:**
+
+- `hsdp_replicate_size × hsdp_shard_size == world_size`
+- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+
+#### Operating Modes
+
+HSDP can work in two modes:
+
+- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
+- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
+
+#### Offline Inference
+
+**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_shard_size=4,  # Shard across 4 GPUs
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Combined HSDP + Sequence Parallel**:
+
+```python
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,  # Sequence parallel
+        use_hsdp=True,     # HSDP overlays on SP
+    ),
+)
+```
+
+#### Online Serving
+
+**Standalone HSDP** (shard model across 4 GPUs):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
+```
+
+**Combined with Sequence Parallel**:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
+```
+
+#### Adding HSDP Support to New Models
+
+For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).

docs/user_guide/diffusion_acceleration.md

@@ -30,6 +30,8 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 
 5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
 
+6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
+
 ## Quick Comparison
 
 ### Cache Methods
@@ -69,9 +71,9 @@ The following table shows which models are currently supported by each accelerat
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
 
 ### Quantization
 
@@ -265,6 +267,30 @@ outputs = omni.generate(
 )
 ```
 
+### Using HSDP
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_replicate_size=1,  # Number of replica groups
+        hsdp_shard_size=8,      # Number of GPUs to shard across (or -1 for auto)
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
 ### Using CFG-Parallel
 
 Run image-to-image:
@@ -319,3 +345,4 @@ For detailed information on each acceleration method:
 - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
 - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
 - **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
+- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.

examples/offline_inference/image_to_video/README.md

@@ -66,5 +66,10 @@ Key arguments:
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
+- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
+- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
+
+
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -80,6 +80,24 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--ulysses-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism.",
+    )
+    parser.add_argument(
+        "--ring-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ring sequence parallelism.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
@@ -92,6 +110,29 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--use-hsdp",
+        action="store_true",
+        help=("Enable Hybrid Sharded Data Parallel to shard model weights across GPUs. "),
+    )
+    parser.add_argument(
+        "--hsdp-shard-size",
+        type=int,
+        default=-1,
+        help=(
+            "Number of GPUs to shard model weights across within each replica group. "
+            "-1 (default) auto-calculates as world_size / replicate_size. "
+        ),
+    )
+    parser.add_argument(
+        "--hsdp-replicate-size",
+        type=int,
+        default=1,
+        help=(
+            "Number of replica groups for HSDP. Each replica holds a full sharded copy. "
+            "Default 1 means pure sharding (no replication). "
+        ),
+    )
     return parser.parse_args()
 
 
@@ -128,7 +169,13 @@ def main():
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
     )
     omni = Omni(
         model=args.model,

tests/diffusion/distributed/test_hsdp.py

@@ -0,0 +1,225 @@
+# 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.
+
+These tests verify HSDP configuration logic without requiring a distributed environment.
+"""
+
+import pytest
+import torch.nn as nn
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu]
+
+
+class TestHSDPInferenceConfig:
+    """Tests for HSDPInferenceConfig dataclass."""
+
+    def test_default_values(self):
+        """Test default configuration values."""
+        config = HSDPInferenceConfig()
+        assert config.enabled is False
+        assert config.hsdp_replicate_size == 1
+        assert config.hsdp_shard_size == -1
+
+    def test_custom_values(self):
+        """Test custom configuration values."""
+        config = HSDPInferenceConfig(
+            enabled=True,
+            hsdp_replicate_size=2,
+            hsdp_shard_size=4,
+        )
+        assert config.enabled is True
+        assert config.hsdp_replicate_size == 2
+        assert config.hsdp_shard_size == 4
+
+
+class TestDiffusionParallelConfigHSDP:
+    """Tests for HSDP settings in DiffusionParallelConfig."""
+
+    def test_hsdp_disabled_by_default(self):
+        """HSDP should be disabled by default."""
+        config = DiffusionParallelConfig()
+        assert config.use_hsdp is False
+        assert config.hsdp_shard_size == -1
+        assert config.hsdp_replicate_size == 1
+
+    def test_hsdp_auto_shard_size(self):
+        """Test auto-calculation of hsdp_shard_size when use_hsdp=True."""
+        # ulysses_degree=4 -> world_size=4
+        # hsdp_shard_size should be auto-calculated as 4 // 1 = 4
+        config = DiffusionParallelConfig(
+            ulysses_degree=4,
+            use_hsdp=True,
+        )
+        assert config.world_size == 4
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 1
+
+    def test_hsdp_auto_shard_size_fails_standalone(self):
+        """Test that auto-calculate fails when other parallelism is all 1."""
+        # When all other parallelism is 1, cannot auto-calculate
+        # User must specify hsdp_shard_size explicitly
+        with pytest.raises(ValueError, match="Cannot auto-calculate hsdp_shard_size"):
+            DiffusionParallelConfig(
+                use_hsdp=True,
+                # All other parallelism defaults to 1
+            )
+
+    def test_hsdp_standalone_mode(self):
+        """Test standalone HSDP (HSDP without other parallelism)."""
+        # Standalone HSDP: all other parallelism=1, explicit shard_size
+        config = DiffusionParallelConfig(
+            use_hsdp=True,
+            hsdp_shard_size=4,  # Explicit shard size
+            hsdp_replicate_size=1,
+        )
+        # world_size should be determined by HSDP
+        assert config.world_size == 4
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 1
+
+    def test_hsdp_standalone_with_replicate(self):
+        """Test standalone HSDP with replication."""
+        config = DiffusionParallelConfig(
+            use_hsdp=True,
+            hsdp_shard_size=4,
+            hsdp_replicate_size=2,
+        )
+        # world_size = shard_size * replicate_size
+        assert config.world_size == 8
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 2
+
+    def test_hsdp_with_replicate(self):
+        """Test HSDP with replication (hybrid mode) combined with other parallelism."""
+        # world_size=8, replicate=2 -> shard_size should be 4
+        config = DiffusionParallelConfig(
+            ulysses_degree=8,
+            use_hsdp=True,
+            hsdp_replicate_size=2,
+        )
+        assert config.world_size == 8
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 2
+
+    def test_hsdp_explicit_shard_size_valid(self):
+        """Test explicit hsdp_shard_size that matches world_size."""
+        config = DiffusionParallelConfig(
+            ulysses_degree=4,
+            use_hsdp=True,
+            hsdp_shard_size=4,
+            hsdp_replicate_size=1,
+        )
+        assert config.hsdp_shard_size == 4
+
+    def test_hsdp_explicit_shard_size_invalid(self):
+        """Test that invalid HSDP dimensions raise an error when combined with other parallelism."""
+        with pytest.raises(ValueError, match="HSDP dimensions"):
+            DiffusionParallelConfig(
+                ulysses_degree=4,  # world_size=4
+                use_hsdp=True,
+                hsdp_shard_size=3,  # 1 * 3 != 4
+                hsdp_replicate_size=1,
+            )
+
+    def test_hsdp_replicate_size_exceeds_world_size(self):
+        """Test that replicate_size > world_size raises an error."""
+        with pytest.raises(ValueError, match="replicate_size.*must evenly divide world_size"):
+            DiffusionParallelConfig(
+                ulysses_degree=4,  # world_size=4
+                use_hsdp=True,
+                hsdp_replicate_size=8,  # 8 > 4, invalid
+            )
+
+    def test_hsdp_combined_world_size(self):
+        """Test that combined HSDP matches other parallelism world_size."""
+        config_no_hsdp = DiffusionParallelConfig(ulysses_degree=4)
+        config_with_hsdp = DiffusionParallelConfig(
+            ulysses_degree=4,
+            use_hsdp=True,
+            hsdp_shard_size=4,
+        )
+        # When combined with other parallelism, world_size should match
+        assert config_no_hsdp.world_size == config_with_hsdp.world_size == 4
+
+    def test_hsdp_standalone_world_size(self):
+        """Test that standalone HSDP determines world_size."""
+        config_hsdp = DiffusionParallelConfig(
+            use_hsdp=True,
+            hsdp_shard_size=8,
+        )
+        # Standalone HSDP: world_size is determined by HSDP
+        assert config_hsdp.world_size == 8
+
+    def test_hsdp_cannot_use_with_tp(self):
+        """Test that HSDP and Tensor Parallelism cannot be used together."""
+        with pytest.raises(ValueError, match="cannot be used with Tensor Parallelism"):
+            DiffusionParallelConfig(
+                tensor_parallel_size=2,
+                use_hsdp=True,
+                hsdp_shard_size=4,
+            )
+
+    def test_from_dict_with_hsdp(self):
+        """Test creating config from dict with HSDP settings."""
+        config = DiffusionParallelConfig.from_dict(
+            {
+                "ulysses_degree": 4,
+                "use_hsdp": True,
+                "hsdp_replicate_size": 2,
+            }
+        )
+        assert config.use_hsdp is True
+        assert config.hsdp_replicate_size == 2
+        assert config.hsdp_shard_size == 2  # auto: 4 // 2
+
+
+class TestHSDPShardConditions:
+    """Tests for _hsdp_shard_conditions matching logic."""
+
+    @staticmethod
+    def _is_transformer_block(name: str, module: nn.Module) -> bool:
+        """Example shard condition matching transformer blocks."""
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    def test_condition_matches_blocks(self):
+        """Test that condition matches transformer block patterns."""
+        cond = self._is_transformer_block
+        # Should match
+        assert cond("blocks.0", nn.Linear(10, 10)) is True
+        assert cond("blocks.15", nn.Linear(10, 10)) is True
+        assert cond("transformer.blocks.0", nn.Linear(10, 10)) is True
+        # Should not match
+        assert cond("blocks", nn.Linear(10, 10)) is False
+        assert cond("blocks.norm", nn.Linear(10, 10)) is False
+        assert cond("embeddings", nn.Linear(10, 10)) is False
+
+    def test_model_with_shard_conditions(self):
+        """Test model with _hsdp_shard_conditions attribute."""
+
+        class MockModel(nn.Module):
+            @staticmethod
+            def _is_block(name: str, module: nn.Module) -> bool:
+                return name.startswith("blocks.") and name.split(".")[-1].isdigit()
+
+            _hsdp_shard_conditions = [_is_block]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(10, 10) for _ in range(2)])
+
+        model = MockModel()
+        conditions = getattr(model, "_hsdp_shard_conditions", None)
+        assert conditions is not None
+        assert len(conditions) == 1
+
+        # Verify conditions work on actual model modules
+        matched = []
+        for name, module in model.named_modules():
+            if any(cond(name, module) for cond in conditions):
+                matched.append(name)
+        assert "blocks.0" in matched
+        assert "blocks.1" in matched

tests/diffusion/test_diffusion_model_runner.py

@@ -54,6 +54,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner.od_config = SimpleNamespace(
         cache_backend=cache_backend_name,
         enable_cache_dit_summary=enable_cache_dit_summary,
+        parallel_config=SimpleNamespace(use_hsdp=False),
     )
     runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
     return runner

vllm_omni/diffusion/data.py

@@ -52,6 +52,15 @@ class DiffusionParallelConfig:
     vae_patch_parallel_size: int = 1
     """Number of ranks used for VAE patch/tile parallelism (decode/encode)."""
 
+    use_hsdp: bool = False
+    """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding."""
+
+    hsdp_shard_size: int = -1
+    """Number of GPUs to shard weights across within each replica group. -1 means auto-calculate."""
+
+    hsdp_replicate_size: int = 1
+    """Number of replica groups for HSDP. Each replica holds a full sharded copy."""
+
     @model_validator(mode="after")
     def _validate_parallel_config(self) -> Self:
         """Validates the config relationships among the parallel strategies."""
@@ -68,13 +77,19 @@ class DiffusionParallelConfig:
             "Sequence parallel size must be equal to the product of ulysses degree and ring degree,"
             f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}"
         )
+
+        # Validate HSDP configuration
+        if self.use_hsdp:
+            assert self.hsdp_replicate_size > 0, "HSDP replicate size must be > 0"
+            assert self.hsdp_shard_size > 0, "HSDP shard size must be > 0 (should be set in __post_init__)"
         return self
 
     def __post_init__(self) -> None:
         if self.sequence_parallel_size is None:
             self.sequence_parallel_size = self.ulysses_degree * self.ring_degree
 
-        self.world_size = (
+        # Calculate world_size from other parallelism dimensions
+        other_parallel_world_size = (
             self.pipeline_parallel_size
             * self.data_parallel_size
             * self.tensor_parallel_size
@@ -83,6 +98,51 @@ class DiffusionParallelConfig:
             * self.cfg_parallel_size
         )
 
+        # Handle HSDP configuration
+        # HSDP can work in two modes:
+        # 1. Standalone: when other parallelism is all 1, HSDP determines world_size
+        # 2. Combined: HSDP overlays on top of other parallelism
+        if self.use_hsdp:
+            if self.tensor_parallel_size > 1:
+                raise ValueError(
+                    "HSDP (use_hsdp=True) cannot be used with Tensor Parallelism "
+                    f"(tensor_parallel_size={self.tensor_parallel_size}). "
+                    "Set tensor_parallel_size=1 when using HSDP."
+                )
+            if self.hsdp_shard_size == -1:
+                # Auto-calculate: use other_parallel_world_size as shard_size
+                if self.hsdp_replicate_size <= 0:
+                    raise ValueError("hsdp_replicate_size must be > 0")
+                if other_parallel_world_size == 1:
+                    raise ValueError(
+                        "Cannot auto-calculate hsdp_shard_size when other parallelism is all 1. "
+                        "Please specify hsdp_shard_size explicitly for standalone HSDP."
+                    )
+                if other_parallel_world_size % self.hsdp_replicate_size != 0:
+                    raise ValueError(
+                        f"Invalid HSDP configuration: replicate_size ({self.hsdp_replicate_size}) "
+                        f"must evenly divide world_size ({other_parallel_world_size}) when shard_size is -1."
+                    )
+                self.hsdp_shard_size = other_parallel_world_size // self.hsdp_replicate_size
+                self.world_size = other_parallel_world_size
+            else:
+                # Explicit shard_size: HSDP can work standalone or combined
+                hsdp_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
+                if other_parallel_world_size == 1:
+                    # Standalone HSDP: world_size is determined by HSDP
+                    self.world_size = hsdp_world_size
+                else:
+                    # Combined: HSDP must match other parallelism world_size
+                    if hsdp_world_size != other_parallel_world_size:
+                        raise ValueError(
+                            f"HSDP dimensions "
+                            f"({self.hsdp_replicate_size} × {self.hsdp_shard_size} = {hsdp_world_size}) "
+                            f"must equal world_size from other parallelism ({other_parallel_world_size})"
+                        )
+                    self.world_size = other_parallel_world_size
+        else:
+            self.world_size = other_parallel_world_size
+
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
         """
@@ -280,8 +340,6 @@ class OmniDiffusionConfig:
 
     num_gpus: int | None = None
 
-    hsdp_replicate_dim: int = 1
-    hsdp_shard_dim: int = -1
     dist_timeout: int | None = None  # timeout for torch.distributed
 
     # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
@@ -301,7 +359,6 @@ class OmniDiffusionConfig:
     # Layer-wise offloading (block-level offloading) parameters
     enable_layerwise_offload: bool = False
 
-    use_fsdp_inference: bool = False
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 
     # VAE memory optimization parameters

vllm_omni/diffusion/distributed/__init__.py

@@ -2,6 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Distributed utilities for vLLM-Omni diffusion models."""
 
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    apply_hsdp_to_model,
+)
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_fs_group,
+    get_fully_shard_rank,
+    get_fully_shard_world_size,
+)
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelConfig,
     SequenceParallelInput,
@@ -35,4 +44,10 @@ __all__ = [
     "sp_shard_with_padding",
     "ShardingValidator",
     "get_sharding_validator",
+    "HSDPInferenceConfig",
+    "apply_hsdp_to_model",
+    # FS (Fully Shard) utilities
+    "get_fs_group",
+    "get_fully_shard_rank",
+    "get_fully_shard_world_size",
 ]

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)

vllm_omni/diffusion/distributed/parallel_state.py

@@ -56,6 +56,7 @@ _SP: SequenceParallelGroupCoordinator | None = None
 _PP: PipelineGroupCoordinator | None = None
 _CFG: GroupCoordinator | None = None
 _DP: GroupCoordinator | None = None
+_FS: GroupCoordinator | None = None  # Fully Sharded (HSDP shard dimension)
 _DIT: GroupCoordinator | None = None
 
 
@@ -175,7 +176,8 @@ class RankGenerator:
         pp: int,
         cfg: int,
         dp: int,
-        order: str,
+        fs: int = 1,
+        order: str = "tp-sp-pp-cfg-dp",
         rank_offset: int = 0,
     ) -> None:
         self.tp = tp
@@ -183,6 +185,7 @@ class RankGenerator:
         self.pp = pp
         self.cfg = cfg
         self.dp = dp
+        self.fs = fs
         self.rank_offset = rank_offset
         self.world_size = tp * sp * pp * cfg * dp
 
@@ -192,14 +195,19 @@ class RankGenerator:
             "pp": self.pp,
             "cfg": self.cfg,
             "dp": self.dp,
+            "fs": self.fs,
         }
         order = order.lower()
 
         for name in self.name_to_size.keys():
+            # Skip 'fs' validation - it's handled separately with independent_ranks=True
+            # and doesn't participate in the main orthogonal rank generation
+            if name == "fs":
+                continue
             if name not in order and self.name_to_size[name] != 1:
                 raise RuntimeError(
                     f"The size of ({name}) is ({self.name_to_size[name]}), "
-                    f"but you haven't specified the order ({self.order})."
+                    f"but you haven't specified the order ({order})."
                 )
             elif name not in order:
                 order = order + "-" + name
@@ -218,7 +226,7 @@ class RankGenerator:
             mask[ordered_token.index(t)] = True
         return mask
 
-    def get_ranks(self, token):
+    def get_ranks(self, token, independent_ranks: bool = False):
         """Get rank group by input token.
 
         Arguments:
@@ -227,7 +235,21 @@ class RankGenerator:
                 to obtain multiple parallel types, we can use a hyphen
                 '-' to separate them. For example, if we want to obtain
                 the TP_DP group, the token should be 'tp-dp'.
+            independent_ranks (bool):
+                If True, generate independent rank groups that divide the world
+                into groups of the specified size. Used for FS (fully shard) groups
+                which operate independently from the main parallelism hierarchy.
         """
+        if independent_ranks and token == "fs":
+            # FS groups divide world into groups of size fs
+            # e.g., world_size=8, fs=4 -> [[0,1,2,3], [4,5,6,7]]
+            ranks = []
+            num_groups = self.world_size // self.fs
+            for i in range(num_groups):
+                group = list(range(i * self.fs + self.rank_offset, (i + 1) * self.fs + self.rank_offset))
+                ranks.append(group)
+            return ranks
+
         mask = self.get_mask(self.order, token)
         ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
         if self.rank_offset > 0:
@@ -333,6 +355,22 @@ def get_data_parallel_rank():
     return get_dp_group().rank_in_group
 
 
+# FS (Fully Shard / HSDP shard dimension)
+def get_fs_group() -> GroupCoordinator:
+    assert _FS is not None, "fully shard group is not initialized"
+    return _FS
+
+
+def get_fully_shard_world_size():
+    """Return world size for the fully shard group."""
+    return get_fs_group().world_size
+
+
+def get_fully_shard_rank():
+    """Return my rank for the fully shard group."""
+    return get_fs_group().rank_in_group
+
+
 def is_dp_last_group():
     """Return True if in the last data parallel group, False otherwise."""
     return (
@@ -438,6 +476,7 @@ def init_model_parallel_group(
         "tensor",
         "sequence",
         "classifier_free_guidance",
+        "fully_shard",
     ], f"parallel_mode {parallel_mode} is not supported"
     if parallel_mode == "pipeline":
         return PipelineGroupCoordinator(
@@ -629,6 +668,7 @@ def initialize_model_parallel(
     ring_degree: int = 1,
     tensor_parallel_size: int = 1,
     pipeline_parallel_size: int = 1,
+    fully_shard_degree: int = 1,
     backend: str | None = None,
 ) -> None:
     if backend is None:
@@ -645,6 +685,7 @@ def initialize_model_parallel(
         ring_degree: number of GPUs used for ring sequence parallelism.
         tensor_parallel_size: number of GPUs used for tensor parallelism.
         pipeline_parallel_size: number of GPUs used for pipeline parallelism.
+        fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension).
         backend: distributed backend of pytorch collective comm.
 
     Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
@@ -716,7 +757,8 @@ def initialize_model_parallel(
         pipeline_parallel_size,
         cfg_parallel_size,
         data_parallel_size,
-        "tp-sp-pp-cfg-dp",
+        fs=fully_shard_degree,
+        order="tp-sp-pp-cfg-dp",
     )
     sp_group_ranks = rank_generator.get_ranks("sp")
     global _DP
@@ -772,6 +814,16 @@ def initialize_model_parallel(
         backend=backend,
         parallel_mode="tensor",
     )
+
+    global _FS
+    assert _FS is None, "fully shard group is already initialized"
+    _FS = init_model_parallel_group(
+        group_ranks=rank_generator.get_ranks("fs", independent_ranks=True),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="fully_shard",
+    )
+
     init_dit_group(dit_parallel_size, backend)
 
 
@@ -801,6 +853,11 @@ def destroy_model_parallel():
         _PP.destroy()
     _PP = None
 
+    global _FS
+    if _FS:
+        _FS.destroy()
+    _FS = None
+
 
 def destroy_distributed_environment():
     global _WORLD

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -26,6 +26,7 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
 from vllm_omni.diffusion.registry import initialize_model
 
 logger = init_logger(__name__)
@@ -219,15 +220,20 @@ class DiffusersPipelineLoader:
         """Load a model with the given configurations."""
         target_device = torch.device(load_device)
         with set_default_torch_dtype(od_config.dtype):
-            with target_device:
-                if load_format == "default":
-                    model = initialize_model(od_config)
-                elif load_format == "custom_pipeline":
-                    model_cls = resolve_obj_by_qualname(custom_pipeline_name)
-                    model = model_cls(od_config=od_config)
-            logger.debug("Loading weights on %s ...", load_device)
-            # Quantization does not happen in `load_weights` but after it
-            self.load_weights(model)
+            if od_config.parallel_config.use_hsdp:
+                model = self._load_model_with_hsdp(
+                    od_config, load_format=load_format, custom_pipeline_name=custom_pipeline_name
+                )
+            else:
+                with target_device:
+                    if load_format == "default":
+                        model = initialize_model(od_config)
+                    elif load_format == "custom_pipeline":
+                        model_cls = resolve_obj_by_qualname(custom_pipeline_name)
+                        model = model_cls(od_config=od_config)
+                logger.debug("Loading weights on %s ...", load_device)
+                # Quantization does not happen in `load_weights` but after it
+                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
@@ -278,3 +284,57 @@ class DiffusersPipelineLoader:
         #             "Following weights were not initialized from "
         #             f"checkpoint: {weights_not_loaded}"
         #         )
+
+    def _load_model_with_hsdp(
+        self,
+        od_config: OmniDiffusionConfig,
+        load_format: str = "default",
+        custom_pipeline_name: str | None = None,
+    ) -> nn.Module:
+        """Load model with HSDP sharding for inference.
+
+        The pipeline contains multiple components (text_encoder, VAE, transformer).
+        Only the transformer is sharded with HSDP. Other components are loaded normally.
+
+        Approach: Load weights first using model's load_weights (handles QKV fusion etc.),
+        then apply HSDP sharding to redistribute weights across GPUs.
+        """
+        from vllm_omni.diffusion.distributed.hsdp import apply_hsdp_to_model
+
+        parallel_config = od_config.parallel_config
+        hsdp_config = HSDPInferenceConfig(
+            enabled=True,
+            hsdp_replicate_size=parallel_config.hsdp_replicate_size,
+            hsdp_shard_size=parallel_config.hsdp_shard_size,
+            param_dtype=od_config.dtype,
+        )
+
+        # Initialize model WITHOUT device context (weights start on CPU).
+        # Unlike the non-HSDP path which uses `with target_device:` to create weights
+        # directly on GPU, HSDP needs weights on CPU first so they can be redistributed
+        # across GPUs by apply_hsdp_to_model. The model's load_weights handles weight
+        # mapping (QKV fusion, etc.).
+        if load_format == "default":
+            model = initialize_model(od_config)
+        elif load_format == "custom_pipeline":
+            model_cls = resolve_obj_by_qualname(custom_pipeline_name)
+            model = model_cls(od_config=od_config)
+        self.load_weights(model)
+
+        # Collect all transformers to shard (some models have transformer_2 for MoE)
+        transformers_to_shard = []
+        transformer = getattr(model, "transformer", None)
+        if transformer is None:
+            raise ValueError("Model has no transformer attribute for HSDP")
+        transformers_to_shard.append(("transformer", transformer))
+
+        # Check for transformer_2 (MoE two-stage models like Wan2.2-I2V)
+        transformer_2 = getattr(model, "transformer_2", None)
+        if transformer_2 is not None:
+            transformers_to_shard.append(("transformer_2", transformer_2))
+
+        # Apply HSDP sharding to all transformers
+        for name, trans in transformers_to_shard:
+            logger.debug("Applying HSDP to %s", name)
+            apply_hsdp_to_model(trans, hsdp_config)
+        return model

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -725,6 +725,13 @@ class WanTransformer3DModel(nn.Module):
         "to_qkv": ["to_q", "to_k", "to_v"],
     }
 
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        """Match transformer blocks for HSDP sharding (e.g., blocks.0, blocks.1)."""
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+
     # Sequence Parallelism for Wan (following diffusers' _cp_plan pattern)
     #
     # The _sp_plan specifies sharding/gathering at module boundaries:

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -166,7 +166,6 @@ class DiffusionModelRunner:
         """Load weights into the pipeline."""
         return self.pipeline.load_weights(weights)
 
-    @torch.inference_mode()
     def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
@@ -176,34 +175,44 @@ class DiffusionModelRunner:
 
         Returns:
             DiffusionOutput with generated results.
+
+        Note:
+            We use torch.no_grad() for HSDP because HSDP2's fully_shard requires access
+            to tensor version counters in pre_forward hooks, which inference tensors do
+            not track. For non-HSDP inference, we use torch.inference_mode() for better
+            performance.
         """
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
         if len(req.prompts) == 0:
             raise ValueError("Cannot execute model with empty request list")
 
-        # The manager handles the check for need_recv_cache internally
-        self.kv_transfer_manager.receive_kv_cache(req, target_device=getattr(self.pipeline, "device", None))
+        # Use no_grad() for HSDP compatibility, inference_mode() otherwise for better perf
+        use_hsdp = self.od_config.parallel_config.use_hsdp
+        grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
+        with grad_context:
+            # The manager handles the check for need_recv_cache internally
+            self.kv_transfer_manager.receive_kv_cache(req, target_device=getattr(self.pipeline, "device", None))
 
-        if req.sampling_params.generator is None and req.sampling_params.seed is not None:
-            if req.sampling_params.generator_device is not None:
-                gen_device = req.sampling_params.generator_device
-            elif self.device.type == "cpu":
-                gen_device = "cpu"
-            else:
-                gen_device = self.device
-            req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed)
+            if req.sampling_params.generator is None and req.sampling_params.seed is not None:
+                if req.sampling_params.generator_device is not None:
+                    gen_device = req.sampling_params.generator_device
+                elif self.device.type == "cpu":
+                    gen_device = "cpu"
+                else:
+                    gen_device = self.device
+                req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed)
 
-        # Refresh cache context if needed
-        if (
-            not getattr(req, "skip_cache_refresh", False)
-            and self.cache_backend is not None
-            and self.cache_backend.is_enabled()
-        ):
-            self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps)
+            # Refresh cache context if needed
+            if (
+                not getattr(req, "skip_cache_refresh", False)
+                and self.cache_backend is not None
+                and self.cache_backend.is_enabled()
+            ):
+                self.cache_backend.refresh(self.pipeline, req.sampling_params.num_inference_steps)
 
-        with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
-            with record_function("pipeline_forward"):
-                output = self.pipeline.forward(req)
+            with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
+                with record_function("pipeline_forward"):
+                    output = self.pipeline.forward(req)
 
             # NOTE:
             if (
@@ -214,4 +223,4 @@ class DiffusionModelRunner:
             ):
                 cache_summary(self.pipeline, details=True)
 
-        return output
+            return output

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -122,6 +122,7 @@ class DiffusionWorker:
                 ring_degree=parallel_config.ring_degree,
                 tensor_parallel_size=parallel_config.tensor_parallel_size,
                 pipeline_parallel_size=parallel_config.pipeline_parallel_size,
+                fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,
             )
             init_workspace_manager(self.device)
 

vllm_omni/entrypoints/async_omni.py

@@ -141,9 +141,20 @@ class AsyncOmni(OmniBase):
             sequence_parallel_size = kwargs.get("sequence_parallel_size")
             tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1
             cfg_parallel_size = kwargs.get("cfg_parallel_size") or 1
+            use_hsdp = kwargs.get("use_hsdp", False)
+            hsdp_shard_size = kwargs.get("hsdp_shard_size", -1)
+            hsdp_replicate_size = kwargs.get("hsdp_replicate_size", 1)
             if sequence_parallel_size is None:
                 sequence_parallel_size = ulysses_degree * ring_degree
-            num_devices = sequence_parallel_size * tensor_parallel_size * cfg_parallel_size
+
+            # Calculate num_devices: consider standalone HSDP
+            other_parallel_size = sequence_parallel_size * tensor_parallel_size * cfg_parallel_size
+            if use_hsdp and other_parallel_size == 1 and hsdp_shard_size > 0:
+                # Standalone HSDP: num_devices is determined by HSDP dimensions
+                num_devices = hsdp_shard_size * hsdp_replicate_size
+            else:
+                num_devices = other_parallel_size
+
             for i in range(1, num_devices):
                 devices += f",{i}"
             parallel_config = DiffusionParallelConfig(
@@ -154,6 +165,9 @@ class AsyncOmni(OmniBase):
                 ulysses_degree=ulysses_degree,
                 ring_degree=ring_degree,
                 cfg_parallel_size=cfg_parallel_size,
+                use_hsdp=use_hsdp,
+                hsdp_shard_size=hsdp_shard_size,
+                hsdp_replicate_size=hsdp_replicate_size,
             )
         default_stage_cfg = [
             {

vllm_omni/entrypoints/cli/serve.py

@@ -210,6 +210,27 @@ class OmniServeCommand(CLISubcommand):
             "Equivalent to setting DiffusionParallelConfig.ring_degree.",
         )
 
+        # HSDP (Hybrid Sharded Data Parallel) parameters
+        omni_config_group.add_argument(
+            "--use-hsdp",
+            dest="use_hsdp",
+            action="store_true",
+            help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models. "
+            "Shards model weights across GPUs to reduce per-GPU memory usage.",
+        )
+        omni_config_group.add_argument(
+            "--hsdp-shard-size",
+            type=int,
+            default=-1,
+            help="Number of GPUs to shard weights across. -1 = auto (world_size / replicate_size).",
+        )
+        omni_config_group.add_argument(
+            "--hsdp-replicate-size",
+            type=int,
+            default=1,
+            help="Number of replica groups for HSDP. Each group holds a full sharded copy.",
+        )
+
         # Cache optimization parameters
         omni_config_group.add_argument(
             "--cache-backend",