From 29357027b0ad8ddc66d554bfdb34b20d17b6c05a Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 9 Mar 2026 14:36:13 +0800 Subject: [PATCH] [Bugfix] Allow to enable HSDP alone (#1567) Signed-off-by: gcanlin --- tests/diffusion/distributed/test_hsdp.py | 163 +++++++++++++++++- vllm_omni/diffusion/data.py | 9 +- .../diffusion/distributed/parallel_state.py | 15 +- .../diffusion/worker/diffusion_worker.py | 1 + 4 files changed, 182 insertions(+), 6 deletions(-) diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 3309155e3..6fcaf1e81 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -156,13 +156,22 @@ class TestDiffusionParallelConfigHSDP: 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"): + with pytest.raises(ValueError, match="cannot be used with TP or DP"): DiffusionParallelConfig( tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=4, ) + def test_hsdp_cannot_use_with_dp(self): + """Test that HSDP and Data Parallelism cannot be used together.""" + with pytest.raises(ValueError, match="cannot be used with TP or DP"): + DiffusionParallelConfig( + data_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( @@ -177,6 +186,158 @@ class TestDiffusionParallelConfigHSDP: assert config.hsdp_shard_size == 2 # auto: 4 // 2 +class TestStandaloneHSDPDetection: + """Tests for standalone HSDP detection and dit_parallel_size calculation. + + These tests verify the logic used in initialize_model_parallel() to detect + standalone HSDP mode and calculate effective parallel sizes. + + Standalone HSDP is when all non-HSDP parallelism dimensions are 1. + """ + + @staticmethod + def compute_standalone_hsdp_params( + data_parallel_size: int = 1, + cfg_parallel_size: int = 1, + sequence_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + tensor_parallel_size: int = 1, + fully_shard_degree: int = 1, + hsdp_replicate_size: int = 1, + ) -> dict: + """Compute standalone HSDP detection parameters. + + This mirrors the logic in initialize_model_parallel(). + """ + dit_parallel_size = ( + data_parallel_size + * cfg_parallel_size + * sequence_parallel_size + * pipeline_parallel_size + * tensor_parallel_size + ) + + # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1 + is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1 + + # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) + if is_standalone_hsdp: + effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size + else: + effective_dit_parallel_size = dit_parallel_size + + effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size + + return { + "original_dit_parallel_size": dit_parallel_size, + "is_standalone_hsdp": is_standalone_hsdp, + "effective_dit_parallel_size": effective_dit_parallel_size, + "effective_dp_size": effective_dp_size, + } + + def test_standalone_hsdp_basic(self): + """Test basic standalone HSDP detection (shard_size=4, replicate=1).""" + result = self.compute_standalone_hsdp_params( + fully_shard_degree=4, + hsdp_replicate_size=1, + ) + assert result["original_dit_parallel_size"] == 1 + assert result["is_standalone_hsdp"] is True + assert result["effective_dit_parallel_size"] == 4 + assert result["effective_dp_size"] == 4 + + def test_standalone_hsdp_with_replicate(self): + """Test standalone HSDP with replication (shard_size=4, replicate=2).""" + result = self.compute_standalone_hsdp_params( + fully_shard_degree=4, + hsdp_replicate_size=2, + ) + assert result["original_dit_parallel_size"] == 1 + assert result["is_standalone_hsdp"] is True + assert result["effective_dit_parallel_size"] == 8 # 4 * 2 + assert result["effective_dp_size"] == 8 + + def test_combined_hsdp_sp_not_standalone(self): + """Test HSDP combined with SP is NOT detected as standalone. + + This is a regression test for the bug where the condition + `dit_parallel_size == fully_shard_degree` incorrectly matched + combined modes like SP=4 + HSDP=4. + """ + result = self.compute_standalone_hsdp_params( + sequence_parallel_size=4, + fully_shard_degree=4, + hsdp_replicate_size=1, + ) + assert result["original_dit_parallel_size"] == 4 + assert result["is_standalone_hsdp"] is False + # Should NOT override dp_size for combined mode + assert result["effective_dp_size"] == 1 # original data_parallel_size + + def test_combined_hsdp_cfg_not_standalone(self): + """Test HSDP combined with CFG is NOT detected as standalone.""" + result = self.compute_standalone_hsdp_params( + cfg_parallel_size=2, + fully_shard_degree=4, + hsdp_replicate_size=1, + ) + assert result["original_dit_parallel_size"] == 2 + assert result["is_standalone_hsdp"] is False + assert result["effective_dp_size"] == 1 + + def test_combined_hsdp_pp_not_standalone(self): + """Test HSDP combined with PP is NOT detected as standalone.""" + result = self.compute_standalone_hsdp_params( + pipeline_parallel_size=2, + fully_shard_degree=4, + hsdp_replicate_size=1, + ) + assert result["original_dit_parallel_size"] == 2 + assert result["is_standalone_hsdp"] is False + assert result["effective_dp_size"] == 1 + + def test_no_hsdp_not_standalone(self): + """Test that no HSDP (fully_shard_degree=1) is NOT standalone.""" + result = self.compute_standalone_hsdp_params( + fully_shard_degree=1, + ) + assert result["original_dit_parallel_size"] == 1 + assert result["is_standalone_hsdp"] is False + assert result["effective_dp_size"] == 1 + + def test_combined_multiple_parallelism_not_standalone(self): + """Test HSDP combined with multiple parallelism is NOT standalone.""" + result = self.compute_standalone_hsdp_params( + sequence_parallel_size=2, + cfg_parallel_size=2, + fully_shard_degree=4, + hsdp_replicate_size=1, + ) + assert result["original_dit_parallel_size"] == 4 # 2 * 2 + assert result["is_standalone_hsdp"] is False + assert result["effective_dp_size"] == 1 + + def test_standalone_hsdp_large_shard(self): + """Test standalone HSDP with large shard size.""" + result = self.compute_standalone_hsdp_params( + fully_shard_degree=8, + hsdp_replicate_size=1, + ) + assert result["is_standalone_hsdp"] is True + assert result["effective_dit_parallel_size"] == 8 + assert result["effective_dp_size"] == 8 + + def test_standalone_hsdp_large_replicate(self): + """Test standalone HSDP with large replicate size.""" + result = self.compute_standalone_hsdp_params( + fully_shard_degree=4, + hsdp_replicate_size=4, + ) + assert result["is_standalone_hsdp"] is True + assert result["effective_dit_parallel_size"] == 16 # 4 * 4 + assert result["effective_dp_size"] == 16 + + class TestHSDPShardConditions: """Tests for _hsdp_shard_conditions matching logic.""" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0a68c2e70..6a1e44472 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -111,11 +111,12 @@ class DiffusionParallelConfig: # 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: + if self.tensor_parallel_size > 1 or self.data_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." + "HSDP (use_hsdp=True) cannot be used with TP or DP " + f"(tensor_parallel_size={self.tensor_parallel_size}, " + f"data_parallel_size={self.data_parallel_size}). " + "Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP." ) if self.hsdp_shard_size == -1: # Auto-calculate: use other_parallel_world_size as shard_size diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index fbcfe0651..4c09fd7a0 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -669,6 +669,7 @@ def initialize_model_parallel( tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, fully_shard_degree: int = 1, + hsdp_replicate_size: int = 1, backend: str | None = None, ) -> None: if backend is None: @@ -740,6 +741,14 @@ def initialize_model_parallel( data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size ) + # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1 + is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1 + + # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) as dit_parallel_size + # This ensures orthogonal rank generation works correctly for all HSDP workers + if is_standalone_hsdp: + dit_parallel_size = fully_shard_degree * hsdp_replicate_size + if world_size < dit_parallel_size: raise RuntimeError( f"world_size ({world_size}) is less than " @@ -751,12 +760,16 @@ def initialize_model_parallel( f"data_parallel_size ({data_parallel_size})" ) + # For standalone HSDP, use (fully_shard_degree * hsdp_replicate_size) as data_parallel_size + # so that RankGenerator.world_size matches the actual number of workers + effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size + rank_generator: RankGenerator = RankGenerator( tensor_parallel_size, sequence_parallel_size, pipeline_parallel_size, cfg_parallel_size, - data_parallel_size, + effective_dp_size, fs=fully_shard_degree, order="tp-sp-pp-cfg-dp", ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 3fd29636e..2f073b368 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -125,6 +125,7 @@ class DiffusionWorker: 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, + hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1, ) init_workspace_manager(self.device)