sgl-project--sglang/python/sglang/srt/debug_utils/dumper.py

import enum
import functools
import json
import os
import random
import re
import socket
import threading
import time
import traceback
from abc import ABC, abstractmethod
from collections.abc import Callable
from contextlib import contextmanager
from copy import deepcopy
from dataclasses import asdict, dataclass, field, fields, replace
from functools import cached_property
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Any, List, Literal, Optional, Union, get_args, get_type_hints

import torch
import torch.distributed as dist
import zmq

from sglang.srt.managers.io_struct import sock_recv, sock_send, wrap_as_pickle

# -------------------------------------- config base ------------------------------------------


@dataclass(frozen=True)
class _BaseConfig(ABC):
    def __post_init__(self) -> None:
        self._verify_types()

    def _verify_types(self) -> None:
        hints = get_type_hints(type(self))
        cls_name = type(self).__name__
        for f in fields(self):
            value = getattr(self, f.name)
            if value is None:
                continue
            expected = self._unwrap_type(hints[f.name])
            if not isinstance(value, expected):
                raise TypeError(
                    f"{cls_name}.{f.name}: expected {expected.__name__}, "
                    f"got {type(value).__name__}"
                )

    @classmethod
    @abstractmethod
    def _env_prefix(cls) -> str: ...

    @classmethod
    def _env_name(cls, field_name: str) -> str:
        return f"{cls._env_prefix()}{field_name.upper()}"

    @classmethod
    def from_env(cls) -> "_BaseConfig":
        return cls(
            **{
                f.name: cls._parse_env_field(cls._env_name(f.name), f.default)
                for f in fields(cls)
            }
        )

    def with_defaults(self, **kwargs) -> "_BaseConfig":
        cls = type(self)
        actual = {
            key: value
            for key, value in kwargs.items()
            if os.getenv(cls._env_name(key)) is None
        }
        return replace(self, **actual) if actual else self

    @staticmethod
    def _unwrap_type(hint) -> type:
        args = get_args(hint)
        if args:
            return next(a for a in args if a is not type(None))
        return hint

    @classmethod
    def _parse_env_field(cls, env_name: str, default):
        return cls._parse_env_value(os.getenv(env_name), default)

    @staticmethod
    def _parse_env_value(raw, default):
        if raw is None or not raw.strip():
            return default
        if isinstance(default, bool):
            return raw.lower() in ("true", "1")
        if isinstance(default, int):
            return int(raw)
        return raw

    @classmethod
    def from_kv_pairs(cls, pairs: Optional[List[str]]) -> "_BaseConfig":
        return cls(**cls._kv_pairs_to_dict(pairs))

    @classmethod
    def _kv_pairs_to_dict(cls, pairs: Optional[List[str]]) -> dict:
        if not pairs:
            return {}

        missing = object()
        defaults = {f.name: f.default for f in fields(cls)}
        result: dict = {}

        for pair in pairs:
            key, sep, value = pair.partition("=")
            if not sep:
                raise ValueError(f"Invalid config pair (missing '='): {pair!r}")
            default = defaults.get(key, missing)
            if default is missing:
                raise ValueError(
                    f"Unknown config key {key!r}. Valid keys: {sorted(defaults)}"
                )
            try:
                result[key] = cls._parse_env_value(value, default)
            except (ValueError, TypeError) as exc:
                field_type = type(default).__name__
                raise TypeError(f"{key}: expected {field_type}, got {value!r}") from exc

        return result


_DEFAULT_EXP_NAME_PREFIX = "dump_"


@dataclass(frozen=True)
class DumperConfig(_BaseConfig):
    enable: bool = False
    filter: Optional[str] = None
    dir: str = "/tmp/dumper"
    enable_output_file: bool = True
    enable_output_console: bool = True
    enable_value: bool = True
    enable_grad: bool = False
    enable_model_value: bool = False
    enable_model_grad: bool = False
    exp_name: Optional[str] = None
    cleanup_previous: bool = False
    collective_timeout: int = 60
    server_port: str = "-1"
    non_intrusive_mode: str = "core"
    source_patcher_config: Optional[str] = None
    grafter_enable: bool = False
    grafter_role: str = ""  # required if enabled: "baseline" or "target"
    grafter_b2t_filter: Optional[str] = None  # names flowing baseline -> target
    grafter_t2b_filter: Optional[str] = None  # names flowing target -> baseline
    grafter_master_address: str = ""  # required if enabled
    grafter_master_port: int = -1  # required if enabled (positive port)
    grafter_baseline_world_size: int = -1  # required if enabled
    grafter_target_world_size: int = -1  # required if enabled
    grafter_backend: str = "nccl"
    grafter_group_name: str = "graft"
    grafter_timeout: int = 300
    # Fully-qualified Python path "pkg.subpkg.module.fn_name"
    # None -> use the default identity-by-rank fallback in _Grafter._default_transform.
    grafter_transform_path: Optional[str] = None
    # When True, append parallel-rank tags (pp_rank/tp_rank/...) to dump filenames so
    # tensors from different ranks do not collide when dumped into a shared directory.
    include_parallel_rank_in_filename: bool = False

    @classmethod
    def _env_prefix(cls) -> str:
        # NOTE: should not be `SGLANG_DUMPER_`, otherwise it is weird when dumping Megatron in Miles
        return "DUMPER_"

    def __post_init__(self) -> None:
        super().__post_init__()
        if self.grafter_enable:
            assert self.grafter_role in ("baseline", "target"), (
                f"grafter_role must be 'baseline' or 'target' when grafter_enable=True, "
                f"got {self.grafter_role!r}"
            )
            assert (
                self.grafter_master_address
            ), "grafter_master_address must be set when grafter_enable=True"
            assert self.grafter_master_port > 0, (
                f"grafter_master_port must be a positive port when grafter_enable=True, "
                f"got {self.grafter_master_port}"
            )
            assert self.grafter_baseline_world_size > 0, (
                f"grafter_baseline_world_size must be > 0 when grafter_enable=True, "
                f"got {self.grafter_baseline_world_size}"
            )
            assert self.grafter_target_world_size > 0, (
                f"grafter_target_world_size must be > 0 when grafter_enable=True, "
                f"got {self.grafter_target_world_size}"
            )
            assert (
                self.grafter_b2t_filter is not None
                or self.grafter_t2b_filter is not None
            ), (
                "grafter_enable=True but neither grafter_b2t_filter nor "
                "grafter_t2b_filter is set; nothing would ever be grafted"
            )

    @property
    def server_port_parsed(self) -> Optional[Union[int, Literal["reuse"]]]:
        raw = self.server_port
        if raw == "reuse":
            return "reuse"
        port = int(raw)
        if port <= 0:
            return None
        return port


# -------------------------------------- dumper core ------------------------------------------


@dataclass
class _DumperState:
    dump_index: int = 0
    step: int = 0
    global_ctx: dict = field(default_factory=dict)
    captured_output_data: Optional[dict] = None
    cleanup_previous_handled: bool = False


class _Dumper:
    """Utility to dump tensors, which can be useful when comparison checking models.

    Example usage:
    dumper.dump("layer_start__hidden_states", hidden_states, layer_id=self.layer_id)
    dumper.step()

    Import from non-SGLang system:
    ```
    import sys
    sys.path.append("/YOUR_PATH/sglang/python/sglang/srt/debug_utils")
    from dumper import dumper
    ```

    Then run the program:
    `DUMPER_ENABLE=1 python ...`

    Auto-cleanup old dumps before first write:
    `DUMPER_CLEANUP_PREVIOUS=1 python ...`

    Alternatively, disable at startup and configure via HTTP:
    1. `python ...`
    2. sglang mode:  `curl -X POST http://localhost:30000/dumper/configure -d '{"enable": true}'`
       standalone:   `curl -X POST http://localhost:40000/dumper/configure -d '{"enable": true}'`
    3. `curl -X POST http://localhost:30000/dumper/configure -d '{"enable": true, "filter": "layer_id=[0-3]"}'`
    4. `curl -X POST http://localhost:30000/dumper/reset`

    Related: `sglang.srt.debug_utils.dump_comparator` for dump comparison
    """

    def __init__(self, *, config: DumperConfig):
        self._config = config
        self._state = _DumperState()
        self._non_intrusives: list[_NonIntrusiveDumper] = []
        self._grafter = _Grafter(config=config)

    # ------------------------------- public :: core ---------------------------------

    @property
    def may_enable(self) -> bool:
        return self._config.enable or self._config.server_port_parsed is not None

    def step(self):
        """This should be called on all ranks at the end of each iteration."""

        self._http_manager  # noqa: B018

        if not self._config.enable:
            return

        # Users may want to `dump` only on some ranks, thus determine name here
        self._ensure_exp_name()

        self._state.step += 1
        _log(f"step={self._state.step}")

    def dump(
        self,
        name: str,
        value,
        save: bool = True,
        dims: Optional[str] = None,
        dims_grad: Optional[str] = None,
        grafter_extras: Optional[dict] = None,
        **kwargs,
    ) -> None:
        value_meta: dict = {}
        grad_meta: dict = {}
        if dims is not None:
            value_meta["dims"] = dims
            grad_meta["dims"] = dims
        if dims_grad is not None:
            value_meta["dims_grad"] = dims_grad
            grad_meta["dims"] = dims_grad

        self._dump_inner(
            name=name,
            value=value,
            extra_kwargs=kwargs,
            save=save,
            enable_value=self._config.enable_value,
            enable_curr_grad=False,
            enable_future_grad=self._config.enable_grad,
            value_tag="Dumper.Value",
            grad_tag="Dumper.Grad",
            value_meta_only_fields=value_meta,
            grad_meta_only_fields=grad_meta,
            grafter_extras=grafter_extras,
        )

    def dump_model(
        self,
        model: "torch.nn.Module",
        name_prefix: str = "param",
        save: bool = True,
        get_grad: Optional[Callable] = None,
        step: Optional[int] = None,
        **kwargs,
    ) -> None:
        for param_name, param in model.named_parameters():
            for plugin in _plugins:
                param_name = (
                    plugin.transform_model_param_name(model, param_name) or param_name
                )
            self._dump_inner(
                name=f"{name_prefix}__{param_name}",
                value=param,
                extra_kwargs=kwargs,
                save=save,
                enable_value=self._config.enable_model_value,
                enable_curr_grad=self._config.enable_model_grad,
                enable_future_grad=False,
                value_tag="Dumper.ParamValue",
                grad_tag="Dumper.ParamGrad",
                get_grad=get_grad,
                step=step,
            )

    def dump_dict(self, name_prefix, data, save: bool = True, **kwargs):
        data = _obj_to_dict(data)
        for name, value in data.items():
            self.dump(f"{name_prefix}_{name}", value, save=save, **kwargs)

    def set_ctx(self, **kwargs):
        """
        Example:

        dumper.configure_default(filter='layer_id=[0-3]')
        dumper.set_ctx(layer_id=self.layer_id)
        ...
        dumper.set_ctx(layer_id=None)
        """
        self._state.global_ctx = {
            k: v for k, v in (self._state.global_ctx | kwargs).items() if v is not None
        }

    def ctx(
        self,
        _extractor: Optional[Callable[..., dict]] = None,
        **static_ctx: Any,
    ) -> Callable:
        """Decorator that sets context before calling the wrapped function and clears it after.

        Two forms:
            @dumper.ctx(lambda self: dict(layer_id=self.layer_id))
            def forward(self, x): ...

            @dumper.ctx(phase="decode")
            def decode_step(self, x): ...
        """
        if _extractor is not None and static_ctx:
            raise ValueError("cannot mix lambda extractor with static kwargs")
        if _extractor is None and not static_ctx:
            raise ValueError("must provide either a lambda or static kwargs")

        def decorator(fn: Callable) -> Callable:
            @functools.wraps(fn)
            def wrapper(*args: Any, **kwargs: Any) -> Any:
                ctx_dict: dict = _extractor(args[0]) if _extractor else static_ctx
                self.set_ctx(**ctx_dict)
                try:
                    return fn(*args, **kwargs)
                finally:
                    self.set_ctx(**{k: None for k in ctx_dict})

            return wrapper

        return decorator

    def apply_source_patches(self) -> None:
        """Apply source patches from DUMPER_SOURCE_PATCHER_CONFIG if set.

        Automatically injects ``from sglang.srt.debug_utils.dumper import dumper``
        into every replacement block so users don't need to write it in YAML.
        """
        config_path = self._config.source_patcher_config
        if not config_path:
            return

        from sglang.srt.debug_utils.source_patcher import apply_patches_from_config

        yaml_content: str = Path(config_path).read_text()
        _log(f"[source_patcher] loading config from {config_path}")
        apply_patches_from_config(
            yaml_content,
            extra_imports=["from sglang.srt.debug_utils.dumper import dumper"],
        )

    def register_non_intrusive_dumper(
        self,
        model: "torch.nn.Module",
    ) -> Optional["_NonIntrusiveDumper"]:
        self._http_manager  # noqa: B018
        mode = self._config.non_intrusive_mode
        if mode == "off":
            return None
        non_intrusive = _NonIntrusiveDumper(dumper=self, model=model, mode=mode)
        self._non_intrusives.append(non_intrusive)
        return non_intrusive

    # ------------------------------- public :: secondary ---------------------------------

    def configure(self, **kwargs) -> None:
        self._config = replace(self._config, **kwargs)

    def configure_default(self, **kwargs) -> None:
        self._config = self._config.with_defaults(**kwargs)

    def reset(self) -> None:
        for non_intrusive in self._non_intrusives:
            non_intrusive.remove()
        self._non_intrusives.clear()
        self._state = _DumperState()

    @contextmanager
    def capture_output(self):
        assert self._state.captured_output_data is None
        self._state.captured_output_data = {}
        try:
            yield self._state.captured_output_data
        finally:
            self._state.captured_output_data = None

    def get_state(self) -> dict:
        return {
            "config": asdict(self._config),
            "dump_index": self._state.dump_index,
            "step": self._state.step,
        }

    @cached_property
    def _http_manager(self) -> Optional["_DumperHttpManager"]:
        if self._config.server_port_parsed is None:
            return None
        return _DumperHttpManager(self)

    # ------------------------- private :: related to dump -----------------------------

    def _dump_inner(
        self,
        *,
        name: str,
        value,
        extra_kwargs: dict,
        save: bool,
        enable_value: bool,
        enable_curr_grad: bool,
        enable_future_grad: bool,
        value_tag: str,
        grad_tag: str,
        value_meta_only_fields: Optional[dict] = None,
        grad_meta_only_fields: Optional[dict] = None,
        grafter_extras: Optional[dict] = None,
        get_grad: Optional[Callable] = None,
        step: Optional[int] = None,
    ) -> None:
        self._http_manager  # noqa: B018

        if not self._config.enable:
            return

        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 not _evaluate_filter(f, tags):
            return

        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)
        self._grafter.maybe_intercept(value=value, tags=tags, extras=grafter_extras)

        if enable_value:
            self._dump_single(
                tag=value_tag,
                tags=tags,
                value=value,
                save=save,
                step=step,
                meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta},
            )

        if enable_curr_grad and isinstance(value, torch.Tensor):
            g = get_grad(value) if get_grad is not None else value.grad
        else:
            g = None

        if g is not None:
            self._dump_single(
                tag=grad_tag,
                tags={**tags, "name": f"grad__{name}"},
                value=g,
                save=save,
                step=step,
                meta_only_fields={**(grad_meta_only_fields or {}), **recompute_meta},
            )

        if enable_future_grad:
            self._register_dump_grad_hook(
                name=name,
                tensor=value,
                extra_kwargs=extra_kwargs,
                save=save,
                meta_only_fields=grad_meta_only_fields or {},
            )

    def _register_dump_grad_hook(
        self,
        *,
        name: str,
        tensor,
        extra_kwargs: dict,
        save: bool,
        meta_only_fields: Optional[dict] = None,
    ) -> None:
        if not isinstance(tensor, torch.Tensor):
            return
        if not tensor.requires_grad:
            return

        captured_step = self._state.step
        captured_tags = dict(
            name=f"grad__{name}",
            **deepcopy(extra_kwargs),
        )
        captured_meta_only = meta_only_fields or {}

        def grad_hook(grad: torch.Tensor) -> None:
            self._dump_single(
                tag="Dumper.Grad",
                tags=captured_tags,
                value=grad,
                save=save,
                step=captured_step,
                meta_only_fields=captured_meta_only,
            )

        tensor.register_hook(grad_hook)

    def _dump_single(
        self,
        *,
        tag: str,
        tags: dict,
        value,
        save: bool,
        step: Optional[int] = None,
        meta_only_fields: Optional[dict] = None,
    ) -> None:
        self._ensure_exp_name()
        self._state.dump_index += 1

        rank = _get_rank()
        full_kwargs = dict(
            step=(step if step is not None else self._state.step),
            rank=rank,
            dump_index=self._state.dump_index,
            **tags,
        )
        if self._config.include_parallel_rank_in_filename:
            full_kwargs.update(_collect_parallel_rank_tags())
        full_filename = _format_tags(full_kwargs) + ".pt"
        path = Path(self._config.dir) / self._config.exp_name / full_filename

        if self._config.enable_output_console:
            _log(
                f"[{tag}] {path} "
                f"type={type(value)} "
                f"shape={value.shape if isinstance(value, torch.Tensor) else None} "
                f"dtype={value.dtype if isinstance(value, torch.Tensor) else None} "
                f"device={value.device if isinstance(value, torch.Tensor) else None} "
                f"id={id(value)} "
                f"sample_value={get_truncated_value(value)}"
            )

        capturing = self._state.captured_output_data is not None
        if save and (self._config.enable_output_file or capturing):
            output_data = {
                "value": value,
                "meta": dict(
                    **full_kwargs,
                    **self._static_meta,
                    **(meta_only_fields or {}),
                ),
            }

            if capturing:
                output_data["value"] = _deepcopy_or_clone(output_data["value"])
                self._state.captured_output_data[tags["name"]] = output_data
            else:
                if (
                    not self._state.cleanup_previous_handled
                    and self._config.cleanup_previous
                ):
                    self._state.cleanup_previous_handled = True
                    _cleanup_old_dumps(
                        Path(self._config.dir), exp_name=self._config.exp_name
                    )

                path.parent.mkdir(parents=True, exist_ok=True)
                _torch_save(output_data, str(path))

    # ------------------------------- private :: misc ---------------------------------

    @cached_property
    def _static_meta(self) -> dict:
        return _compute_static_meta()

    def _ensure_exp_name(self):
        if self._config.exp_name is None:
            name = _get_default_exp_name(
                timeout_seconds=self._config.collective_timeout
            )
            self.configure(exp_name=name)
            _log(f"Choose exp_name={name}")


# -------------------------------------- hook dumper ------------------------------------------


class _NonIntrusiveDumper:
    _NAME_PREFIX = "non_intrusive__"
    _LAYER_NAME_RE = re.compile(r"(?:.+\.)?layers\.(\d+)$")

    def __init__(
        self,
        dumper: _Dumper,
        model: "torch.nn.Module",
        mode: str,
    ):
        self._dumper = dumper
        self._mode = mode
        self._handles: list = []
        self._core_fields: frozenset[str] = frozenset().union(
            *(p.core_fields() for p in _plugins)
        )

        for module_name, module in model.named_modules():
            if ctx := self._detect_module_ctx(module_name, module):
                self._register_ctx_hooks(module, ctx=ctx)

            is_root = module_name == ""
            pre_hook = self._make_forward_pre_hook(
                module_name=module_name, is_root=is_root
            )
            hook = self._make_forward_hook(module_name=module_name, is_root=is_root)
            self._handles += _register_forward_hook_or_replace_fn(
                module,
                pre_hook=pre_hook,
                hook=hook,
                mode="replace_fn" if is_root else "hook",
            )

    def remove(self) -> None:
        for handle in self._handles:
            handle.remove()
        self._handles.clear()

    @classmethod
    def _detect_module_ctx(
        cls, module_name: str, module: "torch.nn.Module"
    ) -> Optional[dict]:
        match = cls._LAYER_NAME_RE.fullmatch(module_name)
        if match:
            for plugin in _plugins:
                layer_id = plugin.detect_layer_id(module)
                if layer_id is not None:
                    return {"layer_id": layer_id}
            return {"layer_id": int(match.group(1))}
        return None

    def _register_ctx_hooks(self, module: "torch.nn.Module", *, ctx: dict) -> None:
        clear_ctx = {k: None for k in ctx}
        self._handles.append(
            module.register_forward_pre_hook(
                lambda _mod, _input, _ctx=ctx: self._dumper.set_ctx(**_ctx)
            )
        )
        self._handles.append(
            module.register_forward_hook(
                lambda _mod, _input, _output, _clear=clear_ctx: self._dumper.set_ctx(
                    **_clear
                )
            )
        )

    def _make_forward_pre_hook(self, *, module_name: str, is_root: bool):
        def _hook(_module, args, kwargs):
            for i, item in enumerate(args):
                self._dump_value(
                    module_name, item, sub_name=f"inputs.{i}", is_root=is_root
                )
            for name, value in kwargs.items():
                self._dump_value(
                    module_name,
                    value,
                    sub_name=f"inputs.{name}",
                    is_root=is_root,
                )

        return _hook

    def _make_forward_hook(self, *, module_name: str, is_root: bool):
        def _hook(_module, input, output):
            if output is not None:
                self._dump_value(module_name, output, sub_name="output", is_root=False)

        return _hook

    def _dump_value(
        self, module_name: str, value: Any, sub_name: str, *, is_root: bool
    ) -> None:
        for key, item in self._convert_value(
            value, skip_forward_batch=(not is_root)
        ).items():
            effective_key = key or sub_name.rsplit(".", 1)[-1]
            if effective_key in self._core_fields:
                self._dumper.dump(effective_key, item)
            elif self._mode == "all":
                parts = [p for p in (module_name, sub_name, key) if p]
                self._dumper.dump(self._NAME_PREFIX + ".".join(parts), item)

    @staticmethod
    def _convert_value(value, *, skip_forward_batch: bool = False) -> dict[str, Any]:
        if isinstance(value, torch.Tensor):
            return {"": value}

        if isinstance(value, (tuple, list)):
            tensors = [t for t in value if isinstance(t, torch.Tensor)]
            if len(tensors) == 1:
                return {"": tensors[0]}
            return {str(i): t for i, t in enumerate(tensors)}

        for plugin in _plugins:
            result = plugin.convert_value(value, skip_forward_batch=skip_forward_batch)
            if result is not None:
                return result

        return {}


def _register_forward_hook_or_replace_fn(
    module: "torch.nn.Module",
    *,
    pre_hook,
    hook,
    mode: str,
) -> list:
    """Attach pre/post forward hooks to *module*.

    mode="hook"       — standard ``register_forward_pre_hook`` / ``register_forward_hook``
                        (fires only via ``__call__``).
    mode="replace_fn" — monkey-patch ``module.forward`` so hooks fire even when
                        callers invoke ``.forward()`` directly (as sglang does for the
                        root model).

    Returns a list of handle objects with a ``.remove()`` method that undoes
    the registration.
    """
    if mode == "hook":
        return [
            module.register_forward_pre_hook(pre_hook, with_kwargs=True),
            module.register_forward_hook(hook),
        ]
    elif mode == "replace_fn":
        original_forward = module.forward

        @functools.wraps(original_forward)
        def _wrapped(*args, **kwargs):
            pre_hook(module, args, kwargs)
            output = original_forward(*args, **kwargs)
            hook(module, args, output)
            return output

        module.forward = _wrapped

        class _Handle:
            def remove(self) -> None:
                assert module.forward is _wrapped
                module.forward = original_forward

        return [_Handle()]
    else:
        raise ValueError(f"Unknown mode {mode!r}")


# -------------------------------------- grafter ------------------------------------------


class _GraftRole(enum.Enum):
    BASELINE = "baseline"
    TARGET = "target"


class _GraftDirection(enum.Enum):
    B2T = "b2t"  # name flows baseline -> target
    T2B = "t2b"  # name flows target -> baseline


@dataclass
class GraftTransformInput:
    """Single argument passed to a user-supplied transform function.

    User transforms have signature::

        def transform(graft_input: GraftTransformInput) -> torch.Tensor: ...

    The dataclass shape lets us add fields (e.g., direction, sender ranks)
    later without breaking existing transforms.
    """

    # Full dumper.dump tags dict (name + recompute_status + extra_kwargs + ctx).
    tags: "dict[str, Any]"
    # One tensor per sender rank, in sender-rank order.
    received_list: "list[torch.Tensor]"
    # Parallel list of per-sender `grafter_extras` (the dict passed to
    # dumper.dump on each sender; None if the sender omitted it).
    received_extras_list: "list[Optional[dict]]"
    # Recv side's local tensor that will be copy_'d into.
    target: "torch.Tensor"


class _Grafter:
    """Cross-system tensor transplant. Triggered silently from dumper.dump.

    Both sides set the SAME grafter_b2t_filter (names that flow baseline ->
    target) and grafter_t2b_filter (names that flow target -> baseline). The
    only per-side difference is grafter_role ("baseline" | "target"), which
    determines whether a name match means send or recv on this side.

    Graft global rank layout: baseline occupies ranks 0..baseline_world-1;
    target occupies ranks baseline_world..baseline_world+target_world-1. Each
    side derives its own rank from its local default PG via dist.get_rank().

    Please refer to TestGrafterE2eExample in tests for an example.
    """

    def __init__(self, *, config: DumperConfig):
        self._config = config
        self._pg = None

    @property
    def enabled(self) -> bool:
        return self._config.grafter_enable

    def maybe_intercept(
        self, *, value: Any, tags: dict, extras: Optional[dict] = None
    ) -> None:
        """Intercept a dumper.dump call. `extras` is per-call auxiliary data
        (e.g., shard layout, dtype hint) that the sender attaches and the
        recv side's transform receives as `received_extras_list`."""
        cfg = self._config
        if not cfg.grafter_enable:
            return

        direction = self._classify_direction(tags)
        if direction is None:
            return

        if not isinstance(value, torch.Tensor):
            _log(
                f"[Grafter] tags={tags} matched grafter_{direction.value}_filter but "
                f"value is not a torch.Tensor (got type={type(value).__name__}); "
                f"skipping graft. Common cause: dumper.dump called with a non-tensor "
                f"value (dict, list, ...) on this name. Either narrow the filter or "
                f"wrap the value in a tensor."
            )
            return

        self._ensure_group()
        role = _GraftRole(cfg.grafter_role)
        is_send = self._is_sender(role=role, direction=direction)

        # all-gather over the graft world; sender ranks contribute (value,
        # extras) tuples, recv ranks contribute None (their local target is
        # private and shouldn't leak). all_gather_object is pickle-routed,
        # so tensor shapes may differ across sender ranks.
        total_world = cfg.grafter_baseline_world_size + cfg.grafter_target_world_size
        my_contribution = (value, extras) if is_send else None
        gathered: list = [None] * total_world
        dist.all_gather_object(gathered, my_contribution, group=self._pg)

        if is_send:
            _log(
                f"[Grafter] send role={role.value} dir={direction.value} "
                f"tags={tags} extras={extras} local={get_tensor_info(value)}"
            )
            return

        sender_contribs = self._sender_slice(direction=direction, gathered=gathered)
        # Pickled CUDA tensors are restored on their original-device name;
        # that may not match this process's local device, so normalize.
        sender_tensors = [
            (c[0].to(value.device) if isinstance(c[0], torch.Tensor) else c[0])
            for c in sender_contribs
        ]
        sender_extras = [c[1] for c in sender_contribs]

        # Transform + copy_ are wrapped: a buggy user transform must NOT
        # crash the whole training/inference run. On error we log the full
        # traceback and skip this graft point; downstream sees the recv
        # side's original tensor unchanged.
        info_before_overridden = get_tensor_info(value)
        try:
            value_to_override = self._apply_transform(
                tags=tags,
                received_list=sender_tensors,
                received_extras_list=sender_extras,
                target=value,
            )
            diff = _compare_tensors_quick(value, value_to_override)
            _log(
                f"[Grafter] recv role={role.value} dir={direction.value} "
                f"tags={tags} n_senders={len(sender_tensors)} "
                f"sender_extras={sender_extras} "
                f"before_overridden={info_before_overridden} "
                f"to_override={get_tensor_info(value_to_override)} "
                f"diff_pre_vs_new={diff}"
            )
            value.copy_(value_to_override)
        except Exception as e:
            _log(
                f"[Grafter] recv role={role.value} dir={direction.value} "
                f"tags={tags} transform/copy_ raised {type(e).__name__}: {e}; "
                f"skipping graft for this call (target tensor unchanged)\n"
                f"{traceback.format_exc()}"
            )

    def _classify_direction(self, tags: dict) -> Optional["_GraftDirection"]:
        cfg = self._config
        match_b2t = self._match(cfg.grafter_b2t_filter, tags)
        match_t2b = self._match(cfg.grafter_t2b_filter, tags)
        if match_b2t and match_t2b:
            raise RuntimeError(
                f"[Grafter] tags={tags} matched BOTH grafter_b2t_filter and grafter_t2b_filter"
            )
        if match_b2t:
            return _GraftDirection.B2T
        if match_t2b:
            return _GraftDirection.T2B
        return None

    @staticmethod
    def _is_sender(*, role: "_GraftRole", direction: "_GraftDirection") -> bool:
        # baseline is the sender for B2T names; target is the sender for T2B.
        return (role == _GraftRole.BASELINE) == (direction == _GraftDirection.B2T)

    def _sender_slice(self, *, direction: "_GraftDirection", gathered: list) -> list:
        cfg = self._config
        if direction == _GraftDirection.B2T:
            return gathered[: cfg.grafter_baseline_world_size]
        return gathered[cfg.grafter_baseline_world_size :]

    @staticmethod
    def _match(expr: Optional[str], tags: dict) -> bool:
        if expr is None:
            return False
        return _evaluate_filter(expr, tags)

    def _ensure_group(self) -> None:
        if self._pg is not None:
            return

        cfg = self._config
        assert (
            dist.is_initialized()
        ), "[Grafter] default torch.distributed must be initialized"
        role = _GraftRole(cfg.grafter_role)
        local_world = dist.get_world_size()
        local_rank = dist.get_rank()
        if role == _GraftRole.BASELINE:
            assert local_world == cfg.grafter_baseline_world_size, (
                f"[Grafter] grafter_baseline_world_size={cfg.grafter_baseline_world_size} "
                f"but dist.get_world_size()={local_world}; they must match on the baseline side"
            )
            global_rank = local_rank
        else:
            assert local_world == cfg.grafter_target_world_size, (
                f"[Grafter] grafter_target_world_size={cfg.grafter_target_world_size} "
                f"but dist.get_world_size()={local_world}; they must match on the target side"
            )
            global_rank = cfg.grafter_baseline_world_size + local_rank
        total_world = cfg.grafter_baseline_world_size + cfg.grafter_target_world_size
        init_method = f"tcp://{cfg.grafter_master_address}:{cfg.grafter_master_port}"
        _log(
            f"[Grafter] init group: role={role.value} "
            f"baseline_world={cfg.grafter_baseline_world_size} "
            f"target_world={cfg.grafter_target_world_size} "
            f"rank={global_rank} init_method={init_method} "
            f"backend={cfg.grafter_backend} name={cfg.grafter_group_name}"
        )
        self._pg = _collective_with_timeout(
            lambda: _init_custom_process_group(
                backend=cfg.grafter_backend,
                init_method=init_method,
                world_size=total_world,
                rank=global_rank,
                group_name=cfg.grafter_group_name,
            ),
            operation_name="_init_custom_process_group in _Grafter",
            timeout_seconds=cfg.grafter_timeout,
        )

    def _apply_transform(
        self,
        *,
        tags: dict,
        received_list: list,
        received_extras_list: list,
        target: torch.Tensor,
    ) -> torch.Tensor:
        # TODO: integrate with dump_comparator unsharder annotations once
        # full inverse (sharded -> global -> sharded) transforms exist.
        graft_input = GraftTransformInput(
            tags=tags,
            received_list=received_list,
            received_extras_list=received_extras_list,
            target=target,
        )
        path = self._config.grafter_transform_path
        fn = self._default_transform if path is None else _load_function(path)
        return fn(graft_input)

    @staticmethod
    def _default_transform(graft_input: GraftTransformInput) -> torch.Tensor:
        """Identity-by-rank fallback. Requires #senders == #recvs and
        shape(received_list[my_recv_rank]) == shape(target). Otherwise raises
        and asks the user for a transform."""
        received_list = graft_input.received_list
        target = graft_input.target
        my_recv_rank = dist.get_rank()
        recv_world_size = dist.get_world_size()
        if len(received_list) != recv_world_size:
            raise RuntimeError(
                _Grafter._default_transform_error(
                    f"requires #senders == #recvs but got "
                    f"#senders={len(received_list)} vs #recvs={recv_world_size}"
                )
            )
        candidate = received_list[my_recv_rank]
        if candidate.shape != target.shape:
            raise RuntimeError(
                _Grafter._default_transform_error(
                    f"requires matching shapes but "
                    f"received_list[{my_recv_rank}].shape={tuple(candidate.shape)} "
                    f"!= target.shape={tuple(target.shape)}"
                )
            )
        return candidate

    @staticmethod
    def _default_transform_error(detail: str) -> str:
        return (
            f"[Grafter] no grafter_transform_path set; default identity-by-rank "
            f"{detail}. Provide a transform via "
            f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol defining "
            f"`transform(graft_input: GraftTransformInput) -> Tensor`."
        )


# -------------------------------------- util fn ------------------------------------------


def _torch_save(value, path: str):
    value = _clone_if_view(value)
    try:
        try:
            return torch.save(value, path)
        except RuntimeError as e:
            if "not pickleable" in str(e):
                stripped = _strip_parameter(value)
                if stripped is not value:
                    _log(f"Observe error={e} and try pickling .data")
                    return _torch_save(stripped, path)
            raise
    except Exception as e:
        _log(f"Observe error={e} when saving data, skip the tensor")


def _map_tensor(value, fn: Callable[[torch.Tensor], torch.Tensor]):
    if isinstance(value, dict):
        return {k: _map_tensor(v, fn) for k, v in value.items()}
    if isinstance(value, torch.Tensor):
        return fn(value)
    return value


def _clone_if_view(value):
    def _fn(t: torch.Tensor) -> torch.Tensor:
        if t.untyped_storage().nbytes() > t.nelement() * t.element_size():
            return t.clone()
        return t

    return _map_tensor(value, _fn)


def _strip_parameter(value):
    def _fn(t: torch.Tensor) -> torch.Tensor:
        if isinstance(t, torch.nn.Parameter):
            return t.data
        return t

    return _map_tensor(value, _fn)


def _collective_with_timeout(fn, operation_name: str, timeout_seconds: int = 60):
    completed = threading.Event()

    def watchdog():
        if not completed.wait(timeout=timeout_seconds):
            _log(
                f"WARNING: '{operation_name}' has not completed after "
                f"{timeout_seconds}s. This usually means not all ranks are "
                f"participating in this collective operation."
            )

    thread = threading.Thread(target=watchdog, daemon=True)
    thread.start()
    try:
        return fn()
    finally:
        completed.set()


def _get_default_exp_name(timeout_seconds: int = 60):
    rank = _get_rank()
    now = time.time()
    ms = int((now % 1) * 1000)
    rand_suffix = random.randint(0, 999)
    object_list = [
        (
            (
                f"{_DEFAULT_EXP_NAME_PREFIX}"
                f"{time.strftime('%Y%m%d_%H%M%S', time.gmtime(now))}"
                f"_{ms:03d}{rand_suffix:03d}"
            )
            if rank == 0
            else None
        )
    ]

    if dist.is_initialized():
        _collective_with_timeout(
            lambda: dist.broadcast_object_list(object_list, device="cuda"),
            operation_name="broadcast_object_list in _get_default_exp_name",
            timeout_seconds=timeout_seconds,
        )

    return object_list[0]


def _cleanup_old_dumps(base_dir: Path, exp_name: Optional[str] = None) -> None:
    import shutil

    if _get_rank() == 0:
        targets = {entry for entry in base_dir.glob(f"{_DEFAULT_EXP_NAME_PREFIX}*")}
        if exp_name:
            targets.add(base_dir / exp_name)
        targets = {d for d in targets if d.is_dir()}

        for entry in targets:
            shutil.rmtree(entry)
            _log(f"Cleaned up {entry}")

    if dist.is_initialized():
        _collective_with_timeout(
            dist.barrier,
            operation_name="barrier in _cleanup_old_dumps",
        )


def _get_rank():
    if dist.is_initialized():
        return dist.get_rank()
    else:
        return 0


def _get_world_size():
    if dist.is_initialized():
        return dist.get_world_size()
    else:
        return 1


def _log(msg: str) -> None:
    """Print a log line tagged with the current rank and wall-clock time."""
    print(f"[Dumper, rank={_get_rank()}, t={time.time():.3f}] {msg}", flush=True)


def _compare_tensors_quick(a: "torch.Tensor", b: "torch.Tensor") -> str:
    """One-line summary of how close two tensors are. Inspired by
    sglang.srt.debug_utils.dump_comparator._compute_and_print_diff;
    intentionally inlined here to keep dumper.py free of cross-file imports.

    Different dtypes are fine — we unify by casting both to fp32, which is
    enough for the order-of-magnitude diff summary we log."""
    if a.shape != b.shape:
        return f"shape mismatch (a={tuple(a.shape)} vs b={tuple(b.shape)})"
    if a.numel() == 0:
        return "empty"
    a_float = a.detach().to(torch.float32)
    b_float = b.detach().to(torch.float32)
    raw_abs = (a_float - b_float).abs()
    max_abs = raw_abs.max().item()
    mean_abs = raw_abs.mean().item()
    rel_diff = _calc_rel_diff(a_float, b_float).item()
    return f"rel_diff={rel_diff:.6g} max_abs={max_abs:.6g} mean_abs={mean_abs:.6g}"


# Copied verbatim from sglang.srt.debug_utils.dump_comparator (originally from
# DeepGEMM). Kept inline here so dumper.py has no cross-file imports.
def _calc_rel_diff(x: "torch.Tensor", y: "torch.Tensor"):
    x, y = x.double(), y.double()
    denominator = (x * x + y * y).sum()
    sim = 2 * (x * y).sum() / denominator
    return 1 - sim


def _obj_to_dict(obj):
    if isinstance(obj, dict):
        return obj
    ret = {}
    for k in dir(obj):
        if k.startswith("__") and k.endswith("__"):
            continue
        try:
            v = getattr(obj, k)
            if not callable(v):
                ret[k] = v
        except Exception:
            # Skip attributes that raise an exception on access
            continue
    return ret


def _materialize_value(value):
    if callable(value):
        value = value()
    return value


_PARALLEL_RANK_KEYS = ("pp_rank", "tp_rank", "cp_rank", "ep_rank", "etp_rank")


def _collect_parallel_rank_tags() -> dict[str, int]:
    """Collect parallel-rank tags from framework plugins for use in dump filenames.

    Merges the ``_PARALLEL_RANK_KEYS`` reported by each plugin's
    ``collect_parallel_info()``; the first plugin to report a given key wins.
    """
    result: dict[str, int] = {}
    for plugin in _plugins:
        info = plugin.collect_parallel_info()
        if not info:
            continue
        for key in _PARALLEL_RANK_KEYS:
            if key in info and key not in result:
                result[key] = info[key]
    return result


def _format_tags(kwargs: dict) -> str:
    return "___".join(f"{k}={v}" for k, v in kwargs.items())


class _DefaultNoneDict(dict):
    """dict subclass that returns None for missing keys, for filter expression eval."""

    def __missing__(self, key: str):
        return None


_FILTER_BUILTINS: dict[str, Any] = {"search": re.search, "match": re.match}


def _evaluate_filter(filter_expr: str, tags: dict[str, Any]) -> bool:
    """Evaluate a Python filter expression against the tags dict.

    Unknown tag keys resolve to None, so `layer_id is None` works when layer_id is absent.
    `re.search` and `re.match` are available as `search()` and `match()`.
    """
    namespace = _DefaultNoneDict(tags)
    namespace.update(_FILTER_BUILTINS)
    return bool(eval(filter_expr, {"__builtins__": {}}, namespace))


def _deepcopy_or_clone(x):
    if isinstance(x, torch.Tensor):
        return x.clone()
    return deepcopy(x)


# -------------------------------------- static meta ------------------------------------------


def _compute_static_meta():
    result = {
        "world_rank": _get_rank(),
        "world_size": _get_world_size(),
    }

    for plugin in _plugins:
        if info := plugin.collect_parallel_info():
            result[f"{plugin.name}_parallel_info"] = info

    for plugin in _plugins:
        tokenizer_path: Optional[str] = plugin.get_tokenizer_path()
        if tokenizer_path is not None:
            result["tokenizer_path"] = tokenizer_path
            break

    return result


# -------------------------------------- http manager ------------------------------------------


class _DumperHttpManager:
    def __init__(self, dumper: "_Dumper"):
        self._dumper = dumper
        http_port = self._dumper._config.server_port_parsed

        rpc_broadcast = _create_zmq_rpc_broadcast(
            self,
            timeout_seconds=self._dumper._config.collective_timeout,
        )

        if _get_rank() == 0:
            assert rpc_broadcast is not None
            self._rpc_broadcast = rpc_broadcast

            if http_port == "reuse":
                _log("Standalone HTTP server disabled, reusing existing ports")
            else:
                _start_http_server(prefix="/dumper/", target=self, http_port=http_port)
                _log(f"HTTP server started on port {http_port}")

    # ------------------------------- public ---------------------------------

    def handle_request(self, *, method: str, body: dict[str, Any]) -> list[dict]:
        return self._rpc_broadcast._handle_request_inner(method=method, body=body)

    # ------------------------------- private ---------------------------------

    def _handle_request_inner(self, *, method: str, body: dict[str, Any]) -> dict:
        if method == "get_state":
            return self._dumper.get_state()
        elif method == "configure":
            self._dumper.configure(**body)
            return {}
        elif method == "reset":
            self._dumper.reset()
            return {}
        else:
            raise ValueError(f"Unknown dumper control method: {method!r}")


# -------------------------------------- http control server ------------------------------------------


def _start_http_server(*, prefix: str, target: object, http_port: int):
    handler_class = _make_http_handler(prefix=prefix, target=target)
    server = HTTPServer(("0.0.0.0", http_port), handler_class)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()


def _make_http_handler(*, prefix: str, target):
    class _HTTPHandler(BaseHTTPRequestHandler):
        def do_POST(self):
            if not self.path.startswith(prefix):
                self.send_error(404)
                return
            method = self.path[len(prefix) :]
            try:
                req_body = self._get_request_body()
                _log(f"HTTP {self.path} {req_body=}")
                result = target.handle_request(method=method, body=req_body)
                resp_body = json.dumps(result).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(resp_body)))
                self.end_headers()
                self.wfile.write(resp_body)
            except Exception as e:
                self.send_error(400, str(e))

        def _get_request_body(self) -> dict:
            content_length = int(self.headers.get("Content-Length", 0))
            if content_length == 0:
                return {}
            return json.loads(self.rfile.read(content_length))

    return _HTTPHandler


# -------------------------------------- zmq rpc ------------------------------------------


def _create_zmq_rpc_broadcast(
    handler, timeout_seconds: int = 60
) -> Optional["_ZmqRpcBroadcast"]:
    """A general-purpose minimal RPC to support broadcasting executions to multi processes"""
    rank = _get_rank()
    world_size = dist.get_world_size() if dist.is_initialized() else 1

    ctx = zmq.Context()
    sock = ctx.socket(zmq.REP)
    sock.bind("tcp://*:0")
    bound_port = int(sock.getsockopt_string(zmq.LAST_ENDPOINT).rsplit(":", 1)[1])
    local_addr = f"tcp://{_get_local_ip_by_remote()}:{bound_port}"

    def serve_loop():
        while True:
            try:
                req = sock_recv(sock)
                result = getattr(handler, req["method"])(*req["args"], **req["kwargs"])
                resp = {"result": result, "error": None}
            except Exception as e:
                _log(f"[ZmqRpc] error inside handler: {e}")
                resp = {"result": None, "error": str(e)}
            sock_send(sock, wrap_as_pickle(resp))

    thread = threading.Thread(target=serve_loop, daemon=True)
    thread.start()
    _log(f"[ZmqRpc] server started at {local_addr}")

    if dist.is_initialized():
        all_addresses = [None] * world_size
        _collective_with_timeout(
            lambda: dist.all_gather_object(all_addresses, local_addr),
            operation_name="all_gather_object in _create_zmq_rpc_broadcast",
            timeout_seconds=timeout_seconds,
        )
    else:
        all_addresses = [local_addr]
    _log(f"[ZmqRpc] all_addresses={all_addresses}")

    if rank == 0:
        handles = []
        for i, addr in enumerate(all_addresses):
            req_socket = ctx.socket(zmq.REQ)
            req_socket.connect(addr)
            handles.append(_ZmqRpcHandle(req_socket, debug_name=f"rank-{i}"))
        return _ZmqRpcBroadcast(handles)
    else:
        return None


class _ZmqRpcHandle:
    """Proxy object to call remote handler methods via ZMQ."""

    def __init__(self, socket, debug_name: str):
        self._socket = socket
        self._debug_name = debug_name

    def __getattr__(self, method_name: str):
        def call(*args, **kwargs):
            sock_send(
                self._socket,
                wrap_as_pickle(
                    {
                        "method": method_name,
                        "args": args,
                        "kwargs": kwargs,
                    }
                ),
            )
            response = sock_recv(self._socket)
            if response["error"]:
                raise RuntimeError(
                    f"RPC error on {self._debug_name}: {response['error']}"
                )
            return response["result"]

        return call


class _RpcBroadcastBase:
    """Base for broadcasting method calls to dumper instance(s)."""

    def __getattr__(self, method_name: str):
        raise NotImplementedError

    def __init__(self, handles: List[_ZmqRpcHandle]):
        self._handles = handles


class _ZmqRpcBroadcast(_RpcBroadcastBase):
    """Broadcasts method calls to all ZMQ RPC handles.

    Returns a list of results, one per rank (ordered by rank).
    """

    def __init__(self, handles: List[_ZmqRpcHandle]):
        self._handles = handles

    def __getattr__(self, method_name: str):
        def call(*args, **kwargs):
            return [
                getattr(handle, method_name)(*args, **kwargs)
                for handle in self._handles
            ]

        return call


# --------------------------------- copied code (avoid dependency) --------------------------------------


def _get_local_ip_by_remote() -> Optional[str]:
    # try ipv4
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))  # Doesn't need to be reachable
        return s.getsockname()[0]
    except Exception:
        pass

    try:
        hostname = socket.gethostname()
        ip = socket.gethostbyname(hostname)
        if ip and ip != "127.0.0.1" and ip != "0.0.0.0":
            return ip
    except Exception:
        pass

    # try ipv6
    try:
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        # Google's public DNS server, see
        # https://developers.google.com/speed/public-dns/docs/using#addresses
        s.connect(("2001:4860:4860::8888", 80))  # Doesn't need to be reachable
        return s.getsockname()[0]
    except Exception:
        _log("Can not get local ip by remote")
    return None


def _load_function(path: str) -> Callable:
    """Resolve a fully-qualified Python path 'pkg.module.symbol' to its object.

    Copied (verbatim, minus the function-registry branch) from
    miles.utils.misc.load_function — kept inline so dumper.py has no
    cross-package dependency.
    """
    import importlib

    module_path, _, attr = path.rpartition(".")
    if not module_path:
        raise ValueError(
            f"_load_function expects 'pkg.module.symbol', got {path!r} "
            f"(missing dotted prefix)"
        )
    module = importlib.import_module(module_path)
    return getattr(module, attr)


def _init_custom_process_group(
    *,
    backend: str,
    init_method: str,
    world_size: int,
    rank: int,
    group_name: str,
    timeout=None,
):
    """Build a fresh torch.distributed process group, separate from the default
    one and any other custom groups (e.g. RLHF weight-update groups). Used by
    the grafter to bridge baseline and target systems.

    Adapted from sglang.srt.utils.common.init_custom_process_group; inlined
    here to keep dumper.py free of cross-file imports.
    """
    from torch.distributed.distributed_c10d import (
        Backend,
        PrefixStore,
        _new_process_group_helper,
        _world,
        default_pg_timeout,
        rendezvous,
    )

    if timeout is None:
        timeout = default_pg_timeout

    rendezvous_iterator = rendezvous(init_method, rank, world_size, timeout=timeout)
    store, rank, world_size = next(rendezvous_iterator)
    store.set_timeout(timeout)
    store = PrefixStore(group_name, store)

    backend_obj = Backend(backend)
    # PyTorch 2.6 renamed `pg_options` to `backend_options`.
    torch_major_minor = tuple(
        int(x) for x in torch.__version__.split("+")[0].split(".")[:2]
    )
    pg_options_param_name = (
        "backend_options" if torch_major_minor >= (2, 6) else "pg_options"
    )
    pg, _ = _new_process_group_helper(
        world_size,
        rank,
        [],
        backend_obj,
        store,
        group_name=group_name,
        **{pg_options_param_name: None},
        timeout=timeout,
    )
    _world.pg_group_ranks[pg] = {i: i for i in range(world_size)}
    return pg


# -------------------------------------- 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
    def name(self) -> str: ...

    @abstractmethod
    def collect_parallel_info(self) -> dict: ...

    @abstractmethod
    def convert_value(
        self, value: Any, *, skip_forward_batch: bool
    ) -> Optional[dict[str, Any]]:
        """Return converted dict, or None if this plugin doesn't handle the value."""
        ...

    @abstractmethod
    def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]:
        """Return 0-indexed layer_id, or None if not detectable."""
        ...

    def core_fields(self) -> frozenset[str]:
        return frozenset()

    def get_tokenizer_path(self) -> Optional[str]:
        return None

    def detect_recompute_status(self) -> _RecomputeStatus:
        return _RecomputeStatus.DISABLED

    def transform_model_param_name(
        self, model: "torch.nn.Module", param_name: str
    ) -> Optional[str]:
        """Return a rewritten parameter name, or None to keep the original.

        Used by ``dump_model`` to canonicalize parameter names across parallel
        layouts (e.g. mapping pipeline-local layer indices to global ones).
        """
        return None


class _SGLangPlugin(_FrameworkPlugin):
    _available = True
    try:
        from sglang.srt import distributed as _dist
        from sglang.srt.layers import dp_attention as _dp_attn
        from sglang.srt.layers.logits_processor import LogitsProcessorOutput
        from sglang.srt.model_executor.forward_batch_info import (
            ForwardBatch,
            PPProxyTensors,
        )
    except ImportError:
        _available = False

    @property
    def name(self) -> str:
        return "sglang"

    def collect_parallel_info(self) -> dict:
        if not self._available:
            return {}

        info = {}

        from sglang.srt.runtime_context import get_parallel

        try:
            parallel = get_parallel()
            info["tp_rank"] = parallel.tp_rank
            info["tp_size"] = parallel.tp_size
            info["pp_rank"] = parallel.pp_rank
            info["pp_size"] = parallel.pp_size
            info["moe_ep_rank"] = parallel.moe_ep_rank
            info["moe_ep_size"] = parallel.moe_ep_size
            info["moe_tp_rank"] = parallel.moe_tp_rank
            info["moe_tp_size"] = parallel.moe_tp_size
            info["moe_dp_rank"] = parallel.moe_dp_rank
            info["moe_dp_size"] = parallel.moe_dp_size
        except (AttributeError, AssertionError):
            info["distributed_error"] = True

        try:
            parallel = get_parallel()
            info["enable_dp_attention"] = self._dp_attn.is_dp_attention_enabled()
            info["attn_tp_rank"] = parallel.attn_tp_rank
            info["attn_tp_size"] = parallel.attn_tp_size
            info["attn_dp_rank"] = self._dp_attn.get_attention_dp_rank()
            info["attn_dp_size"] = self._dp_attn.get_attention_dp_size()
            info["attn_cp_rank"] = parallel.attn_cp_rank
            info["attn_cp_size"] = parallel.attn_cp_size
        except (AttributeError, AssertionError):
            info["dp_attention_error"] = True

        return info

    def convert_value(
        self, value: Any, *, skip_forward_batch: bool
    ) -> Optional[dict[str, Any]]:
        if not self._available:
            return None

        if isinstance(value, self.LogitsProcessorOutput):
            return {"next_token_logits": value.next_token_logits}
        if isinstance(value, self.ForwardBatch):
            if skip_forward_batch:
                return {}
            result = {
                "input_ids": value.input_ids,
                "seq_lens": value.seq_lens,
                "positions": value.positions,
                "req_pool_indices": value.req_pool_indices,
            }
            if value.rids is not None:
                result["rids"] = value.rids
            return result
        if isinstance(value, self.PPProxyTensors):
            return {k: v for k, v in value.tensors.items()}

        return None

    def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]:
        if hasattr(module, "layer_id"):
            return module.layer_id
        return None

    def core_fields(self) -> frozenset[str]:
        return frozenset(
            {"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"}
        )

    def get_tokenizer_path(self) -> Optional[str]:
        if not self._available:
            return None

        try:
            from sglang.srt.runtime_context import get_server_args

            args = get_server_args()
            if args is None:
                return None

            return args.tokenizer_path
        except Exception:
            return None


class _MegatronPlugin(_FrameworkPlugin):
    _available = True
    try:
        from megatron.core import parallel_state as _mpu
        from megatron.core.packed_seq_params import PackedSeqParams
    except ImportError:
        _available = False

    @property
    def name(self) -> str:
        return "megatron"

    def collect_parallel_info(self) -> dict:
        if not self._available:
            return {}

        info = {}
        try:
            info["tp_rank"] = self._mpu.get_tensor_model_parallel_rank()
            info["tp_size"] = self._mpu.get_tensor_model_parallel_world_size()
            info["pp_rank"] = self._mpu.get_pipeline_model_parallel_rank()
            info["pp_size"] = self._mpu.get_pipeline_model_parallel_world_size()
            info["dp_rank"] = self._mpu.get_data_parallel_rank()
            info["dp_size"] = self._mpu.get_data_parallel_world_size()
            info["cp_rank"] = self._mpu.get_context_parallel_rank()
            info["cp_size"] = self._mpu.get_context_parallel_world_size()
            info["vpp_rank"] = self._mpu.get_virtual_pipeline_model_parallel_rank()
            info["vpp_size"] = (
                self._mpu.get_virtual_pipeline_model_parallel_world_size()
            )
            info["ep_rank"] = self._mpu.get_expert_model_parallel_rank()
            info["ep_size"] = self._mpu.get_expert_model_parallel_world_size()
            info["etp_rank"] = self._mpu.get_expert_tensor_parallel_rank()
            info["etp_size"] = self._mpu.get_expert_tensor_parallel_world_size()
            info["edp_rank"] = self._mpu.get_expert_data_parallel_rank()
            info["edp_size"] = self._mpu.get_expert_data_parallel_world_size()
            info["tcp_rank"] = self._mpu.get_tensor_and_context_parallel_rank()
            info["tcp_size"] = self._mpu.get_tensor_and_context_parallel_world_size()
            info["etmp_rank"] = self._mpu.get_expert_tensor_and_model_parallel_rank()
            info["etmp_size"] = (
                self._mpu.get_expert_tensor_and_model_parallel_world_size()
            )
            info["tp_src_rank"] = self._mpu.get_tensor_model_parallel_src_rank()
            info["mp_src_rank"] = self._mpu.get_model_parallel_src_rank()
            info["dp_src_rank"] = self._mpu.get_data_parallel_src_rank()
        except (AttributeError, AssertionError):
            info["megatron_error"] = True

        # Megatron sequence parallel reuses the TP group (no dedicated parallel state API).
        # When sequence_parallel=True, inject sp_rank/sp_size for the comparator unsharder.
        try:
            from megatron.training.global_vars import get_args

            args = get_args()
            if getattr(args, "sequence_parallel", False) and "tp_rank" in info:
                info["sp_rank"] = info["tp_rank"]
                info["sp_size"] = info["tp_size"]
        except (ImportError, AssertionError, AttributeError):
            pass

        return info

    def convert_value(
        self, value: Any, *, skip_forward_batch: bool
    ) -> Optional[dict[str, Any]]:
        if not self._available:
            return None
        if isinstance(value, self.PackedSeqParams):
            return {
                "cu_seqlens_q": value.cu_seqlens_q,
                "cu_seqlens_kv": value.cu_seqlens_kv,
                "qkv_format": value.qkv_format,
            }
        return None

    def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]:
        if hasattr(module, "layer_number"):
            return module.layer_number - 1
        return None

    def core_fields(self) -> frozenset[str]:
        return frozenset(
            {"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

    def transform_model_param_name(
        self, model: "torch.nn.Module", param_name: str
    ) -> Optional[str]:
        """Rewrite pipeline-local layer indices to global ones in a param name.

        With pipeline parallelism, ``model.named_parameters()`` reports layer
        indices local to the current PP stage (e.g. ``layers.0`` on every stage).
        Adding the stage's ``get_transformer_layer_offset`` makes the dumped
        names globally unique and comparable across stages. Returns None (keep the
        original name) when not applicable.
        """
        if not self._available:
            return None

        try:
            pp_size = self._mpu.get_pipeline_model_parallel_world_size()
        except (AttributeError, AssertionError):
            return None
        if pp_size <= 1:
            return None

        config = self._get_model_config(model)
        if config is None:
            return None

        offset = self._get_transformer_layer_offset(config)
        if not offset:
            return None

        def _add_offset(match: "re.Match") -> str:
            return f"layers.{int(match.group(1)) + offset}"

        return re.sub(r"layers\.(\d+)", _add_offset, param_name)

    @staticmethod
    def _get_transformer_layer_offset(config) -> int:
        """Return the PP-stage layer offset for ``config``, or 0 if unavailable."""
        try:
            from megatron.core.transformer.transformer_layer import (
                get_transformer_layer_offset,
            )

            return get_transformer_layer_offset(config)
        except (ImportError, AttributeError, AssertionError):
            return 0

    @staticmethod
    def _get_model_config(model: "torch.nn.Module"):
        """Unwrap nested ``.module`` wrappers to reach the Megatron model config."""
        inner = model
        for _ in range(10):
            if hasattr(inner, "config"):
                return inner.config
            if hasattr(inner, "module"):
                inner = inner.module
            else:
                break
        return None


_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 ------------------------------------------


dumper = _Dumper(config=DumperConfig.from_env())


# -------------------------------------- other utility functions ------------------------------------------


def get_truncated_value(value):
    if value is None:
        return None

    if isinstance(value, tuple):
        return [get_truncated_value(x) for x in value]

    if not isinstance(value, torch.Tensor):
        return value

    if value.numel() < 200:
        return value

    slices = [slice(0, 5) if dim_size > 50 else slice(None) for dim_size in value.shape]
    return value[tuple(slices)]


def get_tensor_info(x):
    """
    from sglang.srt.debug_utils.dumper import get_tensor_info
    """
    if not isinstance(x, torch.Tensor):
        return f"type={type(x)} value={x}"
    min = x.float().min() if x.numel() > 0 else None
    max = x.float().max() if x.numel() > 0 else None
    mean = x.float().mean() if x.numel() > 0 else None
    torch.set_printoptions(precision=10)
    x_sample_head = str(x.flatten()[:5])
    x_sample_tail = str(x.flatten()[-5:])
    torch.set_printoptions(precision=4)
    return (
        f"type={type(x)} "
        f"shape={x.shape} "
        f"dtype={x.dtype} "
        f"device={x.device} "
        f"stride={x.stride()} "
        f"req_grad={x.requires_grad} "
        f"min={min} "
        f"max={max} "
        f"mean={mean} "
        f"x_sample_head={x_sample_head} "
        f"x_sample_tail={x_sample_tail}"
    )