diff --git a/docs/design/feature/hsdp.md b/docs/design/feature/hsdp.md new file mode 100644 index 000000000..958ad192e --- /dev/null +++ b/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 | + +--- diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index f8d00e2df..8002f97c9 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/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). diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 4a7ba65ee..9fd14a2b0 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/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. diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index 6a547df9d..2643084f1 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/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. diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 5a91c4c73..a491ddfc3 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/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, diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py new file mode 100644 index 000000000..3309155e3 --- /dev/null +++ b/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 diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index d8ecfb492..6ff74154f 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index cbfa1b1bc..703dba258 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/distributed/__init__.py b/vllm_omni/diffusion/distributed/__init__.py index c791ae1be..668ace06d 100644 --- a/vllm_omni/diffusion/distributed/__init__.py +++ b/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", ] diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py new file mode 100644 index 000000000..69506c928 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 26112af46..fbcfe0651 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/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 diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7806789b1..4f9a37c7a 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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 diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 5efdbc185..231a001f7 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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: diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 5fee65645..a9fd48261 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 1f6f2c04f..d4389f4d2 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 79dc7a8b7..19c8b07ac 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 = [ { diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 74a4ff7f0..57e20e409 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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",