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

56
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py
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=}")

2
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py
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

80
python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py
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(

7
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
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(

40
python/sglang/srt/debug_utils/comparator/output_types.py
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

6
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
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,

23
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
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] = []