From f829fb3ff7c8f6115f6e2422a8e263413fb8bd25 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 4 Aug 2026 08:35:03 +0800 Subject: [PATCH] [diffusion] Fix component accuracy topology reuse (#33317) --- .../device_communicators/ipc_a2a.py | 12 +- .../runtime/distributed/parallel_state.py | 23 ++++ .../component_accuracy/config.py | 5 +- .../component_accuracy/hooks.py | 8 +- .../component_accuracy/utils.py | 6 + .../unit/test_component_accuracy_inputs.py | 45 +++++++ ...est_component_accuracy_parallel_runtime.py | 115 ++++++++++++++++++ .../test/unit/test_ipc_a2a_lifecycle.py | 41 +++++++ 8 files changed, 250 insertions(+), 5 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_component_accuracy_inputs.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_ipc_a2a_lifecycle.py diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py index 70fee44be..e0eff04be 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py @@ -99,9 +99,14 @@ class IpcA2AState: self.max_buffers = 0 self.calls = 0 self.rank = None + self.group = None self.failed = False self.inited = False + def reset(self) -> None: + """Drop mappings that belong to a model-parallel group being replaced.""" + self.__init__() + def _share(self, t, group): """Exchange `t` with the peer via torch IPC, re-opening the handle in the LOCAL device context (the mapping is only dereferenceable from the @@ -136,6 +141,7 @@ class IpcA2AState: from torch.utils.cpp_extension import load_inline self.rank = dist.get_rank(group=group) + self.group = group dev = torch.cuda.current_device() if not torch.cuda.can_device_access_peer(dev, 1 - dev): raise _Unsupported("no peer-to-peer access between the two devices") @@ -260,7 +266,11 @@ def ipc_a2a_ready(group) -> bool: from sglang.multimodal_gen.runtime.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.platforms import current_platform - if not envs.SGLANG_DIFFUSION_IPC_A2A or IPC_A2A.failed: + if not envs.SGLANG_DIFFUSION_IPC_A2A: + return False + if IPC_A2A.group is not None and IPC_A2A.group is not group: + IPC_A2A.reset() + if IPC_A2A.failed: return False # TP+Ulysses groups are strided in global-rank order, while this transport # supports the adjacent two-device topology used by TP1+U2. Reject the diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 2f45c3a65..fa2392733 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -901,10 +901,33 @@ def destroy_model_parallel() -> None: """Set the groups to none and destroy them.""" global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE + # The IPC transport keeps CUDA mappings associated with the current + # Ulysses group. Drop them before tearing down the process groups. + from .device_communicators.ipc_a2a import IPC_A2A + from .parallel_groups import PROCESS_GROUP + + IPC_A2A.reset() + for group in (_TP, _SP, _DP, _CFG, _PP, _VAE_DECODE): if group is not None: group.destroy() + # Ulysses and Ring groups are created separately from the SP coordinator, + # so GroupCoordinator.destroy() does not own or release them. Explicitly + # destroy them here; otherwise repeated Ulysses/Ring topology switches leak + # NCCL communicators and their CUDA memory. + destroyed_sequence_groups = [] + for group in (PROCESS_GROUP.ULYSSES_PG, PROCESS_GROUP.RING_PG): + if ( + group is not None + and group is not torch.distributed.group.WORLD + and all(group is not destroyed for destroyed in destroyed_sequence_groups) + ): + torch.distributed.destroy_process_group(group) + destroyed_sequence_groups.append(group) + PROCESS_GROUP.ULYSSES_PG = None + PROCESS_GROUP.RING_PG = None + for group in (_DIT, _VAE): if group is not None: torch.distributed.destroy_process_group(group) diff --git a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py index 137f0d108..44dad074d 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py +++ b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py @@ -49,8 +49,9 @@ CASE_THRESHOLDS: Dict[str, Dict[ComponentType, float]] = { "wan2_2_t2v_a14b_2gpu": {ComponentType.TRANSFORMER: 0.99}, "wan2_2_t2v_a14b_teacache_2gpu": {ComponentType.TRANSFORMER: 0.99}, "wan2_2_t2v_a14b_lora_2gpu": {ComponentType.TRANSFORMER: 0.99}, - "zimage_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.9935}, - "zimage_image_t2i_2_gpus_non_square": {ComponentType.TRANSFORMER: 0.9935}, + # H100/H200 FlashAttention runs vary between roughly 0.99338 and 0.99350. + "zimage_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.993}, + "zimage_image_t2i_2_gpus_non_square": {ComponentType.TRANSFORMER: 0.993}, } # Active skip policy. Keep this limited to cases with current, concrete evidence diff --git a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py index 970bb2f31..a6fc0f2b1 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py +++ b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py @@ -362,10 +362,14 @@ def _build_transformer_hook_inputs( inputs["pooled_projections"] = rng.randn( (1, pooled_channels), device, torch.bfloat16 ) - if ( + supports_text_attention_mask = ( "encoder_attention_mask" in param_names or "encoder_hidden_states_mask" in param_names - ): + ) + uses_ring_parallel = (case.server_args.ring_degree or 1) > 1 + if supports_text_attention_mask and not uses_ring_parallel: + # This synthetic mask is all True and has no semantic effect. Omit it + # for ring-parallel cases, whose masked attention path is unsupported. attention_mask = torch.ones( 1, DEFAULT_TEXT_SEQ_LEN, device=device, dtype=torch.bool ) diff --git a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/utils.py b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/utils.py index bf7d5e20d..453e94b22 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/utils.py +++ b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/utils.py @@ -16,8 +16,10 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( destroy_model_parallel, get_classifier_free_guidance_world_size, get_data_parallel_world_size, + get_ring_parallel_world_size, get_sequence_parallel_world_size, get_tensor_model_parallel_world_size, + get_ulysses_parallel_world_size, maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, ) @@ -305,11 +307,15 @@ def initialize_parallel_runtime(sgl_args: ServerArgs) -> None: if model_parallel_is_initialized(): current_tp = get_tensor_model_parallel_world_size() current_sp = get_sequence_parallel_world_size() + current_ulysses = get_ulysses_parallel_world_size() + current_ring = get_ring_parallel_world_size() current_dp = get_data_parallel_world_size() current_cfg = get_classifier_free_guidance_world_size() if ( current_tp == tp_size and current_sp == sp_degree + and current_ulysses == ulysses_degree + and current_ring == ring_degree and current_dp == dp_size and current_cfg == cfg_degree ): diff --git a/python/sglang/multimodal_gen/test/unit/test_component_accuracy_inputs.py b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_inputs.py new file mode 100644 index 000000000..66200ea7d --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_inputs.py @@ -0,0 +1,45 @@ +from types import SimpleNamespace + +import torch.nn as nn + +from sglang.multimodal_gen.test.single_test_file.component_accuracy.hooks import ( + _build_transformer_hook_inputs, +) + + +class _TransformerWithOptionalMask(nn.Module): + def forward( + self, + hidden_states, + encoder_hidden_states, + timestep, + encoder_hidden_states_mask=None, + ): + raise NotImplementedError + + +def _case(*, ring_degree: int) -> SimpleNamespace: + return SimpleNamespace( + server_args=SimpleNamespace( + model_path="test/model", + ring_degree=ring_degree, + ) + ) + + +def test_omits_noop_attention_mask_for_ring_parallel_case(): + inputs = _build_transformer_hook_inputs( + _case(ring_degree=2), _TransformerWithOptionalMask(), "cpu" + ) + + assert "encoder_hidden_states_mask" not in inputs + assert "encoder_attention_mask" not in inputs + + +def test_keeps_attention_mask_for_non_ring_case(): + inputs = _build_transformer_hook_inputs( + _case(ring_degree=1), _TransformerWithOptionalMask(), "cpu" + ) + + assert inputs["encoder_hidden_states_mask"].all() + assert inputs["encoder_attention_mask"].all() diff --git a/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py new file mode 100644 index 000000000..5ea5b96fd --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py @@ -0,0 +1,115 @@ +from contextlib import ExitStack +from types import SimpleNamespace +from unittest.mock import call, patch + +from sglang.multimodal_gen.runtime.distributed import parallel_state +from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IPC_A2A, +) +from sglang.multimodal_gen.runtime.distributed.parallel_groups import PROCESS_GROUP +from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + initialize_parallel_runtime, +) + +_UTILS = "sglang.multimodal_gen.test.single_test_file.component_accuracy.utils" + + +def _server_args(*, ulysses_degree: int, ring_degree: int) -> SimpleNamespace: + return SimpleNamespace( + tp_size=1, + sp_degree=2, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + dp_size=1, + cfg_parallel_degree=1, + ) + + +def _patch_current_topology( + stack: ExitStack, *, ulysses_degree: int, ring_degree: int +) -> None: + for context_manager in ( + patch(f"{_UTILS}.model_parallel_is_initialized", return_value=True), + patch(f"{_UTILS}.get_tensor_model_parallel_world_size", return_value=1), + patch(f"{_UTILS}.get_sequence_parallel_world_size", return_value=2), + patch( + f"{_UTILS}.get_ulysses_parallel_world_size", + return_value=ulysses_degree, + ), + patch(f"{_UTILS}.get_ring_parallel_world_size", return_value=ring_degree), + patch(f"{_UTILS}.get_data_parallel_world_size", return_value=1), + patch(f"{_UTILS}.get_classifier_free_guidance_world_size", return_value=1), + ): + stack.enter_context(context_manager) + + +def test_reinitializes_when_sp_decomposition_changes(): + with ExitStack() as stack: + _patch_current_topology(stack, ulysses_degree=2, ring_degree=1) + stack.enter_context( + patch(f"{_UTILS}.torch.distributed.is_initialized", return_value=True) + ) + barrier = stack.enter_context(patch(f"{_UTILS}.torch.distributed.barrier")) + destroy = stack.enter_context(patch(f"{_UTILS}.destroy_model_parallel")) + initialize = stack.enter_context( + patch(f"{_UTILS}.maybe_init_distributed_environment_and_model_parallel") + ) + initialize_parallel_runtime(_server_args(ulysses_degree=1, ring_degree=2)) + + destroy.assert_called_once_with() + initialize.assert_called_once_with( + tp_size=1, + sp_size=2, + cfg_degree=1, + ulysses_degree=1, + ring_degree=2, + dp_size=1, + ) + assert barrier.call_count == 2 + + +def test_reuses_matching_sp_decomposition(): + with ExitStack() as stack: + _patch_current_topology(stack, ulysses_degree=1, ring_degree=2) + stack.enter_context( + patch(f"{_UTILS}.torch.distributed.is_initialized", return_value=True) + ) + destroy = stack.enter_context(patch(f"{_UTILS}.destroy_model_parallel")) + initialize = stack.enter_context( + patch(f"{_UTILS}.maybe_init_distributed_environment_and_model_parallel") + ) + initialize_parallel_runtime(_server_args(ulysses_degree=1, ring_degree=2)) + + destroy.assert_not_called() + initialize.assert_not_called() + + +def test_destroy_releases_sequence_parallel_subgroups_after_partial_init(): + ulysses_group = object() + ring_group = object() + + with ExitStack() as stack: + for name in ( + "_TP", + "_SP", + "_DP", + "_CFG", + "_PP", + "_VAE_DECODE", + "_DIT", + "_VAE", + ): + stack.enter_context(patch.object(parallel_state, name, None)) + stack.enter_context(patch.object(PROCESS_GROUP, "ULYSSES_PG", ulysses_group)) + stack.enter_context(patch.object(PROCESS_GROUP, "RING_PG", ring_group)) + reset_ipc = stack.enter_context(patch.object(IPC_A2A, "reset")) + destroy_group = stack.enter_context( + patch.object(parallel_state.torch.distributed, "destroy_process_group") + ) + + parallel_state.destroy_model_parallel() + + reset_ipc.assert_called_once_with() + assert destroy_group.call_args_list == [call(ulysses_group), call(ring_group)] + assert PROCESS_GROUP.ULYSSES_PG is None + assert PROCESS_GROUP.RING_PG is None diff --git a/python/sglang/multimodal_gen/test/unit/test_ipc_a2a_lifecycle.py b/python/sglang/multimodal_gen/test/unit/test_ipc_a2a_lifecycle.py new file mode 100644 index 000000000..211a9f33c --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_ipc_a2a_lifecycle.py @@ -0,0 +1,41 @@ +from unittest.mock import patch + +from sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a import ( + IpcA2AState, + ipc_a2a_ready, +) + +_IPC = "sglang.multimodal_gen.runtime.distributed.device_communicators.ipc_a2a" + + +def test_reinitializes_ipc_transport_for_replaced_process_group(): + state = IpcA2AState() + old_group = object() + new_group = object() + state.inited = True + state.group = old_group + state.calls = 7 + + def initialize(group): + state.inited = True + state.group = group + + with ( + patch(f"{_IPC}.IPC_A2A", state), + patch(f"{_IPC}.envs.SGLANG_DIFFUSION_IPC_A2A", True), + patch( + "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.distributed.get_tp_world_size", + return_value=1, + ), + patch.object(state, "init", side_effect=initialize) as init, + patch(f"{_IPC}.torch.cuda.is_current_stream_capturing", return_value=False), + ): + assert ipc_a2a_ready(new_group) + + init.assert_called_once_with(new_group) + assert state.group is new_group + assert state.calls == 0