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

Support multi sharding group on the same dimension in dump comparator (#19601)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:36:48 +08:00
committed by GitHub
parent 46960e65cf
commit ea6ff7b01f
14 changed files with 469 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py
+7 -7
View File
@@ -85,8 +85,8 @@ class TestComputePerStepSubPlans:
def test_zigzag_returns_both_plans(self) -> None:
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
metas=[
_make_meta(dims="b s(cp,zigzag) h", cp_rank=0, cp_size=2),
_make_meta(dims="b s(cp,zigzag) h", cp_rank=1, cp_size=2),
_make_meta(dims="b s(cp:zigzag) h", cp_rank=0, cp_size=2),
_make_meta(dims="b s(cp:zigzag) h", cp_rank=1, cp_size=2),
]
)
unsharder_plans: list[UnsharderPlan] = [
@@ -177,33 +177,33 @@ class TestComputeAlignerPlan:
class TestComputePerStepSubPlansThd:
def test_thd_zigzag_returns_thd_plans(self) -> None:
"""t(cp,zigzag) h(tp) generates THD-typed unsharder + reorderer plans."""
"""t(cp:zigzag) h(tp) generates THD-typed unsharder + reorderer plans."""
thd_global_seq_lens: list[int] = [100, 64, 92]
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
metas=[
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=0,
cp_size=2,
tp_rank=0,
tp_size=2,
),
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=0,
cp_size=2,
tp_rank=1,
tp_size=2,
),
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=1,
cp_size=2,
tp_rank=0,
tp_size=2,
),
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=1,
cp_size=2,
tp_rank=1,
test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py
+91 -13
View File
@@ -25,8 +25,8 @@ register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestComputeReordererPlans:
def test_compute_reorderer_plans_zigzag(self) -> None:
"""s(cp,zigzag) produces a ReordererPlan."""
dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
"""s(cp:zigzag) produces a ReordererPlan."""
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),
@@ -43,8 +43,8 @@ class TestComputeReordererPlans:
assert plans[0].params.cp_size == 2
def test_compute_reorderer_plans_thd_zigzag(self) -> None:
"""t(cp,zigzag) produces a ZigzagToNaturalThdParams plan."""
dim_specs = parse_dims("t(cp,zigzag) h(tp)")
"""t(cp:zigzag) produces a ZigzagToNaturalThdParams plan."""
dim_specs = parse_dims("t(cp:zigzag) h(tp)")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -64,8 +64,8 @@ class TestComputeReordererPlans:
assert plans[0].params.seq_lens == [100, 64, 92]
def test_non_seq_dim_still_raises(self) -> None:
"""Zigzag on non-sequence/non-token dim (e.g. h(cp,zigzag)) raises ValueError."""
dim_specs = parse_dims("h(cp,zigzag) d")
"""Zigzag on non-sequence/non-token dim (e.g. h(cp:zigzag)) raises ValueError."""
dim_specs = parse_dims("h(cp:zigzag) d")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)},
]
@@ -73,8 +73,8 @@ class TestComputeReordererPlans:
compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
def test_thd_zigzag_without_seq_lens_raises(self) -> None:
"""t(cp,zigzag) without thd_global_seq_lens raises ValueError."""
dim_specs = parse_dims("t(cp,zigzag) h(tp)")
"""t(cp:zigzag) without thd_global_seq_lens 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),
@@ -85,8 +85,8 @@ class TestComputeReordererPlans:
compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
def test_thd_natural_no_reorder(self) -> None:
"""t(cp,natural) and t(cp) produce no reorder plans."""
for dims_str in ["t(cp,natural) h(tp)", "t(cp) h(tp)"]:
"""t(cp:natural) and t(cp) produce no reorder plans."""
for dims_str in ["t(cp:natural) h(tp)", "t(cp) h(tp)"]:
dim_specs = parse_dims(dims_str)
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
@@ -100,8 +100,8 @@ class TestComputeReordererPlans:
assert plans == []
def test_compute_reorderer_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)"]:
"""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]] = [
{
@@ -141,7 +141,7 @@ class TestCpZigzagTpE2E:
}
)
dim_specs: list[DimSpec] = parse_dims("b s(cp,zigzag) h(tp)")
dim_specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp)")
dim_names: list[str] = [s.name for s in dim_specs]
unsharder_plans = compute_unsharder_plan(
@@ -166,5 +166,83 @@ class TestCpZigzagTpE2E:
assert torch.allclose(current[0].rename(None), full_tensor)
class TestCpZigzagSpSameDimE2E:
"""E2E test for t(cp:zigzag,sp) — two axes sharding the same token dim."""
def test_cp2_sp2_zigzag_e2e(self) -> None:
"""CP=2 zigzag + SP=2 on same token dim: full unshard + reorder round-trip.
Shard order (outer to inner, matching left-to-right in dims annotation):
1. CP zigzag splits token dim into 2 CP chunks (zigzag order)
2. SP splits each CP chunk into 2 SP chunks
Unshard order (inner to outer, right-to-left):
1. SP concat (inner): merge SP chunks back
2. CP concat (outer): merge CP chunks back
3. Zigzag reorder: restore natural token order
"""
torch.manual_seed(42)
total_tokens: int = 16
hidden: int = 8
full_tensor: torch.Tensor = torch.randn(total_tokens, hidden)
# Step 1: CP zigzag split — split into 2*cp_size=4 natural chunks, reorder by zigzag
cp_size: int = 2
sp_size: int = 2
n_natural_chunks: int = cp_size * 2
natural_chunks: list[torch.Tensor] = list(
full_tensor.chunk(n_natural_chunks, dim=0)
)
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged: torch.Tensor = torch.cat(
[natural_chunks[i] for i in zigzag_order], dim=0
)
cp_chunks: list[torch.Tensor] = list(zigzagged.chunk(cp_size, dim=0))
# Step 2: SP split within each CP chunk
tensors: list[torch.Tensor] = []
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(cp_size):
sp_chunks: list[torch.Tensor] = list(
cp_chunks[cp_rank].chunk(sp_size, dim=0)
)
for sp_rank in range(sp_size):
tensors.append(sp_chunks[sp_rank])
parallel_infos.append(
{
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=cp_size),
ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=sp_size),
}
)
dim_specs: list[DimSpec] = parse_dims("t(cp:zigzag,sp) h")
dim_names: list[str] = [s.name for s in dim_specs]
unsharder_plans = compute_unsharder_plan(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
reorderer_plans = compute_reorderer_plans(
dim_specs=dim_specs,
parallel_infos=parallel_infos,
thd_global_seq_lens=[total_tokens],
)
all_plans = [*unsharder_plans, *reorderer_plans]
assert len(unsharder_plans) == 2 # SP concat, CP concat
assert unsharder_plans[0].axis == ParallelAxis.SP
assert unsharder_plans[1].axis == ParallelAxis.CP
assert len(reorderer_plans) == 1 # zigzag reorder
current: list[torch.Tensor] = [t.refine_names(*dim_names) for t in tensors]
for plan in all_plans:
if isinstance(plan, ReordererPlan):
current = execute_reorderer_plan(plan, current)
else:
current = execute_unsharder_plan(plan, current).tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py
+9 -9
View File
@@ -90,7 +90,7 @@ class TestEnsureDimsInMetas:
assert result is metas
def test_cp_sharded_sglang_input_ids_infers_dims(self):
"""CP + input_ids in sglang infers dims 't(cp,zigzag)'."""
"""CP + input_ids in sglang infers dims 't(cp:zigzag)'."""
metas: list[dict] = [
self._make_meta(cp_size=2, cp_rank=0),
self._make_meta(cp_size=2, cp_rank=1),
@@ -99,11 +99,11 @@ class TestEnsureDimsInMetas:
name="input_ids", plugin=_sglang_plugin, metas=metas, ndim=1
)
assert result is not metas
assert result[0]["dims"] == "t(cp,zigzag)"
assert result[1]["dims"] == "t(cp,zigzag)"
assert result[0]["dims"] == "t(cp:zigzag)"
assert result[1]["dims"] == "t(cp:zigzag)"
def test_cp_sharded_sglang_positions_infers_dims(self):
"""CP + positions in sglang infers dims 't(cp,zigzag)'."""
"""CP + positions in sglang infers dims 't(cp:zigzag)'."""
metas: list[dict] = [
self._make_meta(cp_size=2, cp_rank=0),
self._make_meta(cp_size=2, cp_rank=1),
@@ -111,10 +111,10 @@ class TestEnsureDimsInMetas:
result = _ensure_dims_in_metas(
name="positions", plugin=_sglang_plugin, metas=metas, ndim=1
)
assert result[0]["dims"] == "t(cp,zigzag)"
assert result[0]["dims"] == "t(cp:zigzag)"
def test_cp_sharded_megatron_input_ids_infers_dims_1d(self):
"""CP + input_ids in megatron (1D) infers dims 't(cp,zigzag)'."""
"""CP + input_ids in megatron (1D) infers dims 't(cp:zigzag)'."""
metas: list[dict] = [
{"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
{"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
@@ -122,10 +122,10 @@ class TestEnsureDimsInMetas:
result = _ensure_dims_in_metas(
name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=1
)
assert result[0]["dims"] == "t(cp,zigzag)"
assert result[0]["dims"] == "t(cp:zigzag)"
def test_cp_sharded_megatron_input_ids_infers_dims_2d(self):
"""CP + input_ids in megatron (2D) infers dims 'b s(cp,zigzag)'."""
"""CP + input_ids in megatron (2D) infers dims 'b s(cp:zigzag)'."""
metas: list[dict] = [
{"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
{"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
@@ -133,7 +133,7 @@ class TestEnsureDimsInMetas:
result = _ensure_dims_in_metas(
name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=2
)
assert result[0]["dims"] == "b s(cp,zigzag)"
assert result[0]["dims"] == "b s(cp:zigzag)"
def test_cp_non_sharded_name_returns_metas_unchanged(self):
"""CP + non-sharded tensor name (seq_lens) returns metas as-is."""
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py
+6 -6
View File
@@ -218,19 +218,19 @@ class TestInferCpShardedDims:
"""Tests for infer_cp_sharded_dims on each plugin."""
def test_megatron_infer_1d(self) -> None:
"""Megatron 1D → 't(cp,zigzag)'."""
"""Megatron 1D → 't(cp:zigzag)'."""
result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1)
assert result == "t(cp,zigzag)"
assert result == "t(cp:zigzag)"
def test_megatron_infer_2d(self) -> None:
"""Megatron 2D → 'b s(cp,zigzag)'."""
"""Megatron 2D → 'b s(cp:zigzag)'."""
result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=2)
assert result == "b s(cp,zigzag)"
assert result == "b s(cp:zigzag)"
def test_sglang_infer_1d(self) -> None:
"""SGLang 1D → 't(cp,zigzag)'."""
"""SGLang 1D → 't(cp:zigzag)'."""
result: str = _sglang_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1)
assert result == "t(cp,zigzag)"
assert result == "t(cp:zigzag)"
def test_megatron_infer_3d_raises(self) -> None:
"""Megatron 3D raises ValueError."""
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
+5 -5
View File
@@ -654,7 +654,7 @@ class TestReduceSum:
part_a = full_tensor * 0.6
part_b = full_tensor * 0.4
dim_specs = parse_dims("h(tp,partial) d")
dim_specs = parse_dims("h(tp:partial) d")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -676,7 +676,7 @@ class TestReduceSum:
full_tensor = torch.randn(4, 8)
parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]
dim_specs = parse_dims("h(tp,partial) d")
dim_specs = parse_dims("h(tp:partial) d")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -710,7 +710,7 @@ class TestReduceSum:
}
)
dim_specs = parse_dims("b s(cp) h(tp,partial)")
dim_specs = parse_dims("b s(cp) h(tp:partial)")
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -739,7 +739,7 @@ class TestReduceSum:
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
]
dim_specs = parse_dims("h(tp,partial) d")
dim_specs = parse_dims("h(tp:partial) d")
plans = compute_unsharder_plan(dim_specs, parallel_infos)
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
@@ -752,7 +752,7 @@ class TestReduceSum:
def test_reduce_preserves_named_dims(self) -> None:
"""Named tensor dimensions are preserved through reduce_sum."""
dim_specs = parse_dims("h(tp,partial) d")
dim_specs = parse_dims("h(tp:partial) d")
part_a = torch.randn(4, 8).refine_names("h", "d")
part_b = torch.randn(4, 8).refine_names("h", "d")
test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py
+77 -6
View File
@@ -175,7 +175,7 @@ class TestComputeUnsharderPlan:
compute_unsharder_plan(dim_specs, parallel_infos)
def test_reduction_partial_returns_reduce_sum(self) -> None:
dim_specs = parse_dims("h(tp,partial)")
dim_specs = parse_dims("h(tp:partial)")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -188,7 +188,7 @@ class TestComputeUnsharderPlan:
def test_reduction_partial_tp4(self) -> None:
"""TP=4 with partial reduction produces a single ReduceSumParams step."""
dim_specs = parse_dims("h(tp,partial)")
dim_specs = parse_dims("h(tp:partial)")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -200,7 +200,7 @@ class TestComputeUnsharderPlan:
def test_multi_axis_with_reduction_on_one(self) -> None:
"""CP concat + TP reduce produces a 2-step plan."""
dim_specs = parse_dims("s(cp) h(tp,partial)")
dim_specs = parse_dims("s(cp) h(tp:partial)")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for tp_rank in range(2):
@@ -221,7 +221,7 @@ class TestComputeUnsharderPlan:
def test_reduction_scrambled_ranks(self) -> None:
"""Scrambled world_rank order with partial reduction."""
dim_specs = parse_dims("h(tp,partial)")
dim_specs = parse_dims("h(tp:partial)")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
@@ -235,7 +235,7 @@ class TestComputeUnsharderPlan:
assert plans[0].groups == [[1, 3, 0, 2]]
def test_ordering_zigzag_accepted(self) -> None:
dim_specs = parse_dims("s(cp,zigzag)")
dim_specs = parse_dims("s(cp:zigzag)")
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -244,7 +244,7 @@ class TestComputeUnsharderPlan:
assert plans[0].axis == ParallelAxis.CP
def test_ordering_natural_accepted(self) -> None:
dim_specs = parse_dims("s(cp,natural)")
dim_specs = parse_dims("s(cp:natural)")
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -288,6 +288,77 @@ class TestComputeUnsharderPlan:
assert len(plans[2].groups) == 1
assert len(plans[2].groups[0]) == 2
def test_same_dim_cp_sp_plan(self) -> None:
"""t(cp:zigzag,sp) with CP=2 SP=2: SP unshards first (inner), then CP."""
dim_specs = parse_dims("t(cp:zigzag,sp) 1 h")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for sp_rank in range(2):
parallel_infos.append(
{
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=2),
}
)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
# SP unshards first (rightmost modifier = innermost shard)
sp_plan = plans[0]
assert sp_plan.axis == ParallelAxis.SP
assert isinstance(sp_plan.params, ConcatParams)
assert sp_plan.params.dim_name == "t"
assert len(sp_plan.groups) == 2
for group in sp_plan.groups:
assert len(group) == 2
# CP unshards second (leftmost modifier = outermost shard)
cp_plan = plans[1]
assert cp_plan.axis == ParallelAxis.CP
assert isinstance(cp_plan.params, ConcatParams)
assert cp_plan.params.dim_name == "t"
assert len(cp_plan.groups) == 1
assert len(cp_plan.groups[0]) == 2
def test_same_dim_cp_sp_with_thd(self) -> None:
"""t(cp:zigzag,sp) with THD: SP → ConcatParams, CP → CpThdConcatParams."""
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
CpThdConcatParams,
)
dim_specs = parse_dims("t(cp:zigzag,sp) h")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for sp_rank in range(2):
parallel_infos.append(
{
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=2),
}
)
thd_global_seq_lens: list[int] = [100, 64]
plans = compute_unsharder_plan(
dim_specs, parallel_infos, thd_global_seq_lens=thd_global_seq_lens
)
assert len(plans) == 2
# SP unshards first: plain concat (SP is not CP, no THD special handling)
sp_plan = plans[0]
assert sp_plan.axis == ParallelAxis.SP
assert isinstance(sp_plan.params, ConcatParams)
assert sp_plan.params.dim_name == "t"
# CP unshards second: THD concat because dim is 't' + axis is CP + thd_global_seq_lens provided
cp_plan = plans[1]
assert cp_plan.axis == ParallelAxis.CP
assert isinstance(cp_plan.params, CpThdConcatParams)
assert cp_plan.params.dim_name == "t"
assert cp_plan.params.seq_lens_per_rank == [50, 32]
def test_sp_in_dims_but_not_in_parallel_info(self) -> None:
"""s(sp) in dims but SP absent from parallel_info (SP disabled), should auto-skip."""
dim_specs = parse_dims("s(sp) b h(tp)")