From 039808e0d97d7969a2cb102d074c1e45b8ceef0b Mon Sep 17 00:00:00 2001 From: Xenoryn Date: Sun, 6 Sep 2026 16:23:30 +0800 Subject: [PATCH] [Perf][Diffusion] Skip the Ulysses length all-gather when shards are padded equal (#5717) Signed-off-by: Xenoryn Co-authored-by: Claude --- tests/diffusion/attention/test_ulysses_uaa.py | 377 ++++++++++++++++++ .../distributed/test_sp_plan_hooks.py | 114 ++++++ .../diffusion/attention/parallel/ulysses.py | 33 +- vllm_omni/diffusion/forward_context.py | 15 +- .../diffusion/hooks/sequence_parallel.py | 8 +- 5 files changed, 537 insertions(+), 10 deletions(-) diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 041bb941a..230eb8b71 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -27,6 +27,202 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_equal_rank_contract_skips_length_gather( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + observed = [] + + def fake_all_to_all(pg, tensor, **kwargs): + observed.append((kwargs["seq_lens"], kwargs["padded_head_cnt"])) + return tensor, tensor.shape[2] + + monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr( + ulysses, + "_all_gather_int", + lambda *args, **kwargs: pytest.fail("length gather must be skipped"), + ) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + with set_forward_context(): + get_forward_context()._sp_equal_pad_stack.append(True) + strategy.pre_attention( + torch.zeros(1, 3, 28, 4), + torch.zeros(1, 3, 7, 4), + torch.zeros(1, 3, 7, 4), + None, + ) + + assert observed == [([3, 3], 32), ([3, 3], 8), ([3, 3], 8)] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_hybrid_equal_rank_contract_skips_ulysses_gather_keeps_ring_check( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Contract + ring: the fast path skips only the Ulysses gather. + + The ring equality check is a separate collective on ring_group and must + still run, receiving s_global derived from the contract. + """ + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ring_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 2 + + sp_group = FakeGroup() + observed = [] + ring_checks = [] + + def fake_all_to_all(pg, tensor, **kwargs): + observed.append((kwargs["seq_lens"], kwargs["padded_head_cnt"])) + return tensor, tensor.shape[2] + + def fake_gather_int(pg, value, *, device): + if pg is sp_group.ulysses_group: + pytest.fail("Ulysses length gather must be skipped under the contract") + assert pg is sp_group.ring_group + ring_checks.append(int(value)) + return [int(value), int(value)] + + monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all) + monkeypatch.setattr(ulysses, "_all_gather_int", fake_gather_int) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + strategy = ulysses.UlyssesParallelAttention( + sp_group, + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + with set_forward_context(): + get_forward_context()._sp_equal_pad_stack.append(True) + strategy.pre_attention( + torch.zeros(1, 3, 28, 4), + torch.zeros(1, 3, 7, 4), + torch.zeros(1, 3, 7, 4), + None, + ) + + # Ring check ran: s_global = ulysses_world_size * local = 2 * 3. + assert ring_checks == [6] + assert observed == [([3, 3], 32), ([3, 3], 8), ([3, 3], 8)] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_without_contract_keeps_dynamic_length_gather( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No auto_pad boundary means shards may differ, so the gather must stay.""" + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + observed = [] + monkeypatch.setattr( + ulysses, + "_all_gather_int", + lambda *args, **kwargs: [3, 4], + ) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + + def fake_all_to_all(pg, tensor, **kwargs): + observed.append(kwargs["seq_lens"]) + return tensor, tensor.shape[2] + + monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + with set_forward_context(): + strategy.pre_attention( + torch.zeros(1, 3, 28, 4), + torch.zeros(1, 3, 7, 4), + torch.zeros(1, 3, 7, 4), + None, + ) + + assert observed == [[3, 4], [3, 4], [3, 4]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_mixed_boundaries_keep_the_dynamic_length_gather( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """One manual boundary inside an auto_pad region voids the contract.""" + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + gathered = [] + + def fake_gather(*args, **kwargs): + gathered.append(args) + return [3, 4] + + monkeypatch.setattr(ulysses, "_all_gather_int", fake_gather) + monkeypatch.setattr( + ulysses, + "_ulysses_all_to_all_any_qkv", + lambda pg, tensor, **kwargs: (tensor, tensor.shape[2]), + ) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + with set_forward_context(): + ctx = get_forward_context() + ctx._sp_equal_pad_stack.extend([True, False]) + assert not ctx.sp_rank_local_seq_lens_equal + strategy.pre_attention( + torch.zeros(1, 3, 28, 4), + torch.zeros(1, 3, 7, 4), + torch.zeros(1, 3, 7, 4), + None, + ) + + assert len(gathered) == 1 + + @pytest.mark.core_model @pytest.mark.cpu def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio( @@ -634,6 +830,187 @@ def test_ulysses_uaa_matches_baseline( pass +def _run_contract_skip_case( + local_rank: int, + world_size: int, + init_method: str, + input_file: str, + num_heads: int, + head_size: int, + num_kv_heads: int | None, +) -> None: + """Per-rank worker: slow path vs contract fast path on real rank pairs. + + Phase 1 leaves the equal-pad stack empty so the Ulysses length all-gather + really runs (counted); phase 2 pushes the split hook's auto_pad marker and + replaces _all_gather_int with a raising guard. With ring_degree=1 any call + to it is the Ulysses length gather, so the guard proves the skip, and the + per-rank output comparison proves the fast path is numerically equivalent + to the slow path it replaces. + """ + from vllm_omni.diffusion.attention.parallel import ulysses as ulysses_mod + + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + + os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" + + init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=world_size, + ring_degree=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=world_size, + ring_degree=1, + cfg_parallel_size=1, + ulysses_mode="advanced_uaa", + ) + od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config) + + with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config): + attn = Attention( + num_heads=num_heads, + head_size=head_size, + causal=False, + softmax_scale=1.0 / (head_size**0.5), + num_kv_heads=num_kv_heads, + ).to(device=device, dtype=torch.float32) + + with np.load(input_file, allow_pickle=False) as payload: + q_full = torch.from_numpy(payload["q"]).to(device=device) + k_full = torch.from_numpy(payload["k"]).to(device=device) + v_full = torch.from_numpy(payload["v"]).to(device=device) + + def make_shards() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # Fresh contiguous shards for each phase; seq_len is divisible by + # world_size, so every rank gets the same local length. + return ( + torch.tensor_split(q_full, world_size, dim=1)[local_rank].contiguous(), + torch.tensor_split(k_full, world_size, dim=1)[local_rank].contiguous(), + torch.tensor_split(v_full, world_size, dim=1)[local_rank].contiguous(), + ) + + # The Attention layer only enables SP communication when + # ForwardContext.sp_active is True; in production the split hook sets + # this (and pushes the auto_pad equal-length marker onto + # _sp_equal_pad_stack). Here we drive it directly. + ctx = get_forward_context() + ctx._sp_shard_depth = 1 + + def make_sdp_ctx() -> object: + # Fresh instance per phase: the @contextmanager instance releases + # its state on first exit and cannot be re-entered. + return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True) + + # Phase 1: contract absent -> the length all-gather must really run. + real_gather = ulysses_mod._all_gather_int + gather_calls: list[object] = [] + + def counting_gather(pg: object, value: int, *, device: torch.device) -> list[int]: + gather_calls.append(pg) + return real_gather(pg, value, device=device) + + ulysses_mod._all_gather_int = counting_gather + try: + q, k, v = make_shards() + with make_sdp_ctx(): + out_slow = attn(q, k, v, attn_metadata=None).contiguous() + finally: + ulysses_mod._all_gather_int = real_gather + assert len(gather_calls) >= 1, "slow path must call the Ulysses length all-gather" + + # Phase 2: replicate the split hook's auto_pad marker so the fast + # path fires; any _all_gather_int call in this phase is the Ulysses + # length gather and must not happen. + ctx._sp_equal_pad_stack.append(True) + + def fail_on_gather(pg: object, value: int, *, device: torch.device) -> list[int]: + raise AssertionError("the Ulysses length all-gather must be skipped under the equal-pad contract") + + ulysses_mod._all_gather_int = fail_on_gather + try: + q, k, v = make_shards() + with make_sdp_ctx(): + out_fast = attn(q, k, v, attn_metadata=None).contiguous() + finally: + ulysses_mod._all_gather_int = real_gather + ctx._sp_equal_pad_stack.pop() + + torch.testing.assert_close(out_fast, out_slow, atol=1e-5, rtol=1e-5) + + destroy_distributed_env() + + +@pytest.mark.parametrize( + "num_heads,num_kv_heads", + [ + (3, None), # MHA: heads padded 3 -> 4 by ulysses_degree=2 + (28, 7), # GQA: derived-Q padding 28/7 -> 32/8, the main case for the skip + ], +) +def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: int, num_kv_heads: int | None) -> None: + """Real 2-rank parity: contract fast path == slow path, gather skipped. + + The CPU-only contract tests replace the all-to-all with a fake; this test + runs the real strategy on two ranks. Each rank first computes the slow + path (empty equal-pad stack, real length all-gather), then the fast path + with the split hook's auto_pad marker pushed and _all_gather_int replaced + by a raising guard -- which also proves the collective was truly skipped, + not just silently replaced. + """ + sp_world_size = 2 + if current_omni_platform.get_device_count() < sp_world_size: + pytest.skip(f"Test requires {sp_world_size} GPUs") + + batch_size = 2 + head_size = 8 + seq_len = 6 # divisible by 2 so both ranks hold equal-length shards + + init_method = get_distributed_init_method() + + with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: + input_file = f_in.name + + try: + torch.manual_seed(0) + kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + np.savez( + input_file, + q=torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32).numpy(), + k=torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32).numpy(), + v=torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32).numpy(), + ) + + torch.multiprocessing.spawn( + _run_contract_skip_case, + args=( + sp_world_size, + init_method, + input_file, + num_heads, + head_size, + num_kv_heads, + ), + nprocs=sp_world_size, + ) + finally: + try: + os.remove(input_file) + except OSError: + pass + + @pytest.mark.parametrize( "num_heads,num_kv_heads,joint_len", [ diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py index 4ae6ec064..9a3c0ee82 100644 --- a/tests/diffusion/distributed/test_sp_plan_hooks.py +++ b/tests/diffusion/distributed/test_sp_plan_hooks.py @@ -1113,6 +1113,120 @@ class TestStrictModeSplitValidation: hook._prepare_sp_input(x, metadata["x"], (), {}) +@pytest.mark.cpu +class TestEqualPadContract: + """The stack of per-boundary auto_pad flags that lets attention skip a gather. + + A boundary is recorded as equal-length only when the framework padded it, + because only then is every rank's shard guaranteed the same size. + """ + + def _split_hook(self, *, auto_pad: bool, monkeypatch): + from vllm_omni.diffusion.attention import selector + from vllm_omni.diffusion.distributed import parallel_state, sp_sharding + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import ( + SequenceParallelSplitHook, + ) + + class MaskBackend: + supports_attention_mask = True + + @staticmethod + def get_name(): + return "test" + + monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2) + monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1) + monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1) + monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None)) + # sp_shard, used by the manual path, binds these at import time. + monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 2) + monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: 1) + + hook = SequenceParallelSplitHook( + { + 0: SequenceParallelInput( + split_dim=1, + expected_dims=3, + split_output=True, + auto_pad=auto_pad, + ) + }, + SequenceParallelConfig(ulysses_degree=2), + ) + hook.initialize_hook(nn.Identity()) + return hook + + def test_auto_pad_split_claims_equal_rank_lengths(self, monkeypatch): + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + set_forward_context, + ) + + hook = self._split_hook(auto_pad=True, monkeypatch=monkeypatch) + with set_forward_context(): + hook.pre_forward(nn.Identity()) + shard = hook.post_forward(nn.Identity(), torch.arange(5).reshape(1, 5, 1)) + ctx = get_forward_context() + + # 5 tokens padded to 6, so both ranks hold 3. + assert shard.shape[1] == 3 + assert ctx._sp_equal_pad_stack == [True] + assert ctx.sp_rank_local_seq_lens_equal + + def test_manual_split_does_not_claim_equal_rank_lengths(self, monkeypatch): + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + set_forward_context, + ) + + hook = self._split_hook(auto_pad=False, monkeypatch=monkeypatch) + with set_forward_context(): + hook.pre_forward(nn.Identity()) + hook.post_forward(nn.Identity(), torch.arange(6).reshape(1, 6, 1)) + ctx = get_forward_context() + + assert ctx._sp_equal_pad_stack == [False] + assert not ctx.sp_rank_local_seq_lens_equal + + def test_gather_pops_only_its_own_boundary(self, monkeypatch): + """A stack, not a counter: the gather must not resurrect the contract. + + With one auto_pad and one manual boundary open, a counter could not + tell which one a gather closes and would wrongly report equal lengths. + """ + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + set_forward_context, + ) + from vllm_omni.diffusion.hooks import sequence_parallel + from vllm_omni.diffusion.hooks.sequence_parallel import ( + SequenceParallelGatherHook, + ) + + monkeypatch.setattr( + sequence_parallel, + "sp_gather", + lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim), + ) + hook = SequenceParallelGatherHook( + SequenceParallelOutput(gather_dim=1, expected_dims=3), + SequenceParallelConfig(ulysses_degree=2), + ) + + with set_forward_context(): + ctx = get_forward_context() + ctx._sp_shard_depth = 2 + ctx._sp_equal_pad_stack.extend([False, True]) + hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4)) + + assert ctx._sp_shard_depth == 1 + assert ctx._sp_equal_pad_stack == [False] + assert not ctx.sp_rank_local_seq_lens_equal + + @pytest.mark.cpu class TestSequenceParallelConfig: """Test SequenceParallelConfig dataclass.""" diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index d56838196..cbaa8ba62 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -14,7 +14,11 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator -from vllm_omni.diffusion.forward_context import get_ulysses_mode +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + get_ulysses_mode, + is_forward_context_available, +) logger = init_logger(__name__) @@ -142,7 +146,7 @@ def _ulysses_all_to_all_any_qkv( input_split_sizes = [s_local] * world_size output_split_sizes = seq_lens - s_global = int(sum(output_split_sizes)) + s_global = sum(output_split_sizes) out = torch.empty((s_global, bsz, head_cnt_local, head_dim), device=x.device, dtype=x.dtype) dist.all_to_all_single( @@ -177,7 +181,7 @@ def _ulysses_all_to_all_any_o( return x bsz, s_global, head_cnt_local, head_dim = x.shape - s_local = int(local_seq_len) + s_local = local_seq_len # (B, S_global, H_local, D) -> (S_global, B, H_local, D) x_t = x.permute(1, 0, 2, 3).contiguous() @@ -411,9 +415,22 @@ class UlyssesParallelAttention: f"(got scatter_idx={self._scatter_idx}, gather_idx={self._gather_idx})." ) - local_seq_len = int(query.shape[1]) - seq_lens = _all_gather_int(self._ulysses_pg, local_seq_len, device=query.device) - s_global = int(sum(seq_lens)) + if is_forward_context_available() and get_forward_context().sp_rank_local_seq_lens_equal: + # auto_pad already made every rank's shard the same length, so + # the lengths are known without a collective. Keep local_seq_len + # as a SymInt (no int()) so torch.compile(dynamic=True) does not + # specialize on it, and skip the host sync + graph break that + # _all_gather_int would introduce in every attention call. + local_seq_len = query.shape[1] + seq_lens = [local_seq_len] * ulysses_world_size + else: + local_seq_len = int(query.shape[1]) + seq_lens = _all_gather_int( + self._ulysses_pg, + local_seq_len, + device=query.device, + ) + s_global = sum(seq_lens) # In hybrid Ulysses+Ring, Ring attention uses P2P send/recv with fixed-shape # buffers. This requires all ring ranks to have the same seq_len after the @@ -571,8 +588,8 @@ class UlyssesParallelAttention: joint_len=joint_len, joint_strategy=joint_strategy, use_uaa=(mode == "advanced_uaa"), - uaa_seq_lens=tuple(int(x) for x in seq_lens) if mode == "advanced_uaa" else (), - uaa_local_seq_len=int(local_seq_len) if mode == "advanced_uaa" else 0, + uaa_seq_lens=tuple(seq_lens) if mode == "advanced_uaa" else (), + uaa_local_seq_len=local_seq_len if mode == "advanced_uaa" else 0, orig_head_cnt=int(orig_head_cnt) if mode == "advanced_uaa" else 0, joint_orig_head_cnt=int(joint_orig_head_cnt) if mode == "advanced_uaa" else 0, ) diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 8924e175e..1459da495 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -4,7 +4,7 @@ from __future__ import annotations from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any import torch @@ -66,6 +66,10 @@ class ForwardContext: # Tracks the depth of SP sharding - incremented on shard, decremented on gather # Used by attention layers to determine if SP communication should be enabled _sp_shard_depth: int = 0 + # One entry per active SP split boundary: whether it auto-pads, which makes + # every rank's shard the same length and lets attention collectives skip a + # runtime length all-gather. Pushed on split, popped on gather. + _sp_equal_pad_stack: list[bool] = field(default_factory=list) @property def sp_active(self) -> bool: @@ -88,6 +92,15 @@ class ForwardContext: sp_size = self.omni_diffusion_config.parallel_config.sequence_parallel_size return sp_size is not None and sp_size > 1 + @property + def sp_rank_local_seq_lens_equal(self) -> bool: + """Whether every active SP boundary guarantees equal local shard sizes. + + Only framework-managed auto_pad boundaries provide this contract; a + region that also shards manually keeps the dynamic length exchange. + """ + return bool(self._sp_equal_pad_stack) and all(self._sp_equal_pad_stack) + def __post_init__(self): pass diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 20800ffc0..95f5b0f47 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -264,6 +264,7 @@ class SequenceParallelSplitHook(ModelHook): output_list = [output] if is_tensor else list(output) actually_sharded = False + equal_rank_seq_lens = True for index, spm in self.metadata.items(): if not isinstance(index, int): @@ -277,10 +278,13 @@ class SequenceParallelSplitHook(ModelHook): output_list[index] = self._prepare_sp_input(original, spm, self._last_args, self._last_kwargs) if output_list[index] is not original: actually_sharded = True + equal_rank_seq_lens &= isinstance(spm, SequenceParallelInput) and spm.auto_pad # Mark SP as active only if at least one tensor was actually sharded if actually_sharded and is_forward_context_available(): - get_forward_context()._sp_shard_depth += 1 + ctx = get_forward_context() + ctx._sp_shard_depth += 1 + ctx._sp_equal_pad_stack.append(equal_rank_seq_lens) return output_list[0] if is_tensor else type(output)(output_list) @@ -553,6 +557,8 @@ class SequenceParallelGatherHook(ModelHook): if actually_gathered and is_forward_context_available(): ctx = get_forward_context() ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1) + if ctx._sp_equal_pad_stack: + ctx._sp_equal_pad_stack.pop() return output[0] if is_tensor else type(output)(output)