Support singleton dimension squeezing in dump comparator (#19566)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,165 @@
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
|
||||
AxisAlignerPlan,
|
||||
compute_axis_aligner_plan,
|
||||
execute_axis_aligner_plan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestComputeAxisAlignerPlan:
|
||||
def test_no_dims_returns_none(self) -> None:
|
||||
assert compute_axis_aligner_plan(Pair(x=None, y=None)) is None
|
||||
assert compute_axis_aligner_plan(Pair(x="t h d", y=None)) is None
|
||||
assert compute_axis_aligner_plan(Pair(x=None, y="t h d")) is None
|
||||
|
||||
def test_same_order_returns_none(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t h d", y="t h d")
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_different_order(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t h d", y="t d h")
|
||||
)
|
||||
assert result is not None
|
||||
assert result.pattern.x == "t h d -> t d h"
|
||||
assert result.pattern.y is None
|
||||
|
||||
def test_name_mismatch_returns_none_with_warning(self) -> None:
|
||||
with warning_sink.context() as warnings:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t h d", y="t h e")
|
||||
)
|
||||
|
||||
assert result is None
|
||||
assert len(warnings) == 1
|
||||
assert warnings[0].category == "axis_aligner_dim_mismatch"
|
||||
assert "dim name sets differ" in warnings[0].message
|
||||
|
||||
def test_modifiers_ignored_for_name_extraction(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t h(tp) d", y="t d h(tp)")
|
||||
)
|
||||
assert result is not None
|
||||
assert result.pattern.x == "t h d -> t d h"
|
||||
|
||||
def test_squeeze_only_no_swap(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t 1 h", y="t h")
|
||||
)
|
||||
assert result is not None
|
||||
assert result.pattern.x == "t 1 h -> t h"
|
||||
assert result.pattern.y is None
|
||||
|
||||
def test_squeeze_both_sides(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t 1 h", y="1 t h")
|
||||
)
|
||||
assert result is not None
|
||||
assert result.pattern.x == "t 1 h -> t h"
|
||||
assert result.pattern.y == "1 t h -> t h"
|
||||
|
||||
def test_squeeze_plus_swap(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t 1 h d", y="t d h")
|
||||
)
|
||||
assert result is not None
|
||||
assert result.pattern.x == "t 1 h d -> t d h"
|
||||
assert result.pattern.y is None
|
||||
|
||||
def test_squeeze_y_only(self) -> None:
|
||||
result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
|
||||
Pair(x="t h", y="t 1 h")
|
||||
)
|
||||
assert result is not None
|
||||
assert result.pattern.x is None
|
||||
assert result.pattern.y == "t 1 h -> t h"
|
||||
|
||||
|
||||
class TestExecuteAxisAlignerPlan:
|
||||
def test_rearrange(self) -> None:
|
||||
torch.manual_seed(42)
|
||||
tensor: torch.Tensor = torch.randn(4, 8, 16).refine_names("t", "h", "d")
|
||||
plan = AxisAlignerPlan(
|
||||
pattern=Pair(x="t h d -> t d h", y=None),
|
||||
)
|
||||
|
||||
result: torch.Tensor = execute_axis_aligner_plan(
|
||||
tensor=tensor, plan=plan, side="x"
|
||||
)
|
||||
|
||||
assert result.shape == (4, 16, 8)
|
||||
for i in range(4):
|
||||
assert torch.equal(
|
||||
result[i],
|
||||
tensor.rename(None)[i].T,
|
||||
)
|
||||
|
||||
def test_execute_squeeze(self) -> None:
|
||||
torch.manual_seed(42)
|
||||
tensor: torch.Tensor = torch.randn(4, 1, 8).refine_names("t", "singleton0", "h")
|
||||
plan = AxisAlignerPlan(
|
||||
pattern=Pair(x="t 1 h -> t h", y=None),
|
||||
)
|
||||
|
||||
result: torch.Tensor = execute_axis_aligner_plan(
|
||||
tensor=tensor, plan=plan, side="x"
|
||||
)
|
||||
|
||||
assert result.shape == (4, 8)
|
||||
|
||||
def test_execute_squeeze_then_swap(self) -> None:
|
||||
torch.manual_seed(42)
|
||||
tensor: torch.Tensor = torch.randn(4, 1, 8, 16).refine_names(
|
||||
"t", "singleton0", "h", "d"
|
||||
)
|
||||
plan = AxisAlignerPlan(
|
||||
pattern=Pair(x="t 1 h d -> t d h", y=None),
|
||||
)
|
||||
|
||||
result: torch.Tensor = execute_axis_aligner_plan(
|
||||
tensor=tensor, plan=plan, side="x"
|
||||
)
|
||||
|
||||
assert result.shape == (4, 16, 8)
|
||||
|
||||
def test_execute_y_side(self) -> None:
|
||||
torch.manual_seed(42)
|
||||
tensor: torch.Tensor = torch.randn(4, 1, 8).refine_names("t", "singleton0", "h")
|
||||
plan = AxisAlignerPlan(
|
||||
pattern=Pair(x=None, y="t 1 h -> t h"),
|
||||
)
|
||||
|
||||
result: torch.Tensor = execute_axis_aligner_plan(
|
||||
tensor=tensor, plan=plan, side="y"
|
||||
)
|
||||
|
||||
assert result.shape == (4, 8)
|
||||
|
||||
def test_noop_side(self) -> None:
|
||||
torch.manual_seed(42)
|
||||
tensor: torch.Tensor = torch.randn(4, 8, 16).refine_names("t", "h", "d")
|
||||
plan = AxisAlignerPlan(
|
||||
pattern=Pair(x="t h d -> t d h", y=None),
|
||||
)
|
||||
|
||||
result: torch.Tensor = execute_axis_aligner_plan(
|
||||
tensor=tensor, plan=plan, side="y"
|
||||
)
|
||||
|
||||
assert result.shape == (4, 8, 16)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -639,5 +639,119 @@ class TestThdCpConcat:
|
||||
)
|
||||
|
||||
|
||||
class TestThdCpConcat:
|
||||
def test_single_seq(self) -> None:
|
||||
"""Single seq THD unshard: 2 ranks → per-seq concat."""
|
||||
rank0 = torch.tensor([1, 2, 3]).refine_names("t")
|
||||
rank1 = torch.tensor([4, 5, 6]).refine_names("t")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
expected = torch.tensor([1, 2, 3, 4, 5, 6])
|
||||
assert torch.equal(result[0].rename(None), expected)
|
||||
|
||||
def test_multi_seq(self) -> None:
|
||||
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
|
||||
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
|
||||
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
|
||||
seq_a_r0 = torch.arange(0, 50)
|
||||
seq_b_r0 = torch.arange(100, 132)
|
||||
pad_r0 = torch.full((46,), -1)
|
||||
rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
|
||||
|
||||
seq_a_r1 = torch.arange(50, 100)
|
||||
seq_b_r1 = torch.arange(132, 164)
|
||||
pad_r1 = torch.full((46,), -2)
|
||||
rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
unsharded: torch.Tensor = result[0].rename(None)
|
||||
|
||||
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
|
||||
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
# seqB: r0(32) + r1(32) = 64 tokens
|
||||
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
# pad: r0(46) + r1(46) = 92 tokens
|
||||
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
|
||||
|
||||
def test_with_hidden_dim(self) -> None:
|
||||
"""THD unshard with trailing hidden dim: shape [T, H]."""
|
||||
torch.manual_seed(42)
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
|
||||
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
|
||||
seq_a_r0 = torch.randn(3, hidden)
|
||||
seq_b_r0 = torch.randn(2, hidden)
|
||||
rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
|
||||
|
||||
seq_a_r1 = torch.randn(3, hidden)
|
||||
seq_b_r1 = torch.randn(2, hidden)
|
||||
rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
unsharded: torch.Tensor = result[0].rename(None)
|
||||
|
||||
assert unsharded.shape == (10, hidden)
|
||||
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
|
||||
def test_with_leading_batch_dim(self) -> None:
|
||||
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
|
||||
torch.manual_seed(42)
|
||||
batch: int = 2
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
|
||||
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
|
||||
seq_a_r0 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r0 = torch.randn(batch, 2, hidden)
|
||||
rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
|
||||
|
||||
seq_a_r1 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r1 = torch.randn(batch, 2, hidden)
|
||||
rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
with warning_sink.context():
|
||||
result = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(result) == 1
|
||||
unsharded: torch.Tensor = result[0].rename(None)
|
||||
|
||||
assert unsharded.shape == (batch, 10, hidden)
|
||||
# seqA: r0(3) + r1(3) = 6 tokens per batch
|
||||
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
|
||||
# seqB: r0(2) + r1(2) = 4 tokens per batch
|
||||
assert torch.equal(
|
||||
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
Reference in New Issue
Block a user