[REFACT]Refactor diffusion parallel state (#5531)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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.
 
 ---
 

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)
 
 

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

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)

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

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

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

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

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 = []

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",
 ]

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

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

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)

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

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()]
 

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)

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:

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)

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