[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 2026-08-19 11:23:10 +08:00
1 parent 14f2c16fb2
commit 14bd8f8772
19 files changed
+315 -387

No files matched your search

+3 -4
View File
@@ -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)