chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support context parallel zigzag reordering in dump comparator (#19281)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-25 09:46:17 +08:00
committed by GitHub
parent 0de1f4b07b
commit 2e2b18e870
13 changed files with 423 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
python/sglang/srt/debug_utils/comparator/unshard/__init__.py → python/sglang/srt/debug_utils/comparator/aligner/__init__.py
View File

82
python/sglang/srt/debug_utils/comparator/aligner/reorder.py Normal file
View File

@@ -0,0 +1,82 @@
from typing import Literal
import torch
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
from sglang.srt.debug_utils.comparator.utils import _FrozenBase
class ZigzagToNaturalParams(_FrozenBase):
op: Literal["zigzag_to_natural"] = "zigzag_to_natural"
dim: int
cp_size: int
ReorderParams = ZigzagToNaturalParams
class ReorderPlan(_FrozenBase):
params: ReorderParams
_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {"s"}
def compute_reorder_plans(
dim_specs: list[DimSpec],
parallel_infos: list[dict[ParallelAxis, AxisInfo]],
) -> list[ReorderPlan]:
plans: list[ReorderPlan] = []
for dim_index, spec in enumerate(dim_specs):
if (
spec.ordering is not None
and spec.ordering != Ordering.NATURAL
and spec.parallel is not None
):
if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES:
raise ValueError(
f"Zigzag ordering is only supported on sequence dims "
f"(bshd/sbhd format, dim name must be one of "
f"{sorted(_ALLOWED_ZIGZAG_DIM_NAMES)}), "
f"but got dim name {spec.name!r} in {spec}"
)
assert spec.ordering == Ordering.ZIGZAG
axis_size: int = parallel_infos[0][spec.parallel].axis_size
plans.append(
ReorderPlan(
params=ZigzagToNaturalParams(dim=dim_index, cp_size=axis_size),
)
)
return plans
def execute_reorder_plan(
plan: ReorderPlan,
tensors: list[torch.Tensor],
) -> list[torch.Tensor]:
return [
_reorder_zigzag_to_natural(
tensor, dim=plan.params.dim, cp_size=plan.params.cp_size
)
for tensor in tensors
]
def _reorder_zigzag_to_natural(
tensor: torch.Tensor, *, dim: int, cp_size: int
) -> torch.Tensor:
"""Undo CP zigzag interleaving, restoring natural chunk order.
Generalized from Megatron-LM _undo_attention_load_balancing
(megatron/core/ssm/mamba_context_parallel.py:360-373).
"""
num_chunks: int = cp_size * 2
chunks: tuple[torch.Tensor, ...] = tensor.chunk(num_chunks, dim=dim)
order: list[int] = [2 * i for i in range(cp_size)] + [
num_chunks - 2 * i - 1 for i in range(cp_size)
]
return torch.cat([chunks[i] for i in order], dim=dim)

0
python/sglang/srt/debug_utils/comparator/aligner/unshard/__init__.py Normal file
View File

2
python/sglang/srt/debug_utils/comparator/unshard/executor.py → python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py
View File

@@ -1,6 +1,6 @@
import torch
from sglang.srt.debug_utils.comparator.unshard.types import (
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
ConcatParams,
UnshardParams,
UnshardPlan,

2
python/sglang/srt/debug_utils/comparator/unshard/parallel_info.py → python/sglang/srt/debug_utils/comparator/aligner/unshard/parallel_info.py
View File

@@ -1,7 +1,7 @@
from typing import Optional
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
_PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")

8
python/sglang/srt/debug_utils/comparator/unshard/planner.py → python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py
View File

@@ -1,13 +1,13 @@
from collections import defaultdict
from typing import NamedTuple
from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
from sglang.srt.debug_utils.comparator.unshard.types import (
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
AxisInfo,
ConcatParams,
UnshardParams,
UnshardPlan,
)
from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis
# _CoordsList[tensor_index][axis] =
# the axis_rank (shard position) of the tensor_index-th tensor along `axis`
@@ -127,8 +127,4 @@ def _resolve_unshard_params(*, spec: DimSpec, dim_index: int) -> UnshardParams:
raise NotImplementedError(
f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"
)
if spec.ordering is not None and spec.ordering != Ordering.NATURAL:
raise NotImplementedError(
f"Unshard for ordering={spec.ordering} not yet implemented (Phase 2)"
)
return ConcatParams(dim=dim_index)

10
python/sglang/srt/debug_utils/comparator/unshard/types.py → python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py
View File

@@ -16,11 +16,6 @@ class ConcatParams(_FrozenBase):
dim: int
# Phase 2: add ReduceSumParams, CpZigzagParams here, then change UnshardParams to:
# UnshardParams = Annotated[
# Union[ConcatParams, ReduceSumParams, CpZigzagParams],
# Field(discriminator="op"),
# ]
UnshardParams = ConcatParams
@@ -33,8 +28,3 @@ class UnshardPlan(_FrozenBase):
# plan[0] (CP): groups=[[0,2],[1,3]] — 4 tensors → 2 tensors
# plan[1] (TP): groups=[[0,1]] — 2 tensors → 1 tensor
groups: list[list[int]]
# Union of all plan types. Future pipeline components (e.g. reduction,
# reordering) will add their own plan types here.
Plan = UnshardPlan

47
python/sglang/srt/debug_utils/comparator/pipeline.py
View File

@@ -1,22 +1,33 @@
from pathlib import Path
from typing import Any, Optional
from typing import Any, Optional, Union
import torch
from sglang.srt.debug_utils.comparator.aligner.reorder import (
ReorderPlan,
compute_reorder_plans,
execute_reorder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
execute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.parallel_info import (
normalize_parallel_info,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
from sglang.srt.debug_utils.comparator.dims import parse_dims
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
SkipRecord,
)
from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors
from sglang.srt.debug_utils.comparator.unshard.executor import execute_unshard_plan
from sglang.srt.debug_utils.comparator.unshard.parallel_info import (
normalize_parallel_info,
)
from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
from sglang.srt.debug_utils.comparator.unshard.types import Plan, UnshardPlan
from sglang.srt.debug_utils.dump_loader import ValueWithMeta
Plan = Union[UnshardPlan, ReorderPlan]
def process_tensor_group(
*,
@@ -83,7 +94,13 @@ def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]:
dim_specs = parse_dims(dims_str)
parallel_infos = [normalize_parallel_info(meta) for meta in metas]
return compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos)
unshard_plans = compute_unshard_plan(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
reorder_plans = compute_reorder_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
return [*unshard_plans, *reorder_plans]
def _extract_tensors(
@@ -106,10 +123,16 @@ def _execute_plans(
current = tensors
for plan in plans:
if isinstance(plan, UnshardPlan):
current = execute_unshard_plan(plan, current)
else:
raise NotImplementedError(f"Unknown {plan=}")
current = _execute_plan(current, plan)
assert len(current) == 1
return current[0]
def _execute_plan(tensors, plan):
if isinstance(plan, UnshardPlan):
return execute_unshard_plan(plan, tensors)
elif isinstance(plan, ReorderPlan):
return execute_reorder_plan(plan, tensors)
else:
raise NotImplementedError(f"Unknown {plan=}")

158
test/registered/debug_utils/comparator/aligner/test_reorder.py Normal file
View File

@@ -0,0 +1,158 @@
import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.aligner.reorder import (
ReorderPlan,
_reorder_zigzag_to_natural,
compute_reorder_plans,
execute_reorder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
execute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestZigzagToNatural:
def test_zigzag_to_natural_cp2(self) -> None:
"""cp_size=2: zigzag order [0,3,1,2] -> natural [0,1,2,3]."""
natural = torch.arange(24).reshape(4, 6)
chunks = list(natural.chunk(4, dim=0))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=2)
assert torch.equal(result, natural)
def test_zigzag_to_natural_cp3(self) -> None:
"""cp_size=3: zigzag 162534 -> natural 123456 (1-indexed)."""
natural = torch.arange(60).reshape(6, 10)
chunks = list(natural.chunk(6, dim=0))
zigzag_order: list[int] = [0, 5, 1, 4, 2, 3]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=3)
assert torch.equal(result, natural)
def test_zigzag_to_natural_arbitrary_dim(self) -> None:
"""Reorder along dim=1 instead of dim=0."""
natural = torch.arange(48).reshape(3, 4, 4)
chunks = list(natural.chunk(4, dim=1))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=1)
result = _reorder_zigzag_to_natural(zigzagged, dim=1, cp_size=2)
assert torch.equal(result, natural)
class TestComputeReorderPlans:
def test_compute_reorder_plans_zigzag(self) -> None:
"""s(cp,zigzag) produces a ReorderPlan."""
dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
},
]
plans = compute_reorder_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
assert len(plans) == 1
assert plans[0].params.op == "zigzag_to_natural"
assert plans[0].params.dim == 1
assert plans[0].params.cp_size == 2
def test_compute_reorder_plans_non_seq_dim_raises(self) -> None:
"""Zigzag on non-sequence dim (e.g. t(cp,zigzag)) raises ValueError."""
dim_specs = parse_dims("t(cp,zigzag) h(tp)")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
},
]
with pytest.raises(ValueError, match="only supported on sequence dims"):
compute_reorder_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
def test_compute_reorder_plans_natural(self) -> None:
"""s(cp) and s(cp,natural) produce no reorder plans."""
for dims_str in ["b s(cp) h(tp)", "b s(cp,natural) h(tp)"]:
dim_specs = parse_dims(dims_str)
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
},
]
plans = compute_reorder_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
assert plans == []
class TestCpZigzagTpE2E:
def test_cp_zigzag_tp_e2e(self) -> None:
"""CP=2 zigzag + TP=2: full pipeline round-trip."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
# Shard: first split seq dim (dim=1) into CP=2 with zigzag ordering,
# then split hidden dim (dim=2) into TP=2.
natural_cp_chunks = list(full_tensor.chunk(4, dim=1))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([natural_cp_chunks[i] for i in zigzag_order], dim=1)
cp_chunks = list(zigzagged.chunk(2, dim=1))
tensors: list[torch.Tensor] = []
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=2))
for tp_rank in range(2):
tensors.append(tp_chunks[tp_rank])
parallel_infos.append(
{
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
}
)
dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
unshard_plans = compute_unshard_plan(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
reorder_plans = compute_reorder_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
all_plans = [*unshard_plans, *reorder_plans]
assert len(unshard_plans) == 2
assert len(reorder_plans) == 1
current: list[torch.Tensor] = tensors
for plan in all_plans:
if isinstance(plan, ReorderPlan):
current = execute_reorder_plan(plan, current)
else:
current = execute_unshard_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

10
test/registered/debug_utils/comparator/unshard/test_execute.py → test/registered/debug_utils/comparator/aligner/unshard/test_execute.py
View File

@@ -3,13 +3,15 @@ import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
from sglang.srt.debug_utils.comparator.unshard.executor import (
from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
_apply_unshard,
execute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)

6
test/registered/debug_utils/comparator/unshard/test_parallel_info.py → test/registered/debug_utils/comparator/aligner/unshard/test_parallel_info.py
View File

@@ -2,11 +2,11 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
from sglang.srt.debug_utils.comparator.unshard.parallel_info import (
from sglang.srt.debug_utils.comparator.aligner.unshard.parallel_info import (
normalize_parallel_info,
)
from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)

13
test/registered/debug_utils/comparator/unshard/test_plan.py → test/registered/debug_utils/comparator/aligner/unshard/test_plan.py
View File

@@ -2,9 +2,11 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -172,13 +174,14 @@ class TestComputeUnshardPlan:
with pytest.raises(NotImplementedError, match="reduction"):
compute_unshard_plan(dim_specs, parallel_infos)
def test_ordering_not_natural_raises(self) -> None:
def test_ordering_zigzag_accepted(self) -> None:
dim_specs = parse_dims("s(cp,zigzag)")
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
with pytest.raises(NotImplementedError, match="ordering"):
compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unshard_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.CP
def test_ordering_natural_accepted(self) -> None:
dim_specs = parse_dims("s(cp,natural)")

127
test/registered/debug_utils/comparator/test_entrypoint.py
View File

@@ -809,6 +809,74 @@ class TestEntrypointGroupingLogical:
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
def test_cp_zigzag_unshard(self, tmp_path, capsys):
"""CP=2 zigzag reorder is correctly undone through the full pipeline."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8, 6)
full_target = full_baseline + torch.randn(4, 8, 6) * 0.001
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
for side_dir, full_tensor in [
(baseline_dir, full_baseline),
(target_dir, full_target),
]:
_create_cp_zigzag_tp_sharded_dumps(
side_dir,
full_tensor=full_tensor,
name="attn_out",
cp_size=2,
tp_size=1,
seq_dim=1,
head_dim=2,
dims_str="b s(cp,zigzag) h",
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
def test_cp_zigzag_tp_unshard(self, tmp_path, capsys):
"""CP=2 zigzag + TP=2: multi-axis unshard with reorder through full pipeline."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8, 16)
full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
for side_dir, full_tensor in [
(baseline_dir, full_baseline),
(target_dir, full_target),
]:
_create_cp_zigzag_tp_sharded_dumps(
side_dir,
full_tensor=full_tensor,
name="hidden",
cp_size=2,
tp_size=2,
seq_dim=1,
head_dim=2,
dims_str="b s(cp,zigzag) h(tp)",
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
# --------------------------- Assertion helpers -------------------
@@ -1011,6 +1079,65 @@ def _create_ep_cp_tp_sharded_dumps(
return directory / _FIXED_EXP_NAME
def _create_cp_zigzag_tp_sharded_dumps(
directory: Path,
*,
full_tensor: torch.Tensor,
name: str,
cp_size: int,
tp_size: int,
seq_dim: int,
head_dim: int,
dims_str: str,
num_steps: int = 1,
) -> Path:
"""Create CP-zigzag (+optional TP) sharded dump files from a full tensor."""
num_chunks: int = cp_size * 2
natural_chunks: list[torch.Tensor] = list(
full_tensor.chunk(num_chunks, dim=seq_dim)
)
zigzag_order: list[int] = []
for i in range(cp_size):
zigzag_order.append(i)
zigzag_order.append(num_chunks - 1 - i)
zigzagged: torch.Tensor = torch.cat(
[natural_chunks[idx] for idx in zigzag_order], dim=seq_dim
)
cp_chunks: list[torch.Tensor] = list(zigzagged.chunk(cp_size, dim=seq_dim))
rank: int = 0
for cp_rank in range(cp_size):
tp_chunks: list[torch.Tensor] = (
list(cp_chunks[cp_rank].chunk(tp_size, dim=head_dim))
if tp_size > 1
else [cp_chunks[cp_rank]]
)
for tp_rank in range(tp_size):
parallel_info: dict[str, int] = {
"cp_rank": cp_rank,
"cp_size": cp_size,
}
if tp_size > 1:
parallel_info["tp_rank"] = tp_rank
parallel_info["tp_size"] = tp_size
_create_rank_dump(
directory,
rank=rank,
name=name,
tensor=tp_chunks[tp_rank],
dims=dims_str,
parallel_info=parallel_info,
num_steps=num_steps,
)
rank += 1
return directory / _FIXED_EXP_NAME
def _create_tp_sharded_dumps(
directory: Path,
*,