[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
@@ -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