[REFACT]Refactor diffusion parallel state (#5531)
Signed-off-by: bjf-frz <frz123db@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
bjf-frzandHongsheng Liu authored and GitHub committed
1 parent
14f2c16fb2
commit
14bd8f8772
19 files changed
+315
-387
No files matched your search
@@ -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])
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in new issue
Block a user