Support handling arbitrary objects in dump comparator (#19558)

python/sglang/srt/debug_utils/comparator/bundle_comparator.py

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
 from sglang.srt.debug_utils.comparator.dims import apply_dim_names, parse_dim_names
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
+    NonTensorRecord,
     SkipRecord,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
@@ -28,7 +29,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
 )
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
-from sglang.srt.debug_utils.dump_loader import ValueWithMeta
+from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
 
 _FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"}
 
@@ -44,9 +45,9 @@ def compare_bundle_pair(
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
-) -> Union[ComparisonRecord, SkipRecord]:
+) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
     with warning_sink.context() as collected_warnings:
-        result = _compare_bundle_pair_raw(
+        result = _compare_bundle_pair_inner(
             name=name,
             filenames_pair=filenames_pair,
             baseline_path=baseline_path,
@@ -59,7 +60,7 @@ def compare_bundle_pair(
     return result.model_copy(update={"warnings": collected_warnings})
 
 
-def _compare_bundle_pair_raw(
+def _compare_bundle_pair_inner(
     *,
     name: str,
     filenames_pair: Pair[list[str]],
@@ -70,18 +71,49 @@ def _compare_bundle_pair_raw(
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
-) -> Union[ComparisonRecord, SkipRecord]:
-    # 1. Load (tensor + meta, ungrouped)
-    valid_pair: Pair[list[ValueWithMeta]] = Pair(
-        x=_load_valid_tensors(filenames=filenames_pair.x, base_path=baseline_path),
-        y=_load_valid_tensors(filenames=filenames_pair.y, base_path=target_path),
+) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
+    # 1. Load all successfully loaded values
+    all_pair: Pair[list[ValueWithMeta]] = Pair(
+        x=_load_all_values(filenames=filenames_pair.x, base_path=baseline_path),
+        y=_load_all_values(filenames=filenames_pair.y, base_path=target_path),
     )
 
+    if not all_pair.x or not all_pair.y:
+        reason = "baseline_load_failed" if not all_pair.x else "target_load_failed"
+        return SkipRecord(name=name, reason=reason)
+
+    # 2. Check if any side has non-tensor values → non-tensor display path
+    has_non_tensor: bool = any(
+        not isinstance(it.value, torch.Tensor) for it in [*all_pair.x, *all_pair.y]
+    )
+    if has_non_tensor:
+        return _compare_bundle_pair_non_tensor_type(name=name, value_pair=all_pair)
+
+    # 3. All values are tensors → tensor comparison path
+    return _compare_bundle_pair_tensor_type(
+        name=name,
+        valid_pair=all_pair,
+        token_aligner_plan=token_aligner_plan,
+        diff_threshold=diff_threshold,
+        thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+    )
+
+
+def _compare_bundle_pair_tensor_type(
+    *,
+    name: str,
+    valid_pair: Pair[list[ValueWithMeta]],
+    token_aligner_plan: Optional[TokenAlignerPlan],
+    diff_threshold: float,
+    thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
+        x=None, y=None
+    ),
+) -> Union[ComparisonRecord, SkipRecord]:
     if not valid_pair.x or not valid_pair.y:
         reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
         return SkipRecord(name=name, reason=reason)
 
-    # 2. Plan (meta only, no tensor)
+    # Plan (meta only, no tensor)
     metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
         lambda items: [it.meta for it in items]
     )
@@ -91,7 +123,7 @@ def _compare_bundle_pair_raw(
         thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
     )
 
-    # 3. Apply dim names to tensors, then execute
+    # Apply dim names to tensors, then execute
     tensors_pair: Pair[list[torch.Tensor]] = Pair(
         x=_apply_dim_names_from_meta(
             tensors=[it.value for it in valid_pair.x],
@@ -109,10 +141,10 @@ def _compare_bundle_pair_raw(
     if aligner_result.tensors is None:
         assert aligner_result.failed_side_xy is not None
         side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
-        reason = f"{side_name}_load_failed"
+        reason: str = f"{side_name}_load_failed"
         return SkipRecord(name=name, reason=reason)
 
-    # 4. Compare
+    # Compare
     info = compare_tensor_pair(
         x_baseline=aligner_result.tensors.x.rename(None),
         x_target=aligner_result.tensors.y.rename(None),
@@ -122,20 +154,27 @@ def _compare_bundle_pair_raw(
     return ComparisonRecord(**info.model_dump(), aligner_plan=plan)
 
 
-def _apply_dim_names_from_meta(
+def _compare_bundle_pair_non_tensor_type(
     *,
-    tensors: list[torch.Tensor],
-    metas: list[dict[str, Any]],
-) -> list[torch.Tensor]:
-    if not metas:
-        return tensors
+    name: str,
+    value_pair: Pair[list[ValueWithMeta]],
+) -> NonTensorRecord:
+    baseline_value: Any = value_pair.x[0].value
+    target_value: Any = value_pair.y[0].value
 
-    dims_str: Optional[str] = metas[0].get("dims")
-    if dims_str is None:
-        return tensors
+    try:
+        values_equal: bool = bool(baseline_value == target_value)
+    except Exception:
+        values_equal = False
 
-    dim_names: list[str] = parse_dim_names(dims_str)
-    return [apply_dim_names(t, dim_names) for t in tensors]
+    return NonTensorRecord(
+        name=name,
+        baseline_value=repr(baseline_value),
+        target_value=repr(target_value),
+        baseline_type=type(baseline_value).__name__,
+        target_type=type(target_value).__name__,
+        values_equal=values_equal,
+    )
 
 
 def _apply_dim_names_from_meta(
@@ -154,9 +193,9 @@ def _apply_dim_names_from_meta(
     return [apply_dim_names(t, dim_names) for t in tensors]
 
 
-def _load_valid_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
+def _load_all_values(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
     return [
-        x
+        item
         for f in filenames
-        if isinstance((x := ValueWithMeta.load(base_path / f)).value, torch.Tensor)
+        if (item := ValueWithMeta.load(base_path / f)).value is not LOAD_FAILED
     ]

python/sglang/srt/debug_utils/comparator/display.py

@@ -12,7 +12,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
     RankInfoRecord,
     print_record,
 )
-from sglang.srt.debug_utils.dump_loader import ValueWithMeta
+from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
 
 _PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"]
 
@@ -96,7 +96,7 @@ def _collect_input_ids_and_positions(
     for row in filtered.to_dicts():
         key: tuple[int, int] = (row["step"], row["rank"])
         item: ValueWithMeta = ValueWithMeta.load(dump_dir / row["filename"])
-        if item.value is not None:
+        if item.value is not LOAD_FAILED:
             data_by_step_rank[key][row["name"]] = item.value
 
     table_rows: list[dict[str, Any]] = []

python/sglang/srt/debug_utils/comparator/entrypoint.py

@@ -25,6 +25,7 @@ from sglang.srt.debug_utils.comparator.display import emit_display_records
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     ConfigRecord,
+    NonTensorRecord,
     SkipRecord,
     SummaryRecord,
     print_record,
@@ -137,7 +138,7 @@ def _compare_bundle_pairs(
     token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold: float,
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]],
-) -> Iterator[Union[ComparisonRecord, SkipRecord]]:
+) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]:
     for bundle_info_pair in bundle_info_pairs:
         if not bundle_info_pair.y:
             continue
@@ -159,7 +160,7 @@ def _compare_bundle_pairs(
 
 def _consume_comparison_records(
     *,
-    comparison_records: Iterator[Union[ComparisonRecord, SkipRecord]],
+    comparison_records: Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]],
     output_format: str,
 ) -> None:
     counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}

python/sglang/srt/debug_utils/comparator/output_types.py

@@ -140,6 +140,31 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
         return body
 
 
+class NonTensorRecord(_OutputRecord):
+    type: Literal["non_tensor"] = "non_tensor"
+    name: str
+    baseline_value: str
+    target_value: str
+    baseline_type: str
+    target_type: str
+    values_equal: bool
+
+    @property
+    def category(self) -> str:
+        if self.warnings:
+            return "failed"
+        return "passed" if self.values_equal else "failed"
+
+    def _format_body(self) -> str:
+        if self.values_equal:
+            return f"NonTensor: {self.name} = {self.baseline_value} ({self.baseline_type}) [equal]"
+        return (
+            f"NonTensor: {self.name}\n"
+            f"  baseline = {self.baseline_value} ({self.baseline_type})\n"
+            f"  target   = {self.target_value} ({self.target_type})"
+        )
+
+
 class SummaryRecord(_OutputRecord):
     type: Literal["summary"] = "summary"
     total: int
@@ -207,6 +232,7 @@ AnyRecord = Annotated[
         InputIdsRecord,
         SkipRecord,
         ComparisonRecord,
+        NonTensorRecord,
         SummaryRecord,
         WarningRecord,
     ],

python/sglang/srt/debug_utils/dump_loader.py

@@ -9,6 +9,8 @@ import torch
 
 _TYPED_FIELDS: list[tuple[str, type]] = [("rank", int)]
 
+LOAD_FAILED: object = object()
+
 
 def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
     stem = Path(path).stem
@@ -38,7 +40,7 @@ class ValueWithMeta:
         except Exception as e:
             print(f"Skip load {path} since error {e}")
             return ValueWithMeta(
-                value=None, meta={**meta_from_filename, "filename": path.name}
+                value=LOAD_FAILED, meta={**meta_from_filename, "filename": path.name}
             )
 
         value, meta_from_embedded = _unwrap_dict_format(raw)