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
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py
+38 -18
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Optional
import torch
@@ -21,9 +21,11 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import (
execute_token_aligner,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
UnsharderResult,
execute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
from sglang.srt.debug_utils.comparator.utils import Pair
@@ -31,6 +33,7 @@ from sglang.srt.debug_utils.comparator.utils import Pair
class AlignerResult:
tensors: Optional[Pair[torch.Tensor]]
failed_side_xy: Optional[str] # "x" or "y"; None if success
replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
def execute_aligner_plan(
@@ -39,18 +42,26 @@ def execute_aligner_plan(
plan: AlignerPlan,
) -> AlignerResult:
"""Execute unified unshard/reorder + token-align."""
all_checks: list[ReplicatedCheckResult] = []
# Per-side: unshard + reorder -> dict[step, tensor]
step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans(
step_tensors_x, checks_x = _execute_step_plans(
tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
)
step_tensors_y: dict[int, torch.Tensor] = _execute_step_plans(
all_checks.extend(checks_x)
step_tensors_y, checks_y = _execute_step_plans(
tensors=tensors_pair.y, step_plans=plan.per_step_plans.y
)
all_checks.extend(checks_y)
if not step_tensors_x or not step_tensors_y:
failed_side_xy: str = "x" if not step_tensors_x else "y"
return AlignerResult(tensors=None, failed_side_xy=failed_side_xy)
return AlignerResult(
tensors=None,
failed_side_xy=failed_side_xy,
replicated_checks=all_checks,
)
# Cross-side: token alignment (or direct extraction for single-step)
if plan.token_aligner_plan is not None:
@@ -72,55 +83,64 @@ def execute_aligner_plan(
y=execute_axis_aligner_plan(tensor=combined.y, plan=aligner_plan, side="y"),
)
return AlignerResult(tensors=combined, failed_side_xy=None)
return AlignerResult(
tensors=combined,
failed_side_xy=None,
replicated_checks=all_checks,
)
def _execute_step_plans(
tensors: list[torch.Tensor],
step_plans: list[AlignerPerStepPlan],
) -> dict[int, torch.Tensor]:
) -> tuple[dict[int, torch.Tensor], list[ReplicatedCheckResult]]:
result: dict[int, torch.Tensor] = {}
all_checks: list[ReplicatedCheckResult] = []
for step_plan in step_plans:
step_tensors: list[torch.Tensor] = [
tensors[i] for i in step_plan.input_object_indices
]
tensor: Optional[torch.Tensor] = execute_sub_plans(
tensor, checks = execute_sub_plans(
tensors=step_tensors, plans=step_plan.sub_plans
)
all_checks.extend(checks)
if tensor is not None:
result[step_plan.step] = tensor
return result
return result, all_checks
def execute_sub_plans(
tensors: list[torch.Tensor],
plans: list[AlignerPerStepSubPlan],
) -> Optional[torch.Tensor]:
) -> tuple[Optional[torch.Tensor], list[ReplicatedCheckResult]]:
if not tensors:
return None
return None, []
if not plans:
if len(tensors) != 1:
return None
return tensors[0]
return None, []
return tensors[0], []
current = tensors
current: list[torch.Tensor] = tensors
all_checks: list[ReplicatedCheckResult] = []
for plan in plans:
current = execute_sub_plan(tensors=current, plan=plan)
current, checks = execute_sub_plan(tensors=current, plan=plan)
all_checks.extend(checks)
assert len(current) == 1
return current[0]
return current[0], all_checks
def execute_sub_plan(
tensors: list[torch.Tensor],
plan: AlignerPerStepSubPlan,
) -> list[torch.Tensor]:
) -> tuple[list[torch.Tensor], list[ReplicatedCheckResult]]:
if isinstance(plan, UnsharderPlan):
return execute_unsharder_plan(plan, tensors)
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, tensors)
return unsharder_result.tensors, unsharder_result.replicated_checks
elif isinstance(plan, ReordererPlan):
return execute_reorderer_plan(plan, tensors)
return execute_reorderer_plan(plan, tensors), []
else:
raise NotImplementedError(f"Unknown {plan=}")
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py
+1 -1
View File
@@ -236,7 +236,7 @@ def _load_and_align_aux_tensor(
dim_names: list[str] = resolve_dim_names(dims_str)
tensors = [apply_dim_names(t, dim_names) for t in tensors]
result = execute_sub_plans(tensors=tensors, plans=sub_plans)
result, _replicated_checks = execute_sub_plans(tensors=tensors, plans=sub_plans)
assert result is not None
return result.rename(None) # strip named dims before returning to plugin
python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py
+52 -28
View File
@@ -1,3 +1,4 @@
from dataclasses import dataclass, field
from typing import Optional
import torch
@@ -14,27 +15,37 @@ from sglang.srt.debug_utils.comparator.dims import (
ParallelAxis,
resolve_dim_by_name,
)
from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import compute_diff
_REPLICATED_ATOL: float = 1e-6
@dataclass(frozen=True)
class UnsharderResult:
tensors: list[torch.Tensor]
replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
def execute_unsharder_plan(
plan: UnsharderPlan,
tensors: list[torch.Tensor],
) -> list[torch.Tensor]:
result: list[torch.Tensor] = []
) -> UnsharderResult:
result_tensors: list[torch.Tensor] = []
all_checks: list[ReplicatedCheckResult] = []
for group_idx, group in enumerate(plan.groups):
group_tensors = [tensors[i] for i in group]
tensor = _apply_unshard(
tensor, checks = _apply_unshard(
plan.params,
group_tensors,
axis=plan.axis,
group_index=group_idx,
)
result.append(tensor)
result_tensors.append(tensor)
all_checks.extend(checks)
return result
return UnsharderResult(tensors=result_tensors, replicated_checks=all_checks)
def _apply_unshard(
@@ -43,25 +54,28 @@ def _apply_unshard(
*,
axis: ParallelAxis,
group_index: int,
) -> torch.Tensor:
) -> tuple[torch.Tensor, list[ReplicatedCheckResult]]:
if isinstance(params, PickParams):
_verify_replicated_group(
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
ordered_tensors,
axis=axis,
group_index=group_index,
)
return ordered_tensors[0]
return ordered_tensors[0], checks
if isinstance(params, ConcatParams):
dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
return torch.cat(ordered_tensors, dim=dim)
return torch.cat(ordered_tensors, dim=dim), []
if isinstance(params, CpThdConcatParams):
thd_dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
return _thd_concat(
ordered_tensors,
dim=thd_dim,
seq_lens_per_rank=params.seq_lens_per_rank,
return (
_thd_concat(
ordered_tensors,
dim=thd_dim,
seq_lens_per_rank=params.seq_lens_per_rank,
),
[],
)
if isinstance(params, ReduceSumParams):
@@ -70,7 +84,7 @@ def _apply_unshard(
names: tuple[Optional[str], ...] = ordered_tensors[0].names
if names[0] is not None:
result = result.refine_names(*names)
return result
return result, []
raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
@@ -80,21 +94,31 @@ def _verify_replicated_group(
*,
axis: ParallelAxis,
group_index: int,
) -> None:
baseline = ordered_tensors[0].rename(None)
) -> list[ReplicatedCheckResult]:
baseline: torch.Tensor = ordered_tensors[0].rename(None).float()
checks: list[ReplicatedCheckResult] = []
for i in range(1, len(ordered_tensors)):
other = ordered_tensors[i].rename(None)
if not torch.allclose(baseline, other, atol=1e-6):
warning_sink.add(
ReplicatedMismatchWarning(
axis=axis.value,
group_index=group_index,
differing_index=i,
baseline_index=0,
max_abs_diff=(baseline - other).abs().max().item(),
)
other: torch.Tensor = ordered_tensors[i].rename(None).float()
diff_info = compute_diff(
x_baseline=baseline,
x_target=other,
diff_threshold=_REPLICATED_ATOL,
)
passed: bool = diff_info.max_abs_diff <= _REPLICATED_ATOL
checks.append(
ReplicatedCheckResult(
axis=axis.value,
group_index=group_index,
compared_index=i,
baseline_index=0,
passed=passed,
atol=_REPLICATED_ATOL,
diff=diff_info,
)
)
return checks
def _thd_concat(
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
+6 -1
View File
@@ -160,6 +160,7 @@ def _compare_bundle_pair_tensor_type(
aligner_result: AlignerResult = execute_aligner_plan(
tensors_pair=tensors_pair, plan=plan
)
replicated_checks = aligner_result.replicated_checks
if aligner_result.tensors is None:
assert aligner_result.failed_side_xy is not None
@@ -183,7 +184,11 @@ def _compare_bundle_pair_tensor_type(
diff_threshold=diff_threshold,
seq_dim=seq_dim,
)
record = ComparisonRecord(**info.model_dump(), aligner_plan=plan)
record = ComparisonRecord(
**info.model_dump(),
aligner_plan=plan,
replicated_checks=replicated_checks,
)
if viz_output_dir is not None:
_try_generate_viz(
python/sglang/srt/debug_utils/comparator/output_types.py
+20 -20
View File
@@ -8,8 +8,10 @@ from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_valida
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison,
format_replicated_checks,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
)
from sglang.srt.debug_utils.comparator.utils import _StrictBase
@@ -20,22 +22,6 @@ if TYPE_CHECKING:
)
class ReplicatedMismatchWarning(_StrictBase):
kind: Literal["replicated_mismatch"] = "replicated_mismatch"
axis: str
group_index: int
differing_index: int
baseline_index: int
max_abs_diff: float
def to_text(self) -> str:
return (
f"Replicated along {self.axis}: group {self.group_index}, "
f"index {self.differing_index} differs from {self.baseline_index} "
f"(max_abs_diff={self.max_abs_diff:.6e})"
)
class GeneralWarning(_StrictBase):
kind: Literal["general"] = "general"
category: str
@@ -45,10 +31,19 @@ class GeneralWarning(_StrictBase):
return self.message
AnyWarning = Annotated[
Union[ReplicatedMismatchWarning, GeneralWarning],
Discriminator("kind"),
]
# Type alias — currently only GeneralWarning exists.
# When adding new warning types, convert back to Union + Discriminator("kind").
AnyWarning = GeneralWarning
class ReplicatedCheckResult(_StrictBase):
axis: str
group_index: int
compared_index: int
baseline_index: int
passed: bool
atol: float
diff: DiffInfo
class _OutputRecord(_StrictBase):
@@ -126,15 +121,20 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
type: Literal["comparison"] = "comparison"
aligner_plan: Optional[AlignerPlan] = None
replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
@property
def category(self) -> str:
if self.warnings:
return "failed"
if any(not check.passed for check in self.replicated_checks):
return "failed"
return "passed" if self.diff is not None and self.diff.passed else "failed"
def _format_body(self) -> str:
body: str = format_comparison(self)
if self.replicated_checks:
body += "\n" + format_replicated_checks(self.replicated_checks)
if self.aligner_plan is not None:
body += "\n" + _format_aligner_plan(self.aligner_plan)
return body
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
+3 -3
View File
@@ -57,7 +57,7 @@ def compare_tensor_pair(
downcast_dtype: Optional[torch.dtype] = None
if not shape_mismatch:
diff = _compute_diff(
diff = compute_diff(
x_baseline=x_baseline_f,
x_target=x_target_f,
diff_threshold=diff_threshold,
@@ -74,7 +74,7 @@ def compare_tensor_pair(
Pair(x=baseline_original_dtype, y=target_original_dtype)
)
if downcast_dtype is not None:
diff_downcast = _compute_diff(
diff_downcast = compute_diff(
x_baseline=x_baseline_f.to(downcast_dtype),
x_target=x_target_f.to(downcast_dtype),
diff_threshold=diff_threshold,
@@ -121,7 +121,7 @@ def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
return {p: torch.quantile(x_float, p / 100.0).item() for p in DEFAULT_PERCENTILES}
def _compute_diff(
def compute_diff(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
diff_threshold: float = 1e-3,
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
+23
View File
@@ -1,9 +1,16 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorStats,
)
if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
def format_comparison(info: TensorComparisonInfo) -> str:
lines: list[str] = []
@@ -54,6 +61,22 @@ def format_comparison(info: TensorComparisonInfo) -> str:
return "\n".join(lines)
def format_replicated_checks(checks: list[ReplicatedCheckResult]) -> str:
lines: list[str] = ["Replicated checks:"]
for check in checks:
marker: str = "✅" if check.passed else "❌"
lines.append(
f" {marker} axis={check.axis} group={check.group_index} "
f"idx={check.compared_index} vs {check.baseline_index}: "
f"rel_diff={check.diff.rel_diff:.6e} "
f"max_abs_diff={check.diff.max_abs_diff:.6e} "
f"mean_abs_diff={check.diff.mean_abs_diff:.6e}"
)
return "\n".join(lines)
def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]:
lines: list[str] = []
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__":