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

Enhance replicated tensor checker in dump comparator (#19597)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:34:34 +08:00
committed by GitHub
parent ec08240a6a
commit e41164af1c
15 changed files with 514 additions and 313 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py
+13 -14
View File
@@ -1,5 +1,4 @@
import sys
from typing import Optional
import pytest
import torch
@@ -32,22 +31,25 @@ register_cpu_ci(est_time=15, suite="default", nightly=True)
class TestExecuteSubPlans:
def test_empty_tensors_returns_none(self) -> None:
result: Optional[torch.Tensor] = execute_sub_plans(tensors=[], plans=[])
result, checks = execute_sub_plans(tensors=[], plans=[])
assert result is None
assert checks == []
def test_no_plans_single_tensor_passthrough(self) -> None:
tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
result: Optional[torch.Tensor] = execute_sub_plans(tensors=[tensor], plans=[])
result, checks = execute_sub_plans(tensors=[tensor], plans=[])
assert result is not None
assert torch.equal(result, tensor)
assert checks == []
def test_no_plans_multiple_tensors_returns_none(self) -> None:
tensors: list[torch.Tensor] = [
torch.tensor([1.0]),
torch.tensor([2.0]),
]
result: Optional[torch.Tensor] = execute_sub_plans(tensors=tensors, plans=[])
result, checks = execute_sub_plans(tensors=tensors, plans=[])
assert result is None
assert checks == []
def test_with_unsharder_plan(self) -> None:
t0: torch.Tensor = torch.tensor([[1.0, 2.0]]).refine_names("b", "h")
@@ -59,13 +61,12 @@ class TestExecuteSubPlans:
groups=[[0, 1]],
)
result: Optional[torch.Tensor] = execute_sub_plans(
tensors=[t0, t1], plans=[plan]
)
result, checks = execute_sub_plans(tensors=[t0, t1], plans=[plan])
assert result is not None
expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
assert torch.equal(result.rename(None), expected)
assert checks == []
class TestExecuteSubPlan:
@@ -74,7 +75,7 @@ class TestExecuteSubPlan:
pass
with pytest.raises(NotImplementedError, match="Unknown"):
execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan()) # type: ignore[arg-type]
execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan())
class TestExecuteStepPlans:
@@ -90,11 +91,10 @@ class TestExecuteStepPlans:
sub_plans=[],
)
result: dict[int, torch.Tensor] = _execute_step_plans(
tensors=tensors, step_plans=[step_plan]
)
result, checks = _execute_step_plans(tensors=tensors, step_plans=[step_plan])
assert result == {}
assert checks == []
def test_single_step_passthrough(self) -> None:
tensor: torch.Tensor = torch.tensor([1.0, 2.0])
@@ -105,12 +105,11 @@ class TestExecuteStepPlans:
sub_plans=[],
)
result: dict[int, torch.Tensor] = _execute_step_plans(
tensors=[tensor], step_plans=[step_plan]
)
result, checks = _execute_step_plans(tensors=[tensor], step_plans=[step_plan])
assert 5 in result
assert torch.equal(result[5], tensor)
assert checks == []
class TestExecuteAlignerPlan:
test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py
+4 -9
View File
@@ -20,7 +20,6 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
UnsharderPlan,
)
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
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=10, suite="default", nightly=True)
@@ -227,10 +226,8 @@ class TestThdCpZigzagE2E:
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=seq_lens_per_rank),
groups=[[0, 1]],
)
with warning_sink.context():
unsharded: list[torch.Tensor] = execute_unsharder_plan(
unshard_plan, rank_tensors
)
unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors)
unsharded: list[torch.Tensor] = unsharder_result.tensors
assert len(unsharded) == 1
# Step 2: THD reorder
@@ -266,10 +263,8 @@ class TestThdCpZigzagE2E:
),
groups=[list(range(cp_size))],
)
with warning_sink.context():
unsharded: list[torch.Tensor] = execute_unsharder_plan(
unshard_plan, rank_tensors
)
unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors)
unsharded: list[torch.Tensor] = unsharder_result.tensors
assert len(unsharded) == 1
# Step 2: THD reorder
test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py
+5 -7
View File
@@ -18,7 +18,6 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis, parse_dims
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=10, suite="default", nightly=True)
@@ -157,12 +156,11 @@ class TestCpZigzagTpE2E:
assert len(reorderer_plans) == 1
current: list[torch.Tensor] = [t.refine_names(*dim_names) for t in tensors]
with warning_sink.context():
for plan in all_plans:
if isinstance(plan, ReordererPlan):
current = execute_reorderer_plan(plan, current)
else:
current = execute_unsharder_plan(plan, current)
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)
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
+163 -151
View File
@@ -4,6 +4,7 @@ import pytest
import torch
from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
UnsharderResult,
_apply_unshard,
_verify_replicated_group,
execute_unsharder_plan,
@@ -23,7 +24,7 @@ from sglang.srt.debug_utils.comparator.dims import (
ParallelAxis,
parse_dims,
)
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -49,11 +50,12 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 1
named_shards: list[torch.Tensor] = _name_tensors(shards, dim_specs)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], named_shards)
assert len(result) == 1
assert torch.allclose(result[0].rename(None), full_tensor)
assert warnings == []
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_shards
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
assert unsharder_result.replicated_checks == []
def test_scrambled_world_ranks_correct_result(self) -> None:
full_tensor = torch.randn(4, 8)
@@ -79,11 +81,12 @@ class TestExecuteUnsharderPlan:
dim_specs,
)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], tensors_ordered_by_world_rank)
assert len(result) == 1
assert torch.allclose(result[0].rename(None), full_tensor)
assert warnings == []
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], tensors_ordered_by_world_rank
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
assert unsharder_result.replicated_checks == []
def test_single_step_reduces_tensor_count(self) -> None:
"""8 tensors with 2 groups of 4 produce 2 output tensors."""
@@ -113,13 +116,15 @@ class TestExecuteUnsharderPlan:
tensors.append(source[tp_rank])
named_tensors: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
intermediate = execute_unsharder_plan(plans[0], named_tensors)
assert len(intermediate) == 4
intermediate_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_tensors
)
assert len(intermediate_result.tensors) == 4
with warning_sink.context():
final = execute_unsharder_plan(plans[1], intermediate)
assert len(final) == 1
final_result: UnsharderResult = execute_unsharder_plan(
plans[1], intermediate_result.tensors
)
assert len(final_result.tensors) == 1
def test_cp_tp_concat(self) -> None:
"""CP=2 + TP=2: multi-step unshard reconstructs original tensor."""
@@ -145,9 +150,9 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 2
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
@@ -187,9 +192,9 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 2
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
@@ -241,9 +246,9 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 3
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
@@ -290,9 +295,9 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 3
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
@@ -312,11 +317,12 @@ class TestPickOperation:
assert len(plans) == 1
assert isinstance(plans[0].params, PickParams)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], [tensor, tensor.clone()])
assert len(result) == 1
assert torch.allclose(result[0].rename(None), tensor)
assert warnings == []
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], [tensor, tensor.clone()]
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), tensor)
assert all(c.passed for c in unsharder_result.replicated_checks)
def test_pick_multiple_groups(self) -> None:
"""PickParams with multiple groups picks one from each."""
@@ -348,10 +354,11 @@ class TestPickOperation:
tensor = torch.randn(4)
tensors = [tensor.clone() for _ in range(4)]
with warning_sink.context() as warnings:
result = execute_unsharder_plan(pick_plans[0], tensors)
assert len(result) == 2
assert warnings == []
unsharder_result: UnsharderResult = execute_unsharder_plan(
pick_plans[0], tensors
)
assert len(unsharder_result.tensors) == 2
assert all(c.passed for c in unsharder_result.replicated_checks)
def test_replicated_tp_sharded_cp_e2e(self) -> None:
"""CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip."""
@@ -376,9 +383,9 @@ class TestPickOperation:
assert len(plans) == 2
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
@@ -406,44 +413,44 @@ class TestPickOperation:
assert all(isinstance(p.params, PickParams) for p in plans)
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
class TestVerifyReplicatedGroup:
def test_warns_on_mismatch(self) -> None:
"""_verify_replicated_group produces warning when replicas differ."""
def test_fails_on_mismatch(self) -> None:
"""_verify_replicated_group returns failed check when replicas differ."""
tensor_a = torch.ones(4)
tensor_b = torch.ones(4) + 0.1
with warning_sink.context() as warnings:
_verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(warnings) == 1
assert warnings[0].axis == "tp"
assert warnings[0].group_index == 0
assert warnings[0].differing_index == 1
assert warnings[0].baseline_index == 0
assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(checks) == 1
assert checks[0].axis == "tp"
assert checks[0].group_index == 0
assert checks[0].compared_index == 1
assert checks[0].baseline_index == 0
assert not checks[0].passed
assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
def test_no_warn_when_identical(self) -> None:
"""_verify_replicated_group produces no warning for identical replicas."""
def test_passes_when_identical(self) -> None:
"""_verify_replicated_group returns passed check for identical replicas."""
tensor = torch.randn(4, 8)
with warning_sink.context() as warnings:
_verify_replicated_group(
[tensor, tensor.clone()],
axis=ParallelAxis.TP,
group_index=0,
)
assert warnings == []
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[tensor, tensor.clone()],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(checks) == 1
assert checks[0].passed
def test_multiple_mismatches(self) -> None:
"""_verify_replicated_group reports each differing replica."""
@@ -451,19 +458,20 @@ class TestVerifyReplicatedGroup:
other_a = torch.ones(4)
other_b = torch.ones(4) * 2
with warning_sink.context() as warnings:
_verify_replicated_group(
[baseline, other_a, other_b],
axis=ParallelAxis.CP,
group_index=1,
)
assert len(warnings) == 2
assert warnings[0].differing_index == 1
assert warnings[1].differing_index == 2
assert warnings[1].max_abs_diff == pytest.approx(2.0, abs=1e-5)
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[baseline, other_a, other_b],
axis=ParallelAxis.CP,
group_index=1,
)
assert len(checks) == 2
assert checks[0].compared_index == 1
assert not checks[0].passed
assert checks[1].compared_index == 2
assert not checks[1].passed
assert checks[1].diff.max_abs_diff == pytest.approx(2.0, abs=1e-5)
def test_execute_returns_warnings(self) -> None:
"""execute_unsharder_plan emits warnings for replicated mismatch."""
def test_execute_returns_replicated_checks(self) -> None:
"""execute_unsharder_plan returns replicated checks for mismatch."""
dim_specs = parse_dims("h d")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
@@ -474,56 +482,58 @@ class TestVerifyReplicatedGroup:
tensor_a = torch.zeros(4)
tensor_b = torch.ones(4)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], [tensor_a, tensor_b])
assert len(result) == 1
assert len(warnings) == 1
assert torch.allclose(result[0].rename(None), tensor_a)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], [tensor_a, tensor_b]
)
assert len(unsharder_result.tensors) == 1
assert len(unsharder_result.replicated_checks) == 1
assert not unsharder_result.replicated_checks[0].passed
assert torch.allclose(unsharder_result.tensors[0].rename(None), tensor_a)
def test_atol_boundary_within(self) -> None:
"""Difference exactly at atol (1e-6) -> torch.allclose passes -> no warning."""
"""Difference exactly at atol (1e-6) -> passed."""
baseline = torch.zeros(4)
other = torch.full((4,), 1e-6)
with warning_sink.context() as warnings:
_verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
assert warnings == []
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(checks) == 1
assert checks[0].passed
def test_atol_boundary_exceeded(self) -> None:
"""Difference just above atol (1e-6 + 1e-9) -> torch.allclose fails -> warning."""
"""Difference just above atol (1e-6 + 1e-9) -> failed."""
baseline = torch.zeros(4)
other = torch.full((4,), 1e-6 + 1e-9)
with warning_sink.context() as warnings:
_verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(warnings) == 1
assert warnings[0].differing_index == 1
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(checks) == 1
assert not checks[0].passed
assert checks[0].compared_index == 1
def test_recompute_pseudo_mismatch_warns(self) -> None:
"""_verify_replicated_group produces warning for RECOMPUTE_PSEUDO axis mismatch."""
def test_recompute_pseudo_mismatch(self) -> None:
"""_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch."""
tensor_a = torch.ones(4)
tensor_b = torch.ones(4) + 0.1
with warning_sink.context() as warnings:
_verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.RECOMPUTE_PSEUDO,
group_index=0,
)
assert len(warnings) == 1
assert warnings[0].axis == "recompute_pseudo"
assert warnings[0].group_index == 0
assert warnings[0].differing_index == 1
assert warnings[0].baseline_index == 0
assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.RECOMPUTE_PSEUDO,
group_index=0,
)
assert len(checks) == 1
assert checks[0].axis == "recompute_pseudo"
assert checks[0].group_index == 0
assert checks[0].compared_index == 1
assert checks[0].baseline_index == 0
assert not checks[0].passed
assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
class TestThdCpConcat:
@@ -537,12 +547,11 @@ class TestThdCpConcat:
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
assert len(unsharder_result.tensors) == 1
expected = torch.tensor([1, 2, 3, 4, 5, 6])
assert torch.equal(result[0].rename(None), expected)
assert torch.equal(unsharder_result.tensors[0].rename(None), expected)
def test_multi_seq(self) -> None:
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
@@ -563,11 +572,10 @@ class TestThdCpConcat:
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])
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert len(unsharder_result.tensors) == 1
unsharded: torch.Tensor = unsharder_result.tensors[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]))
@@ -595,11 +603,10 @@ class TestThdCpConcat:
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])
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert len(unsharder_result.tensors) == 1
unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None)
assert unsharded.shape == (10, hidden)
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
@@ -625,11 +632,10 @@ class TestThdCpConcat:
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])
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert len(unsharder_result.tensors) == 1
unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None)
assert unsharded.shape == (batch, 10, hidden)
# seqA: r0(3) + r1(3) = 6 tokens per batch
@@ -657,11 +663,12 @@ class TestReduceSum:
assert isinstance(plans[0].params, ReduceSumParams)
named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
with warning_sink.context():
result = execute_unsharder_plan(plans[0], named_parts)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_parts
)
assert len(result) == 1
assert torch.allclose(result[0].rename(None), full_tensor)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
def test_tp4_reduce(self) -> None:
"""4 partial tensors sum to full tensor."""
@@ -677,11 +684,12 @@ class TestReduceSum:
assert len(plans) == 1
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
with warning_sink.context():
result = execute_unsharder_plan(plans[0], named_parts)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_parts
)
assert len(result) == 1
assert torch.allclose(result[0].rename(None), full_tensor)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
def test_multi_axis_concat_then_reduce(self) -> None:
"""CP concat + TP reduce end-to-end."""
@@ -707,9 +715,9 @@ class TestReduceSum:
assert len(plans) == 2
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(current[0].rename(None), full_tensor)
@@ -735,11 +743,12 @@ class TestReduceSum:
plans = compute_unsharder_plan(dim_specs, parallel_infos)
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
with warning_sink.context():
result = execute_unsharder_plan(plans[0], named_parts)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_parts
)
assert len(result) == 1
assert torch.allclose(result[0].rename(None), full_tensor)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
def test_reduce_preserves_named_dims(self) -> None:
"""Named tensor dimensions are preserved through reduce_sum."""
@@ -752,13 +761,16 @@ class TestReduceSum:
params=ReduceSumParams(),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [part_a, part_b])
unsharder_result: UnsharderResult = execute_unsharder_plan(
plan, [part_a, part_b]
)
assert len(result) == 1
assert result[0].names == ("h", "d")
assert len(unsharder_result.tensors) == 1
assert unsharder_result.tensors[0].names == ("h", "d")
expected = (part_a.rename(None) + part_b.rename(None)).refine_names("h", "d")
assert torch.allclose(result[0].rename(None), expected.rename(None))
assert torch.allclose(
unsharder_result.tensors[0].rename(None), expected.rename(None)
)
if __name__ == "__main__":
test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
+8 -8
View File
@@ -6,9 +6,9 @@ import torch
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
QUANTILE_NUMEL_THRESHOLD,
SAMPLE_DIFF_THRESHOLD,
_compute_diff,
_compute_tensor_stats,
compare_tensor_pair,
compute_diff,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import DiffInfo
from sglang.test.ci.ci_register import register_cpu_ci
@@ -55,7 +55,7 @@ class TestComputeTensorStats:
class TestComputeDiff:
def test_identical_tensors(self):
x = torch.ones(10, 10)
diff = _compute_diff(x_baseline=x, x_target=x)
diff = compute_diff(x_baseline=x, x_target=x)
assert diff.rel_diff == pytest.approx(0.0, abs=1e-5)
assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
@@ -70,7 +70,7 @@ class TestComputeDiff:
y = x.clone()
y[3, 7] = 1.5
diff = _compute_diff(x_baseline=x, x_target=y)
diff = compute_diff(x_baseline=x, x_target=y)
assert diff.max_abs_diff == pytest.approx(0.5, abs=1e-4)
assert diff.max_diff_coord == [3, 7]
@@ -85,14 +85,14 @@ class TestComputeDiff:
def test_large_tensor_skips_diff_quantiles(self):
x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1)
y = x + 0.001
diff = _compute_diff(x_baseline=x, x_target=y)
diff = compute_diff(x_baseline=x, x_target=y)
assert diff.abs_diff_percentiles == {}
def test_rel_diff_value(self):
x = torch.tensor([1.0, 0.0])
y = torch.tensor([0.0, 1.0])
diff = _compute_diff(x_baseline=x, x_target=y)
diff = compute_diff(x_baseline=x, x_target=y)
assert diff.rel_diff == pytest.approx(1.0, abs=1e-5)
assert diff.passed is False
@@ -103,7 +103,7 @@ class TestComputeDiff:
x: torch.Tensor = torch.randn(8, 16)
y: torch.Tensor = x + torch.randn_like(x) * 0.01
diff: DiffInfo = _compute_diff(
diff: DiffInfo = compute_diff(
x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0
)
@@ -117,7 +117,7 @@ class TestComputeDiff:
x: torch.Tensor = torch.randn(8, 16)
y: torch.Tensor = x + torch.randn_like(x) * 0.01
diff: DiffInfo = _compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3)
diff: DiffInfo = compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3)
assert diff.per_token_rel_diff is None
@@ -127,7 +127,7 @@ class TestComputeDiff:
x: torch.Tensor = torch.randn(4, 8)
y: torch.Tensor = x + torch.randn_like(x) * 0.01
diff: DiffInfo = _compute_diff(
diff: DiffInfo = compute_diff(
x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0
)
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
+54 -27
View File
@@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
ConfigRecord,
GeneralWarning,
ReplicatedMismatchWarning,
ReplicatedCheckResult,
SkipRecord,
SummaryRecord,
WarningRecord,
@@ -126,16 +126,24 @@ class TestRecordTypes:
assert restored == record
def _make_warning(**overrides) -> ReplicatedMismatchWarning:
def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
defaults: dict = dict(
axis="tp",
group_index=0,
differing_index=1,
compared_index=1,
baseline_index=0,
max_abs_diff=0.1,
passed=False,
atol=1e-6,
diff=_make_diff(
rel_diff=0.1,
max_abs_diff=0.1,
mean_abs_diff=0.05,
diff_threshold=1e-6,
passed=False,
),
)
defaults.update(overrides)
return ReplicatedMismatchWarning(**defaults)
return ReplicatedCheckResult(**defaults)
class TestWarnings:
@@ -148,7 +156,7 @@ class TestWarnings:
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(passed=True),
warnings=[_make_warning()],
warnings=[GeneralWarning(category="test", message="some warning")],
)
assert record.category == "failed"
@@ -157,18 +165,44 @@ class TestWarnings:
record = SkipRecord(
name="x",
reason="no_baseline",
warnings=[_make_warning()],
warnings=[GeneralWarning(category="test", message="some warning")],
)
assert record.category == "failed"
def test_warnings_json_round_trip(self):
"""warnings survive model_dump_json → parse_record_json round-trip."""
warning = _make_warning(
def test_replicated_checks_all_passed(self):
"""ComparisonRecord with all replicated_checks passed → category=='passed'."""
record = ComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(passed=True),
replicated_checks=[_make_replicated_check(passed=True)],
)
assert record.category == "passed"
def test_replicated_checks_failed_means_record_failed(self):
"""ComparisonRecord with any replicated_check.passed=False → category=='failed'."""
record = ComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(passed=True),
replicated_checks=[_make_replicated_check(passed=False)],
)
assert record.category == "failed"
def test_replicated_check_json_round_trip(self):
"""ReplicatedCheckResult survives JSON round-trip via ComparisonRecord."""
check = _make_replicated_check(
axis="cp",
group_index=2,
differing_index=3,
compared_index=3,
baseline_index=0,
max_abs_diff=0.42,
passed=False,
)
record = ComparisonRecord(
name="mlp",
@@ -177,30 +211,23 @@ class TestWarnings:
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
warnings=[warning],
replicated_checks=[check],
)
restored = parse_record_json(record.model_dump_json())
assert isinstance(restored, ComparisonRecord)
assert len(restored.warnings) == 1
assert len(restored.replicated_checks) == 1
restored_warning = restored.warnings[0]
assert restored_warning.axis == "cp"
assert restored_warning.group_index == 2
assert restored_warning.differing_index == 3
assert restored_warning.baseline_index == 0
assert restored_warning.max_abs_diff == pytest.approx(0.42)
restored_check: ReplicatedCheckResult = restored.replicated_checks[0]
assert restored_check.axis == "cp"
assert restored_check.group_index == 2
assert restored_check.compared_index == 3
assert restored_check.baseline_index == 0
assert not restored_check.passed
def test_any_warning_discriminated_union_round_trip(self):
"""All AnyWarning variants survive JSON round-trip via a WarningRecord."""
all_warnings = [
ReplicatedMismatchWarning(
axis="tp",
group_index=0,
differing_index=1,
baseline_index=0,
max_abs_diff=0.1,
),
GeneralWarning(
category="aux_tensors_missing",
message="Aux tensors missing, skipping token alignment",
test/registered/debug_utils/comparator/test_entrypoint.py
+114 -13
View File
@@ -13,7 +13,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
ConfigRecord,
GeneralWarning,
NonTensorRecord,
ReplicatedMismatchWarning,
ReplicatedCheckResult,
SkipRecord,
SummaryRecord,
WarningRecord,
@@ -910,7 +910,7 @@ class TestEntrypointGroupingLogical:
assert comp.name == "hidden"
def test_recompute_pseudo_mismatch_warning(self, tmp_path, capsys):
"""Recompute pseudo-axis with differing original/recompute → ReplicatedMismatchWarning."""
"""Recompute pseudo-axis with differing original/recompute → failed replicated_checks."""
torch.manual_seed(42)
tensor = torch.randn(4, 8)
mismatched_tensor = tensor + torch.randn(4, 8) * 10.0
@@ -937,12 +937,112 @@ class TestEntrypointGroupingLogical:
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
recompute_warnings = [
w
for w in comparisons[0].warnings
if isinstance(w, ReplicatedMismatchWarning) and w.axis == "recompute_pseudo"
recompute_checks: list[ReplicatedCheckResult] = [
c for c in comparisons[0].replicated_checks if c.axis == "recompute_pseudo"
]
assert len(recompute_warnings) > 0
assert len(recompute_checks) > 0
assert any(not c.passed for c in recompute_checks)
def test_tp_partial_reduction_unshard(self, tmp_path, capsys):
"""TP=2 with partial reduction: element-wise sum reconstructs full tensor."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8)
full_target = full_baseline + torch.randn(4, 8) * 0.001
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
baseline_path = _create_tp_partial_dumps(
baseline_dir,
full_tensor=full_baseline,
name="attn_out",
tp_size=2,
dims_str="b h(tp,partial)",
)
target_path = _create_tp_partial_dumps(
target_dir,
full_tensor=full_target,
name="attn_out",
tp_size=2,
dims_str="b h(tp,partial)",
)
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.total == 1
assert summary.passed == 1
def test_tp_partial_vs_single_rank(self, tmp_path, capsys):
"""Baseline single rank vs target TP=2 partial: unshard target then compare."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8)
target_full = full_tensor + torch.randn(4, 8) * 0.001
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
baseline_path = _create_rank_dump(
baseline_dir, rank=0, name="attn_out", tensor=full_tensor
)
target_path = _create_tp_partial_dumps(
target_dir,
full_tensor=target_full,
name="attn_out",
tp_size=2,
dims_str="b h(tp,partial)",
)
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
def test_cp_concat_tp_partial_reduction(self, tmp_path, capsys):
"""CP=2 concat + TP=2 partial reduction: multi-axis unshard."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8, 16)
full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
for side_dir, full_tensor in [
(tmp_path / "baseline", full_baseline),
(tmp_path / "target", full_target),
]:
side_dir.mkdir()
cp_chunks = list(full_tensor.chunk(2, dim=1))
rank = 0
for cp_rank in range(2):
for tp_rank in range(2):
_create_rank_dump(
side_dir,
rank=rank,
name="hidden",
tensor=cp_chunks[cp_rank] / 2,
dims="b s(cp) h(tp,partial)",
parallel_info={
"cp_rank": cp_rank,
"cp_size": 2,
"tp_rank": tp_rank,
"tp_size": 2,
},
)
rank += 1
args = _make_args(
tmp_path / "baseline" / _FIXED_EXP_NAME,
tmp_path / "target" / _FIXED_EXP_NAME,
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
def test_tp_partial_reduction_unshard(self, tmp_path, capsys):
"""TP=2 with partial reduction: element-wise sum reconstructs full tensor."""
@@ -1233,7 +1333,7 @@ class TestEntrypointReplicatedAxis:
"""Test replicated-axis scenarios through the full entrypoint pipeline."""
def test_replicated_axis_identical_replicas_passed(self, tmp_path, capsys):
"""CP2 TP2, TP replicated and identical → passed, no warnings."""
"""CP2 TP2, TP replicated and identical → passed, replicated_checks all passed."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8, 6)
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
@@ -1264,13 +1364,14 @@ class TestEntrypointReplicatedAxis:
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.warnings == []
assert all(c.passed for c in comp.replicated_checks)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 1
def test_replicated_mismatch_fails(self, tmp_path, capsys):
"""CP2 TP2, TP replicas differ (> atol) → failed with warnings."""
"""CP2 TP2, TP replicas differ (> atol) → failed with replicated_checks."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8, 6)
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
@@ -1303,14 +1404,14 @@ class TestEntrypointReplicatedAxis:
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].category == "failed"
assert len(comparisons[0].warnings) > 0
assert any(not c.passed for c in comparisons[0].replicated_checks)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.failed == 1
def test_summary_counts_failed_from_warnings_only(self, tmp_path, capsys):
"""Diff itself passes but TP replicas differ → summary.failed=1 from warnings."""
def test_summary_counts_failed_from_replicated_checks_only(self, tmp_path, capsys):
"""Diff itself passes but TP replicas differ → summary.failed=1 from replicated_checks."""
torch.manual_seed(42)
full_baseline = torch.randn(4, 8, 6)
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
@@ -1352,7 +1453,7 @@ class TestEntrypointReplicatedAxis:
comp = comparisons[0]
assert comp.diff is not None
assert comp.diff.passed
assert len(comp.warnings) > 0
assert any(not c.passed for c in comp.replicated_checks)
assert comp.category == "failed"
summary = records[-1]
test/registered/debug_utils/comparator/test_warning_sink.py
+10 -13
View File
@@ -3,23 +3,20 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _make_warning(**overrides) -> ReplicatedMismatchWarning:
def _make_warning(**overrides) -> GeneralWarning:
defaults: dict = dict(
axis="tp",
group_index=0,
differing_index=1,
baseline_index=0,
max_abs_diff=0.1,
category="test",
message="test warning",
)
defaults.update(overrides)
return ReplicatedMismatchWarning(**defaults)
return GeneralWarning(**defaults)
class TestWarningSink:
@@ -35,8 +32,8 @@ class TestWarningSink:
def test_nested_contexts(self) -> None:
sink = WarningSink()
outer_warning = _make_warning(group_index=0)
inner_warning = _make_warning(group_index=1)
outer_warning = _make_warning(message="outer")
inner_warning = _make_warning(message="inner")
with sink.context() as outer:
sink.add(outer_warning)
@@ -61,7 +58,7 @@ class TestWarningSink:
sink.add(_make_warning())
captured = capsys.readouterr()
assert "Replicated along tp" in captured.out
assert "test warning" in captured.out
def test_context_captures_instead_of_printing(self, capsys) -> None:
sink = WarningSink()
@@ -96,9 +93,9 @@ class TestWarningSink:
assert len(collected) == 1
sink.add(_make_warning(group_index=99))
sink.add(_make_warning(message="after exception"))
captured = capsys.readouterr()
assert "Replicated along tp" in captured.out
assert "after exception" in captured.out
if __name__ == "__main__":