From ccbc47d6be44ecdebb4ee3d15eaf4b11e1c63623 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 28 Feb 2026 18:06:26 +0800 Subject: [PATCH] Update layer id extraction, diffing, empty handling and error sentinel in dump comparator (#19562) --- .../aligner/unsharder/parallel_info.py | 7 ++++- .../comparator/aligner/unsharder/planner.py | 6 +++- .../tensor_comparator/comparator.py | 29 +++++++++++++++---- .../comparator/tensor_comparator/formatter.py | 12 +++----- .../sglang/srt/debug_utils/dump_comparator.py | 12 +++----- 5 files changed, 43 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py index 7efd79758..719683d4e 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py @@ -6,12 +6,17 @@ from sglang.srt.debug_utils.comparator.dims import ParallelAxis _PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info") +def _is_error_sentinel(value: dict) -> bool: + """Check if a parallel_info dict is an error sentinel (e.g. {'megatron_error': True}).""" + return any(k.endswith("_error") for k in value) + + def normalize_parallel_info(meta: dict) -> dict[ParallelAxis, AxisInfo]: """Extract unified parallel info from dump meta.""" info: Optional[dict] = None for key in _PARALLEL_INFO_KEYS: value = meta.get(key) - if isinstance(value, dict) and value: + if isinstance(value, dict) and value and not _is_error_sentinel(value): if info is not None: raise ValueError( f"Meta contains multiple parallel_info keys among {_PARALLEL_INFO_KEYS}" diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py index 5a525d334..061365705 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py @@ -159,7 +159,11 @@ def _resolve_unshard_params( f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)" ) - if spec.name == TOKEN_DIM_NAME and thd_global_seq_lens is not None: + if ( + spec.name == TOKEN_DIM_NAME + and spec.parallel == ParallelAxis.CP + and thd_global_seq_lens is not None + ): if spec.parallel is None: raise ValueError( f"THD unshard requires a parallel axis on dim '{spec.name}', but got None" diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py index bff72bde4..f6cff930b 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py @@ -90,6 +90,16 @@ def compare_tensor_pair( def _compute_tensor_stats(x: torch.Tensor) -> TensorStats: + if x.numel() == 0: + return TensorStats( + mean=0.0, + abs_mean=0.0, + std=0.0, + min=0.0, + max=0.0, + percentiles={}, + ) + include_quantiles: bool = x.numel() < QUANTILE_NUMEL_THRESHOLD return TensorStats( mean=torch.mean(x).item(), @@ -113,6 +123,19 @@ def _compute_diff( x_target: torch.Tensor, diff_threshold: float = 1e-3, ) -> DiffInfo: + if x_baseline.numel() == 0: + return DiffInfo( + rel_diff=0.0, + max_abs_diff=0.0, + mean_abs_diff=0.0, + abs_diff_percentiles={}, + max_diff_coord=[], + baseline_at_max=0.0, + target_at_max=0.0, + diff_threshold=diff_threshold, + passed=True, + ) + raw_abs_diff = (x_target - x_baseline).abs() max_diff_coord = argmax_coord(raw_abs_diff) @@ -133,9 +156,5 @@ def _compute_diff( baseline_at_max=x_baseline[max_diff_coord].item(), target_at_max=x_target[max_diff_coord].item(), diff_threshold=diff_threshold, - passed=( - rel_diff <= diff_threshold - and max_abs_diff <= diff_threshold - and mean_abs_diff <= diff_threshold - ), + passed=rel_diff <= diff_threshold, ) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py index 91a9c7a2e..fdc14fc57 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py @@ -79,16 +79,12 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]: + rel_diff_marker: str = "❌" if diff.rel_diff > diff.diff_threshold else "✅" lines: list[str] = [ prefix_text - + "\t".join( - f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}" - for name, value in [ - ("rel_diff", diff.rel_diff), - ("max_abs_diff", diff.max_abs_diff), - ("mean_abs_diff", diff.mean_abs_diff), - ] - ), + + f"{rel_diff_marker} rel_diff={diff.rel_diff}\t" + + f"max_abs_diff={diff.max_abs_diff}\t" + + f"mean_abs_diff={diff.mean_abs_diff}", f"max_abs_diff happens at coord={diff.max_diff_coord} with " f"baseline={diff.baseline_at_max} " f"target={diff.target_at_max}", diff --git a/python/sglang/srt/debug_utils/dump_comparator.py b/python/sglang/srt/debug_utils/dump_comparator.py index 5862d6cd5..6f5a3397d 100644 --- a/python/sglang/srt/debug_utils/dump_comparator.py +++ b/python/sglang/srt/debug_utils/dump_comparator.py @@ -194,16 +194,12 @@ def _compute_and_print_diff( mean_abs_diff = raw_abs_diff.mean().item() rel_diff = _calc_rel_diff(x_target, x_baseline) + rel_diff_marker: str = "❌" if rel_diff > diff_threshold else "✅" print( prefix_text - + "\t".join( - f"{'❌' if value > diff_threshold else '✅'} {name}={value}" - for name, value in [ - ("rel_diff", rel_diff), - ("max_abs_diff", max_abs_diff), - ("mean_abs_diff", mean_abs_diff), - ] - ) + + f"{rel_diff_marker} rel_diff={rel_diff}\t" + + f"max_abs_diff={max_abs_diff}\t" + + f"mean_abs_diff={mean_abs_diff}" ) max_diff_coord = _argmax_coord(raw_abs_diff)