From 14bd8f8772dfb2b6ccb2117d6b9078bbb2edbe66 Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Wed, 19 Aug 2026 11:23:10 +0800 Subject: [PATCH] [REFACT]Refactor diffusion parallel state (#5531) Signed-off-by: bjf-frz Co-authored-by: Hongsheng Liu --- docs/user_guide/diffusion/parallelism/hsdp.md | 3 +- tests/config/test_omni_config.py | 7 +- .../test_distributed_vae_executor.py | 8 - .../test_expert_parallel_layout.py | 36 ++- tests/diffusion/distributed/test_hsdp.py | 156 +---------- .../test_diffusion_config_propagation.py | 24 ++ .../test_async_omni_diffusion_config.py | 23 +- vllm_omni/config/omni_config.py | 17 +- vllm_omni/diffusion/data.py | 71 +++-- vllm_omni/diffusion/distributed/__init__.py | 9 - .../autoencoders/distributed_vae_executor.py | 13 +- .../distributed/group_coordinator.py | 24 ++ vllm_omni/diffusion/distributed/hsdp.py | 45 ++-- .../diffusion/distributed/parallel_state.py | 242 ++++++++---------- .../models/minimax_h3/pipeline_minimax_h3.py | 6 +- vllm_omni/diffusion/models/minimax_h3/vae.py | 4 +- vllm_omni/diffusion/offloader/base.py | 2 +- .../diffusion/worker/diffusion_worker.py | 3 +- vllm_omni/engine/async_omni_engine.py | 9 +- 19 files changed, 315 insertions(+), 387 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism/hsdp.md b/docs/user_guide/diffusion/parallelism/hsdp.md index 8a8aa5ef9..9491ed9f5 100644 --- a/docs/user_guide/diffusion/parallelism/hsdp.md +++ b/docs/user_guide/diffusion/parallelism/hsdp.md @@ -152,7 +152,8 @@ In `DiffusionParallelConfig`: **Constraints:** - `hsdp_replicate_size × hsdp_shard_size == world_size` -- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1) +- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism, + Pipeline Parallelism, or Expert Parallelism. --- diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 108d997a5..3ee1d40a1 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -733,8 +733,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding(): cfg = OmniStageDiffusionParallelConfig( - pipeline_parallel_size=2, - ulysses_degree=2, + ulysses_degree=4, use_hsdp=True, hsdp_shard_size=-1, hsdp_replicate_size=2, @@ -745,10 +744,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding(): def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp(): - with pytest.raises(ValueError, match="cannot be used with TP or DP"): + with pytest.raises(ValueError, match="not compatible with TP"): OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2) - with pytest.raises(ValueError, match="cannot be used with TP or DP"): + with pytest.raises(ValueError, match="not compatible with DP"): OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2) diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py index b2ee7c10d..c122e9105 100644 --- a/tests/diffusion/distributed/test_distributed_vae_executor.py +++ b/tests/diffusion/distributed/test_distributed_vae_executor.py @@ -74,14 +74,6 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None) -@pytest.fixture(autouse=True) -def mock_dit_group(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr( - "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group", - lambda: None, - ) - - @pytest.fixture(autouse=True) def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x]) diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py index 74f30f156..d7b66be5d 100644 --- a/tests/diffusion/distributed/test_expert_parallel_layout.py +++ b/tests/diffusion/distributed/test_expert_parallel_layout.py @@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch): monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group) - monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) - for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"): monkeypatch.setattr(parallel_state, name, None) for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) @@ -229,9 +228,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) - monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) - for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"): monkeypatch.setattr(parallel_state, name, None) for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) @@ -259,6 +257,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): ] +@pytest.mark.cpu +@pytest.mark.core_model +def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch): + """A destroyed diffusion PP group must not block the next initialization.""" + + class _DestroyableGroup: + def __init__(self) -> None: + self.destroy_calls = 0 + + def destroy(self) -> None: + self.destroy_calls += 1 + + pipeline_group = _DestroyableGroup() + for name in ("_DP", "_CFG", "_SP"): + monkeypatch.setattr(parallel_state, name, None) + monkeypatch.setattr(parallel_state, "_PP", pipeline_group) + monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None) + for name in ("_DP", "_PCP", "_TP", "_EP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False) + + parallel_state.destroy_model_parallel() + + assert pipeline_group.destroy_calls == 1 + assert parallel_state.vllm_parallel_state._PP is None + + @pytest.mark.cpu @pytest.mark.core_model def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch): @@ -295,9 +320,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch): monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) - monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) - for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"): monkeypatch.setattr(parallel_state, name, None) for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 91c6784fc..0703d19b8 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -234,7 +234,7 @@ 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 TP or DP"): + with pytest.raises(ValueError, match="not compatible with TP"): DiffusionParallelConfig( tensor_parallel_size=2, use_hsdp=True, @@ -243,7 +243,7 @@ class TestDiffusionParallelConfigHSDP: 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"): + with pytest.raises(ValueError, match="not compatible with DP"): DiffusionParallelConfig( data_parallel_size=2, use_hsdp=True, @@ -264,158 +264,6 @@ 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/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index e456e302b..dfe4a41e8 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -92,6 +92,30 @@ class TestParallelConfigPropagation: od = _roundtrip_diffusion_config(model="x", parallel_config=pc) assert od.num_gpus == 2 + def test_dp_is_inferred_from_num_gpus(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2) + od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8) + assert od.parallel_config.data_parallel_size == 2 + assert od.parallel_config.world_size == 8 + + def test_explicit_dp_is_validated_against_num_gpus(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2) + od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4) + assert od.parallel_config.data_parallel_size == 2 + + with pytest.raises(ValueError, match="does not match WORLD-derived value"): + OmniDiffusionConfig.from_kwargs( + model="x", + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2), + num_gpus=8, + ) + + def test_hsdp_does_not_infer_ordinary_dp(self): + pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4) + od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4) + assert od.parallel_config.data_parallel_size == 1 + assert od.parallel_config.world_size == 4 + class TestCreateDefaultDiffusion: """Verify engine_args structure from create_default_diffusion.""" diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1daca1e99..1d5d58011 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -58,6 +58,23 @@ def test_default_stage_devices_from_sequence_parallel(): assert stage_cfg["runtime"]["devices"] == "0,1,2,3" +def test_default_stage_devices_and_dp_from_num_gpus(): + """Resolve omitted DP before deriving devices from an explicit GPU count.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "num_gpus": 8, + "tensor_parallel_size": 2, + "ulysses_degree": 2, + } + )[0] + + parallel_config = stage_cfg["engine_args"]["parallel_config"] + assert parallel_config.data_parallel_size == 2 + assert parallel_config.world_size == 8 + assert stage_cfg["engine_args"]["num_gpus"] == 8 + assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7" + + def test_default_stage_config_uses_parallel_size_kwargs(): """Ensure default diffusion parallel_config uses CLI/API parallel sizes.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( @@ -76,8 +93,8 @@ def test_default_stage_config_uses_parallel_size_kwargs(): assert parallel_config.enable_expert_parallel is True -def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): - """Ensure nullified diffusion parallel-size kwargs fall back to defaults.""" +def test_default_stage_config_preserves_omitted_dp_for_runtime_inference(): + """Keep omitted DP unresolved until the runtime WORLD size is known.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( { "pipeline_parallel_size": None, @@ -92,7 +109,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config.pipeline_parallel_size == 1 - assert parallel_config.data_parallel_size == 1 + assert parallel_config.data_parallel_size is None assert parallel_config.tensor_parallel_size == 1 assert parallel_config.enable_expert_parallel is False assert stage_cfg["engine_args"]["enforce_eager"] is False diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 9cd7ed518..8dd9c3db6 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -600,12 +600,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): * self.cfg_parallel_size ) if self.use_hsdp: - if self.tensor_parallel_size > 1 or self.data_parallel_size > 1: - raise ValueError( - "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})" - ) + incompatible = [] + if self.tensor_parallel_size > 1: + incompatible.append("TP") + if self.data_parallel_size > 1: + incompatible.append("DP") + if self.pipeline_parallel_size > 1: + incompatible.append("PP") + if self.enable_expert_parallel: + incompatible.append("EP") + if incompatible: + raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible)) if self.hsdp_shard_size == -1: if other_parallel_world_size == 1: raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1") diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 34c8ccb96..2abd478c7 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -154,8 +154,8 @@ class DiffusionParallelConfig: pipeline_parallel_size: int = 1 """Number of pipeline parallel stages.""" - data_parallel_size: int = 1 - """Number of data parallel groups.""" + data_parallel_size: int | None = None + """Number of data parallel groups. Inferred from WORLD when omitted.""" tensor_parallel_size: int = 1 """Number of tensor parallel groups.""" @@ -243,7 +243,7 @@ class DiffusionParallelConfig: def _validate_parallel_config(self) -> Self: """Validates the config relationships among the parallel strategies.""" assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0" - assert self.data_parallel_size > 0, "Data parallel size must be > 0" + assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0" assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0" assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0" assert self.ulysses_degree > 0, "Ulysses degree must be > 0" @@ -283,10 +283,12 @@ class DiffusionParallelConfig: self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree ) - # Calculate world_size from other parallelism dimensions + # Until the runtime WORLD size is known, an omitted DP dimension means + # one replica. OmniDiffusionConfig resolves it against num_gpus below. + data_parallel_size = self.data_parallel_size or 1 other_parallel_world_size = ( self.pipeline_parallel_size - * self.data_parallel_size + * data_parallel_size * self.tensor_parallel_size * self.sequence_parallel_size * self.cfg_parallel_size @@ -297,13 +299,17 @@ 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 or self.data_parallel_size > 1: - raise ValueError( - "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." - ) + incompatible = [] + if self.tensor_parallel_size > 1: + incompatible.append("TP") + if data_parallel_size > 1: + incompatible.append("DP") + if self.pipeline_parallel_size > 1: + incompatible.append("PP") + if self.enable_expert_parallel: + incompatible.append("EP") + if incompatible: + raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible)) if self.hsdp_shard_size == -1: # Auto-calculate: use other_parallel_world_size as shard_size if self.hsdp_replicate_size <= 0: @@ -338,6 +344,42 @@ class DiffusionParallelConfig: else: self.world_size = other_parallel_world_size + def resolve_data_parallel_size(self, world_size: int) -> int: + """Resolve or validate diffusion DP against the actual WORLD size.""" + if world_size <= 0: + raise ValueError(f"WORLD size must be > 0, but got {world_size}") + + if self.use_hsdp: + if self.data_parallel_size not in (None, 1): + raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1") + expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size + if world_size != expected_world_size: + raise ValueError( + f"WORLD size ({world_size}) must equal HSDP size " + f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})" + ) + self.data_parallel_size = 1 + self.world_size = world_size + return 1 + + non_dp_size = ( + self.pipeline_parallel_size + * self.tensor_parallel_size + * self.sequence_parallel_size + * self.cfg_parallel_size + ) + if world_size % non_dp_size != 0: + raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})") + inferred_data_parallel_size = world_size // non_dp_size + if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size: + raise ValueError( + f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value " + f"({inferred_data_parallel_size})" + ) + self.data_parallel_size = inferred_data_parallel_size + self.world_size = world_size + return inferred_data_parallel_size + @classmethod def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig": """ @@ -999,10 +1041,7 @@ class OmniDiffusionConfig: else: self.num_gpus = 1 - if self.num_gpus < self.parallel_config.world_size: - raise ValueError( - f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})" - ) + self.parallel_config.resolve_data_parallel_size(self.num_gpus) if self.diffusion_compile_granularity == "full": incompatible_features = [] diff --git a/vllm_omni/diffusion/distributed/__init__.py b/vllm_omni/diffusion/distributed/__init__.py index 668ace06d..2a1159026 100644 --- a/vllm_omni/diffusion/distributed/__init__.py +++ b/vllm_omni/diffusion/distributed/__init__.py @@ -6,11 +6,6 @@ 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, @@ -46,8 +41,4 @@ __all__ = [ "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/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index d804740f9..9d84cdf53 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -9,8 +9,6 @@ import torch import torch.distributed as dist from vllm.logger import init_logger -from vllm_omni.diffusion.distributed.parallel_state import get_dit_group - logger = init_logger(__name__) @@ -45,9 +43,10 @@ class DistributedVaeExecutor: """ def __init__(self): - self.group = get_dit_group() - self.world_size = dist.get_world_size(self.group) - self.rank = dist.get_rank(self.group) + # VAE tiles are scheduled across the complete worker WORLD. + self.group = None + self.world_size = dist.get_world_size() + self.rank = dist.get_rank() self.parallel_size = 1 self.parallel_mode = "tile" @@ -189,7 +188,7 @@ class DistributedVaeMixin: if self.distributed_executor.parallel_size > pp_size: logger.warning( f"vae_patch_parallel_size={self.distributed_executor.parallel_size} " - f"is greater than dit_group={world_size};" - f" using dit_group size={world_size}" + f"is greater than WORLD={world_size};" + f" using WORLD size={world_size}" ) return True diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 14847007f..ac2059f56 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -991,6 +991,18 @@ class PipelineGroupCoordinator(GroupCoordinator): def _pipeline_isend_skip(self, tensor: torch.tensor): return torch.distributed.isend(tensor, dst=self.skip_rank, group=self.skip_device_group) + def destroy(self): + groups = [*self.device_groups, *self.cpu_groups, self.skip_device_group] + seen = {id(self.device_group), id(self.cpu_group)} + for group in groups: + if group is not None and id(group) not in seen: + torch.distributed.destroy_process_group(group) + seen.add(id(group)) + self.device_groups = [] + self.cpu_groups = [] + self.skip_device_group = None + super().destroy() + class SequenceParallelGroupCoordinator(GroupCoordinator): def __init__( @@ -1031,3 +1043,15 @@ class SequenceParallelGroupCoordinator(GroupCoordinator): self.ring_rank = torch.distributed.get_rank(self.ring_group) self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group) self.allgather_rank = torch.distributed.get_rank(self.allgather_group) + + def destroy(self): + groups = [self.ulysses_group, self.ring_group, self.allgather_group] + seen = {id(self.device_group), id(self.cpu_group)} + for group in groups: + if group is not None and id(group) not in seen: + torch.distributed.destroy_process_group(group) + seen.add(id(group)) + self.ulysses_group = None + self.ring_group = None + self.allgather_group = None + super().destroy() diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index cae4356f4..c496804bc 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/vllm_omni/diffusion/distributed/hsdp.py @@ -15,12 +15,7 @@ from torch.distributed.fsdp import ( ) 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.diffusion.distributed.parallel_state import get_world_group from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -54,36 +49,36 @@ class HSDPInferenceConfig: def _create_hsdp_mesh( device_type: str, replicate_size: int, - shard_pg: torch.distributed.ProcessGroup, + shard_size: int, ) -> DeviceMesh: - """Create a 2D DeviceMesh for HSDP using an existing ProcessGroup for the shard dimension. + """Create the FSDP2 DeviceMesh; it owns the HSDP process groups. Args: device_type: The device type (e.g., "cuda", "npu") replicate_size: Number of replica groups - shard_pg: The ProcessGroup for the shard dimension (from FS GroupCoordinator) + shard_size: Number of ranks in each FSDP shard group Returns: A 2D DeviceMesh with dimensions ("replicate", "shard") """ - shard_size = torch.distributed.get_world_size(shard_pg) world_size = replicate_size * shard_size + actual_world_size = torch.distributed.get_world_size() + if world_size != actual_world_size: + raise ValueError( + f"HSDP mesh dimensions ({replicate_size} x {shard_size} = {world_size}) " + f"must equal WORLD size ({actual_world_size})" + ) # Build 2D mesh tensor: shape (replicate_size, shard_size) # Ranks are arranged so that each row is a shard group mesh_tensor = torch.arange(world_size).reshape(replicate_size, shard_size) - # 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, @@ -119,34 +114,24 @@ def apply_hsdp_to_model( 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: + if hsdp_replicate_size * hsdp_shard_size != world_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." + f"HSDP mesh dimensions ({hsdp_replicate_size} x {hsdp_shard_size}) must equal WORLD size ({world_size})" ) logger.info( - "HSDP Inference: replicate_size=%d, shard_size=%d, world_size=%d, rank=%d, fs_world_size=%d, fs_rank=%d", + "HSDP Inference (FSDP2): replicate_size=%d, shard_size=%d, world_size=%d, rank=%d", hsdp_replicate_size, hsdp_shard_size, world_size, rank, - fs_world_size, - fs_rank, ) # When the model contains FP8 parameters (online quantization), let FSDP @@ -163,14 +148,14 @@ def apply_hsdp_to_model( device_type = current_omni_platform.device_type - # Create 2D DeviceMesh for HSDP using the FS group's ProcessGroup for shard dimension + # DeviceMesh is the single source of truth for FSDP2 shard/replicate groups # The mesh shape is (replicate, shard) where: # - replicate: groups of ranks that hold the same shard (for gradient all-reduce in training) # - shard: groups of ranks that each hold different shards (for parameter all-gather) device_mesh = _create_hsdp_mesh( device_type=device_type, replicate_size=hsdp_replicate_size, - shard_pg=fs_group.device_group, + shard_size=hsdp_shard_size, ) hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 55324ccf4..76d458848 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -29,11 +29,11 @@ If you only need to use the distributed environment without model parallelism, """ import inspect +from math import prod import torch import torch.distributed import vllm.distributed.parallel_state as vllm_parallel_state -from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm_omni.diffusion import envs @@ -59,8 +59,6 @@ _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 # Rank-layout metadata for expert parallelism. This is not a process group; # it is reused by platform-specific runtimes that must build companion groups @@ -184,7 +182,6 @@ class RankGenerator: pp: int, cfg: int, dp: int, - fs: int = 1, order: str = "tp-sp-pp-cfg-dp", rank_offset: int = 0, ) -> None: @@ -193,9 +190,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 self.name_to_size = { "tp": self.tp, @@ -203,15 +198,10 @@ 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]}), " @@ -234,7 +224,7 @@ class RankGenerator: mask[ordered_token.index(t)] = True return mask - def get_ranks(self, token, independent_ranks: bool = False): + def get_ranks(self, token): """Get rank group by input token. Arguments: @@ -243,27 +233,11 @@ 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: - for rank_group in ranks: - for i in range(len(rank_group)): - rank_group[i] += self.rank_offset + ranks = generate_masked_orthogonal_rank_groups(prod(self.ordered_size), self.ordered_size, mask) + if self.rank_offset: + return [[rank + self.rank_offset for rank in rank_group] for rank_group in ranks] return ranks @@ -376,22 +350,6 @@ 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 ( @@ -401,17 +359,6 @@ def is_dp_last_group(): ) -def get_dit_world_size(): - """Return world size for the DiT model.""" - return ( - get_data_parallel_world_size() - * get_classifier_free_guidance_world_size() - * get_sequence_parallel_world_size() - * get_pipeline_parallel_world_size() - * get_tensor_model_parallel_world_size() - ) - - # * SET @@ -498,7 +445,6 @@ def init_model_parallel_group( "expert", "sequence", "classifier_free_guidance", - "fully_shard", ], f"parallel_mode {parallel_mode} is not supported" if parallel_mode == "pipeline": return PipelineGroupCoordinator( @@ -546,19 +492,6 @@ def init_vllm_model_parallel_group( ) -def init_dit_group( - dit_parallel_size: int, - backend: str, -): - global _DIT - _DIT = torch.distributed.new_group(ranks=list(range(dit_parallel_size)), backend=backend) - - -def get_dit_group(): - assert _DIT is not None, "DIT group is not initialized" - return _DIT - - # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py def set_seq_parallel_pg( sp_ulysses_degree: int, @@ -737,8 +670,8 @@ def set_seq_parallel_pg( return ulysses_pg, ring_pg, allgather_pg -def initialize_model_parallel( - data_parallel_size: int = 1, +def _initialize_model_parallel( + data_parallel_size: int | None = None, cfg_parallel_size: int = 1, sequence_parallel_size: int | None = None, ulysses_degree: int = 1, @@ -746,9 +679,8 @@ def initialize_model_parallel( allgather_degree: int = 1, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, - fully_shard_degree: int = 1, - hsdp_replicate_size: int = 1, enable_expert_parallel: bool = False, + use_hsdp: bool = False, backend: str | None = None, ) -> None: if backend is None: @@ -768,7 +700,6 @@ def initialize_model_parallel( (causal=False only). Mutually exclusive with ulysses/ring in v1. 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 @@ -798,16 +729,20 @@ def initialize_model_parallel( ranks 8 to 15 belong to the second box. """ # Get world size and rank. Ensure some consistencies. - assert torch.distributed.is_initialized() + if not torch.distributed.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before model parallel initialization") + if model_parallel_is_initialized(): + raise RuntimeError("model parallel groups are already initialized") world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) if allgather_degree > 1: - assert ulysses_degree == 1 and ring_degree == 1, ( - "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " - f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " - f"allgather_degree={allgather_degree}." - ) + if ulysses_degree != 1 or ring_degree != 1: + raise ValueError( + "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " + f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " + f"allgather_degree={allgather_degree}." + ) expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree if sequence_parallel_size is None: @@ -820,42 +755,38 @@ def initialize_model_parallel( f"but got {sequence_parallel_size}" ) - 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) 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 " - f"tensor_parallel_size ({tensor_parallel_size}) x " - f"pipeline_parallel_size ({pipeline_parallel_size}) x" - f"sequence_parallel_size ({sequence_parallel_size}) x" - f"cfg_parallel_size " - f"({cfg_parallel_size}) x" - 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 + non_dp_size = cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size + if world_size % non_dp_size != 0: + raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})") + if use_hsdp: + if data_parallel_size not in (None, 1): + raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1") + if non_dp_size not in (1, world_size): + raise ValueError(f"HSDP non-DP parallel size must be 1 or WORLD size ({world_size}), but got {non_dp_size}") + data_parallel_size = 1 + else: + inferred_data_parallel_size = world_size // non_dp_size + if data_parallel_size is not None and data_parallel_size != inferred_data_parallel_size: + raise ValueError( + f"data_parallel_size ({data_parallel_size}) does not match WORLD-derived value " + f"({inferred_data_parallel_size})" + ) + data_parallel_size = inferred_data_parallel_size rank_generator: RankGenerator = RankGenerator( tensor_parallel_size, sequence_parallel_size, pipeline_parallel_size, cfg_parallel_size, - effective_dp_size, - fs=fully_shard_degree, + data_parallel_size, order="tp-sp-pp-cfg-dp", ) + + def get_rank_groups(token: str) -> list[list[int]]: + if use_hsdp and non_dp_size == 1: + return [[rank] for rank in range(world_size)] + return rank_generator.get_ranks(token) + use_moe_parallel_mapping = False if enable_expert_parallel: od_config = get_forward_context().omni_diffusion_config @@ -863,11 +794,11 @@ def initialize_model_parallel( if not use_moe_parallel_mapping: raise RuntimeError("Expert parallelism enabled for a non-MoE model") - sp_group_ranks = rank_generator.get_ranks("sp") + sp_group_ranks = get_rank_groups("sp") global _DP assert _DP is None, "data parallel group is already initialized" _DP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("dp"), + group_ranks=get_rank_groups("dp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="data", @@ -877,7 +808,7 @@ def initialize_model_parallel( global _CFG assert _CFG is None, "classifier_free_guidance group is already initialized" _CFG = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("cfg"), + group_ranks=get_rank_groups("cfg"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="classifier_free_guidance", @@ -885,7 +816,7 @@ def initialize_model_parallel( global _PP assert _PP is None, "pipeline model parallel group is already initialized" _PP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("pp"), + group_ranks=get_rank_groups("pp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="pipeline", @@ -899,7 +830,7 @@ def initialize_model_parallel( sp_ring_degree=ring_degree, sp_allgather_degree=allgather_degree, rank=get_world_group().rank_in_group, - world_size=dit_parallel_size, + world_size=world_size, sp_group_ranks=sp_group_ranks, ) _SP = init_model_parallel_group( @@ -925,7 +856,7 @@ def initialize_model_parallel( ) assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized" - tp_group_ranks = rank_generator.get_ranks("tp") + tp_group_ranks = get_rank_groups("tp") if use_moe_parallel_mapping: vllm_parallel_state._TP = init_vllm_model_parallel_group( group_ranks=tp_group_ranks, @@ -944,22 +875,14 @@ def initialize_model_parallel( # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP # only when constructing the vLLM EP layout for expert-parallel paths. vllm_parallel_state._DP = init_vllm_model_parallel_group( - group_ranks=rank_generator.get_ranks("cfg-dp"), + group_ranks=get_rank_groups("cfg-dp"), local_rank=get_world_group().local_rank, backend=backend, group_name="dp", ) - global _FS, _EXPERT_PARALLEL_GROUP_RANKS - 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", - ) - - _EXPERT_PARALLEL_GROUP_RANKS = rank_generator.get_ranks("tp-sp-cfg-dp") + global _EXPERT_PARALLEL_GROUP_RANKS + _EXPERT_PARALLEL_GROUP_RANKS = get_rank_groups("tp-sp-cfg-dp") if use_moe_parallel_mapping: vllm_parallel_state._EP = init_vllm_model_parallel_group( group_ranks=_EXPERT_PARALLEL_GROUP_RANKS, @@ -969,12 +892,63 @@ def initialize_model_parallel( use_all2all=True, ) - init_dit_group(dit_parallel_size, backend) + +def initialize_model_parallel( + data_parallel_size: int | None = None, + cfg_parallel_size: int = 1, + sequence_parallel_size: int | None = None, + ulysses_degree: int = 1, + ring_degree: int = 1, + allgather_degree: int = 1, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + enable_expert_parallel: bool = False, + use_hsdp: bool = False, + backend: str | None = None, +) -> None: + """Atomically initialize diffusion parallel groups. + + Configuration is validated before group creation. Any failure tears down + groups created by this attempt, including vLLM-bound coordinators. + """ + existing_groups = { + "dp": _DP, + "cfg": _CFG, + "sp": _SP, + "pp": _PP, + "vllm_tp": vllm_parallel_state._TP, + "vllm_dp": vllm_parallel_state._DP, + "vllm_pp": vllm_parallel_state._PP, + "vllm_ep": vllm_parallel_state._EP, + "vllm_pcp": vllm_parallel_state._PCP, + } + initialized_names = [name for name, group in existing_groups.items() if group is not None] + if initialized_names: + raise RuntimeError( + "model parallel state must be empty before initialization; found " + ", ".join(initialized_names) + ) + try: + _initialize_model_parallel( + data_parallel_size=data_parallel_size, + cfg_parallel_size=cfg_parallel_size, + sequence_parallel_size=sequence_parallel_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + allgather_degree=allgather_degree, + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + enable_expert_parallel=enable_expert_parallel, + use_hsdp=use_hsdp, + backend=backend, + ) + except BaseException: + destroy_model_parallel() + raise def destroy_model_parallel(): """Set the groups to none and destroy them.""" - global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS + global _DP, _CFG, _SP, _PP, _EXPERT_PARALLEL_GROUP_RANKS if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP: vllm_parallel_state._DP.destroy() @@ -1005,14 +979,14 @@ def destroy_model_parallel(): vllm_parallel_state._EP = None _EXPERT_PARALLEL_GROUP_RANKS = None + if vllm_parallel_state._PP and vllm_parallel_state._PP is not _PP: + vllm_parallel_state._PP.destroy() + vllm_parallel_state._PP = None + if _PP: _PP.destroy() _PP = None - if _FS: - _FS.destroy() - _FS = None - def destroy_distributed_environment(): global _WORLD diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index f639f3ebb..57fc26177 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -29,7 +29,7 @@ from vllm_omni.diffusion.cache.cachedit import ( ) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.parallel_state import ( - get_dit_group, + get_world_group, init_world_group, ) from vllm_omni.diffusion.distributed.utils import get_local_device @@ -423,7 +423,7 @@ def _validate_reference_image(image: Image.Image) -> None: def _dit_rank_world() -> tuple[Any, int, int]: if not dist.is_initialized(): return None, 0, 1 - group = get_dit_group() + group = get_world_group().device_group return group, dist.get_rank(group), dist.get_world_size(group) @@ -1715,7 +1715,7 @@ class MiniMaxH3Pipeline( dist.broadcast( has_audio_tensor, src=0, - group=get_dit_group(), + group=get_world_group().device_group, ) has_audio = [bool(value) for value in has_audio_tensor.tolist()] diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index e28450c90..bf75ab7f6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -18,7 +18,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedVaeMixin, ) -from vllm_omni.diffusion.distributed.parallel_state import get_dit_group +from vllm_omni.diffusion.distributed.parallel_state import get_world_group from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from .packed_tokens import minimax_h3_patchify_video_latent @@ -157,7 +157,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): ) -> None: if mode != "tile": raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}") - group = get_dit_group() + group = get_world_group().device_group world_size = dist.get_world_size(group) rank = dist.get_rank(group) parallel_size = int(parallel_size) diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index ea74fbfb7..a750d79f5 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -62,7 +62,7 @@ class OffloadConfig: dp_size = 1 if parallel_config is not None: dp_size = getattr(parallel_config, "data_parallel_size", 1) - # HSDP's fully_shard_degree also contributes to effective DP + # HSDP shard and replica sizes determine the effective group size. hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1 hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1 if use_hsdp and hsdp_shard_size > 0: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index d9030e891..faa35002d 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -279,9 +279,8 @@ class DiffusionWorker: allgather_degree=parallel_config.allgather_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, - hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1, enable_expert_parallel=parallel_config.enable_expert_parallel, + use_hsdp=parallel_config.use_hsdp, ) if ( getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e35315f7e..d30db03cb 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -971,7 +971,7 @@ class AsyncOmniEngine: ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 - data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1 + data_parallel_size = normalized_kwargs.get("data_parallel_size") tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 @@ -1004,6 +1004,11 @@ class AsyncOmniEngine: hsdp_replicate_size=hsdp_replicate_size, ) + num_gpus = normalized_kwargs.get("num_gpus") + if num_gpus is not None: + num_gpus = int(num_gpus) + parallel_config.resolve_data_parallel_size(num_gpus) + num_devices = max(1, int(parallel_config.world_size)) devices = ",".join(str(i) for i in range(num_devices)) model_class_name = kwargs.get("model_class_name", None) @@ -1092,6 +1097,8 @@ class AsyncOmniEngine: else {} ), } + if num_gpus is not None: + stage_engine_args["num_gpus"] = num_gpus # Only set dtype if it was already explicitly passed and normalized if "dtype" in normalized_kwargs: stage_engine_args["dtype"] = normalized_kwargs["dtype"]