Handle recompute and verify closeness in dumper (#19564)

python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py

@@ -24,12 +24,18 @@ def normalize_parallel_info(meta: dict) -> dict[ParallelAxis, AxisInfo]:
             info = value
 
     if info is None:
-        return {}
+        info = {}
 
     result: dict[ParallelAxis, AxisInfo] = {}
     for axis in ParallelAxis:
         axis_rank = info.get(f"{axis.value}_rank")
         axis_size = info.get(f"{axis.value}_size")
+
+        # Recompute pseudo-axis lives at top-level meta, not inside parallel_info
+        if axis_rank is None:
+            axis_rank = meta.get(f"{axis.value}_rank")
+            axis_size = meta.get(f"{axis.value}_size")
+
         if axis_rank is not None and axis_size is not None and axis_size > 1:
             result[axis] = AxisInfo(
                 axis_rank=axis_rank,

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

@@ -20,6 +20,7 @@ class ParallelAxis(Enum):
     CP = "cp"
     EP = "ep"
     SP = "sp"
+    RECOMPUTE_PSEUDO = "recompute_pseudo"
 
 
 class Ordering(Enum):

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

@@ -124,7 +124,7 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
 def _compute_skip_keys(args, *, has_token_aligner_plan: bool):
     skip_keys: set[str] = {"dump_index", "filename"}
     if args.grouping == "logical":
-        skip_keys |= {"rank"}
+        skip_keys |= {"rank", "recompute_status"}
         if has_token_aligner_plan:
             skip_keys |= {"step"}
     return skip_keys

python/sglang/srt/debug_utils/dump_loader.py

@@ -2,12 +2,12 @@ import functools
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 
 import polars as pl
 import torch
 
-_TYPED_FIELDS: list[tuple[str, type]] = [("rank", int)]
+LOAD_FAILED: object = object()
 
 LOAD_FAILED: object = object()
 
@@ -19,9 +19,9 @@ def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
         if "=" in kv:
             k, v = kv.split("=", 1)
             result[k] = v
-    for field, converter in _TYPED_FIELDS:
-        if field in result:
-            result[field] = converter(result[field])
+    for field_name, converter in _TYPED_FIELDS:
+        if field_name in result:
+            result[field_name] = converter(result[field_name])
     return result
 
 
@@ -177,4 +177,9 @@ def read_tokenizer_path(directory: Path) -> Optional[str]:
     return None
 
 
+_TYPED_FIELDS: list[tuple[str, Callable[[str], Any]]] = [
+    ("rank", int),
+]
+
+
 dump_loader = DumpLoader()

python/sglang/srt/debug_utils/dumper.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import json
 import os
@@ -415,7 +416,14 @@ class _Dumper:
         if not self._config.enable:
             return
 
-        tags = dict(name=name, **extra_kwargs, **self._state.global_ctx)
+        recompute_status = _detect_recompute_status()
+        tags = dict(
+            name=name,
+            recompute_status=recompute_status.value,
+            **extra_kwargs,
+            **self._state.global_ctx,
+        )
+
         if (f := self._config.filter) is not None and re.search(
             f, _format_tags(tags)
         ) is None:
@@ -424,6 +432,7 @@ class _Dumper:
         if not (enable_value or enable_curr_grad or enable_future_grad):
             return
 
+        recompute_meta = recompute_status.to_pseudo_parallel_meta()
         value = _materialize_value(value)
 
         if enable_value:
@@ -432,7 +441,7 @@ class _Dumper:
                 tags=tags,
                 value=value,
                 save=save,
-                meta_only_fields=value_meta_only_fields or {},
+                meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta},
             )
 
         if (
@@ -445,7 +454,7 @@ class _Dumper:
                 tags={**tags, "name": f"grad__{name}"},
                 value=g,
                 save=save,
-                meta_only_fields=grad_meta_only_fields or {},
+                meta_only_fields={**(grad_meta_only_fields or {}), **recompute_meta},
             )
 
         if enable_future_grad:
@@ -472,7 +481,10 @@ class _Dumper:
             return
 
         captured_step = self._state.step
-        captured_tags = dict(name=f"grad__{name}", **deepcopy(extra_kwargs))
+        captured_tags = dict(
+            name=f"grad__{name}",
+            **deepcopy(extra_kwargs),
+        )
         captured_meta_only = meta_only_fields or {}
 
         def grad_hook(grad: torch.Tensor) -> None:
@@ -1142,6 +1154,20 @@ def _get_local_ip_by_remote() -> Optional[str]:
 # -------------------------------------- framework plugins ------------------------------------------
 
 
+class _RecomputeStatus(enum.Enum):
+    DISABLED = "disabled"
+    ORIGINAL = "original"  # inside checkpoint, original forward
+    RECOMPUTE = "recompute"  # inside checkpoint, recompute forward
+
+    def to_pseudo_parallel_meta(self) -> dict[str, Any]:
+        if self == _RecomputeStatus.DISABLED:
+            return {}
+        return {
+            "recompute_pseudo_rank": 1 if self == _RecomputeStatus.RECOMPUTE else 0,
+            "recompute_pseudo_size": 2,
+        }
+
+
 class _FrameworkPlugin(ABC):
     @property
     @abstractmethod
@@ -1168,6 +1194,9 @@ class _FrameworkPlugin(ABC):
     def get_tokenizer_path(self) -> Optional[str]:
         return None
 
+    def detect_recompute_status(self) -> _RecomputeStatus:
+        return _RecomputeStatus.DISABLED
+
 
 class _SGLangPlugin(_FrameworkPlugin):
     _available = True
@@ -1353,10 +1382,32 @@ class _MegatronPlugin(_FrameworkPlugin):
             {"input_ids", "position_ids", "cu_seqlens_q", "cu_seqlens_kv", "qkv_format"}
         )
 
+    def detect_recompute_status(self) -> _RecomputeStatus:
+        if not self._available:
+            return _RecomputeStatus.DISABLED
+        try:
+            from megatron.core.tensor_parallel.random import is_checkpointing
+
+            if not is_checkpointing():
+                return _RecomputeStatus.DISABLED
+            if torch.is_grad_enabled():
+                return _RecomputeStatus.RECOMPUTE
+            return _RecomputeStatus.ORIGINAL
+        except (ImportError, AttributeError):
+            return _RecomputeStatus.DISABLED
+
 
 _plugins: list[_FrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()]
 
 
+def _detect_recompute_status() -> _RecomputeStatus:
+    for plugin in _plugins:
+        info = plugin.detect_recompute_status()
+        if info != _RecomputeStatus.DISABLED:
+            return info
+    return _RecomputeStatus.DISABLED
+
+
 # -------------------------------------- singleton ------------------------------------------