import enum import functools import json import os import random import re import socket import threading import time 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 # -------------------------------------- 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 @classmethod def _env_prefix(cls) -> str: # NOTE: should not be `SGLANG_DUMPER_`, otherwise it is weird when dumping Megatron in Miles return "DUMPER_" @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"] = [] # ------------------------------- 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 print(f"[Dumper] [{time.time()}] step={self._state.step}") def dump( self, name: str, value, save: bool = True, dims: Optional[str] = None, dims_grad: Optional[str] = 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, ) def dump_model( self, model: "torch.nn.Module", name_prefix: str = "param", save: bool = True, **kwargs, ) -> None: for param_name, param in model.named_parameters(): 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", ) 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() print(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, ) -> 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) if enable_value: self._dump_single( tag=value_tag, tags=tags, value=value, save=save, meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta}, ) if ( enable_curr_grad and isinstance(value, torch.Tensor) and (g := value.grad) is not None ): self._dump_single( tag=grad_tag, tags={**tags, "name": f"grad__{name}"}, value=g, save=save, 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, ) full_filename = _format_tags(full_kwargs) + ".pt" path = Path(self._config.dir) / self._config.exp_name / full_filename if self._config.enable_output_console: print( f"[{tag}] [{rank}, {time.time()}] {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) print(f"[Dumper] 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}") # -------------------------------------- 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: print(f"[Dumper] Observe error={e} and try pickling .data") return _torch_save(stripped, path) raise except Exception as e: print(f"[Dumper] 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): print( f"\n[Dumper] WARNING: '{operation_name}' has not completed after " f"{timeout_seconds}s. This usually means not all ranks are " f"participating in this collective operation.\n", flush=True, ) 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) print(f"[Dumper] 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 _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 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": print( "[Dumper] Standalone HTTP server disabled, reusing existing ports" ) else: _start_http_server(prefix="/dumper/", target=self, http_port=http_port) print(f"[Dumper] 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() print(f"[Dumper#{_get_rank()}] 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""" import zmq 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_pyobj() result = getattr(handler, req["method"])(*req["args"], **req["kwargs"]) resp = {"result": result, "error": None} except Exception as e: print(f"[Dumper.ZmqRpc] error inside handler: {e}") resp = {"result": None, "error": str(e)} sock.send_pyobj(resp) thread = threading.Thread(target=serve_loop, daemon=True) thread.start() print(f"[Dumper.ZmqRpc] rank={rank} 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] print(f"[Dumper.ZmqRpc] rank={rank} 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): self._socket.send_pyobj( { "method": method_name, "args": args, "kwargs": kwargs, } ) response = self._socket.recv_pyobj() 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: print("Can not get local ip by remote") return None # -------------------------------------- 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 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 = {} try: info["tp_rank"] = self._dist.get_tensor_model_parallel_rank() info["tp_size"] = self._dist.get_tensor_model_parallel_world_size() info["pp_rank"] = self._dist.get_pipeline_model_parallel_rank() info["pp_size"] = self._dist.get_pipeline_model_parallel_world_size() info["moe_ep_rank"] = self._dist.get_moe_expert_parallel_rank() info["moe_ep_size"] = self._dist.get_moe_expert_parallel_world_size() info["moe_tp_rank"] = self._dist.get_moe_tensor_parallel_rank() info["moe_tp_size"] = self._dist.get_moe_tensor_parallel_world_size() info["moe_dp_rank"] = self._dist.get_moe_data_parallel_rank() info["moe_dp_size"] = self._dist.get_moe_data_parallel_world_size() except (AttributeError, AssertionError): info["distributed_error"] = True try: info["enable_dp_attention"] = self._dp_attn.is_dp_attention_enabled() info["attn_tp_rank"] = self._dp_attn.get_attention_tp_rank() info["attn_tp_size"] = self._dp_attn.get_attention_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["local_attn_dp_rank"] = self._dp_attn.get_local_attention_dp_rank() info["local_attn_dp_size"] = self._dp_attn.get_local_attention_dp_size() info["attn_cp_rank"] = self._dp_attn.get_attention_cp_rank() info["attn_cp_size"] = self._dp_attn.get_attention_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.server_args import get_global_server_args args = get_global_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 _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}" )