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

Support directory detection in dump comparator (#19680)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:45:35 +08:00
committed by GitHub
parent 6980416149
commit abdc0ee71f
12 changed files with 280 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
+5 -5
View File
@@ -924,7 +924,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").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
@@ -946,7 +946,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").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
@@ -980,7 +980,7 @@ class TestReduceSum:
}
)
dim_specs = parse_dims("b s(cp) h(tp:partial)")
dim_specs = parse_dims("b s[cp] h[tp:partial]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -1009,7 +1009,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").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
@@ -1022,7 +1022,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").dims
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
-53
View File
@@ -697,58 +697,5 @@ class TestComputeUnsharderPlanFusedDims:
assert isinstance(plans[0].params, ReduceSumParams)
class TestComputeUnsharderPlanFusedDims:
def test_fused_dim_tp2(self) -> None:
"""Fused dim "(num_heads*head_dim)[tp]" should unshard on the fused tensor name."""
dim_specs = parse_dims("t (num_heads*head_dim)[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
assert isinstance(plans[0].params, ConcatParams)
assert plans[0].params.dim_name == "num_heads___head_dim"
assert plans[0].groups == [[0, 1]]
def test_fused_dim_modifier_on_second_sub(self) -> None:
"""Modifier on fused dim: "(a*b)[tp]" should produce concat plan."""
dim_specs = parse_dims("t (a*b)[tp]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
assert isinstance(plans[0].params, ConcatParams)
assert plans[0].params.dim_name == "a___b"
def test_fused_dim_no_modifier(self) -> None:
"""Fused dim without any modifier should have no unshard plans (beyond replicated)."""
dim_specs = parse_dims("t (a*b)").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
# TP not annotated in dims → replicated → pick
assert len(plans) == 1
assert isinstance(plans[0].params, PickParams)
def test_fused_dim_with_reduction(self) -> None:
"""Fused dim with partial reduction: "(a*b)[tp:partial]"."""
dim_specs = parse_dims("t (a*b)[tp:partial]").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
assert isinstance(plans[0].params, ReduceSumParams)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))