From 9a7d8d5eb025d4686f284815ee835cfaf73207eb Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 16 Feb 2026 13:31:19 +0800 Subject: [PATCH] Collect upper level metadata to dump output (#18880) --- python/sglang/srt/debug_utils/dumper.py | 116 ++++++++++++++++++++- test/registered/debug_utils/test_dumper.py | 51 +++++++++ 2 files changed, 166 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index aaa0b937d..3be898034 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -4,6 +4,7 @@ import re import socket import threading import time +from functools import cached_property from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path from typing import List, Optional @@ -166,10 +167,14 @@ class _Dumper: path.parent.mkdir(parents=True, exist_ok=True) output_data = { "value": value.data if isinstance(value, torch.nn.Parameter) else value, - "meta": full_kwargs, + "meta": dict(**full_kwargs, **self._static_meta), } _torch_save(output_data, str(path)) + @cached_property + def _static_meta(self) -> dict: + return _compute_static_meta() + def _torch_save(value, path: str): try: @@ -201,6 +206,13 @@ def _get_rank(): 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 @@ -218,6 +230,108 @@ def _obj_to_dict(obj): return ret +# -------------------------------------- static metadata ------------------------------------------ + + +def _compute_static_meta(): + result = { + "world_rank": _get_rank(), + "world_size": _get_world_size(), + } + + if x := _collect_sglang_parallel_info(): + result["sglang_parallel_info"] = x + if x := _collect_megatron_parallel_info(): + result["megatron_parallel_info"] = x + + return result + + +def _collect_sglang_parallel_info(): + info = {} + + try: + from sglang.srt.distributed import ( + get_moe_expert_parallel_rank, + get_moe_expert_parallel_world_size, + get_moe_tensor_parallel_rank, + get_moe_tensor_parallel_world_size, + get_pipeline_model_parallel_rank, + get_pipeline_model_parallel_world_size, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + + info["tp_rank"] = get_tensor_model_parallel_rank() + info["tp_size"] = get_tensor_model_parallel_world_size() + info["pp_rank"] = get_pipeline_model_parallel_rank() + info["pp_size"] = get_pipeline_model_parallel_world_size() + info["moe_ep_rank"] = get_moe_expert_parallel_rank() + info["moe_ep_size"] = get_moe_expert_parallel_world_size() + info["moe_tp_rank"] = get_moe_tensor_parallel_rank() + info["moe_tp_size"] = get_moe_tensor_parallel_world_size() + except (ImportError, AttributeError, AssertionError): + info["distributed_error"] = True + + try: + from sglang.srt.layers.dp_attention import ( + get_attention_dp_rank, + get_attention_dp_size, + get_attention_tp_rank, + get_attention_tp_size, + get_local_attention_dp_rank, + get_local_attention_dp_size, + is_dp_attention_enabled, + ) + + info["enable_dp_attention"] = is_dp_attention_enabled() + info["attn_tp_rank"] = get_attention_tp_rank() + info["attn_tp_size"] = get_attention_tp_size() + info["attn_dp_rank"] = get_attention_dp_rank() + info["attn_dp_size"] = get_attention_dp_size() + info["local_attn_dp_rank"] = get_local_attention_dp_rank() + info["local_attn_dp_size"] = get_local_attention_dp_size() + except (ImportError, AttributeError, AssertionError): + info["dp_attention_error"] = True + + return info + + +def _collect_megatron_parallel_info(): + info = {} + + try: + from megatron.core import parallel_state as mpu + + info["tp_rank"] = mpu.get_tensor_model_parallel_rank() + info["tp_size"] = mpu.get_tensor_model_parallel_world_size() + info["pp_rank"] = mpu.get_pipeline_model_parallel_rank() + info["pp_size"] = mpu.get_pipeline_model_parallel_world_size() + info["dp_rank"] = mpu.get_data_parallel_rank() + info["dp_size"] = mpu.get_data_parallel_world_size() + info["cp_rank"] = mpu.get_context_parallel_rank() + info["cp_size"] = mpu.get_context_parallel_world_size() + info["vpp_rank"] = mpu.get_virtual_pipeline_model_parallel_rank() + info["vpp_size"] = mpu.get_virtual_pipeline_model_parallel_world_size() + info["ep_rank"] = mpu.get_expert_model_parallel_rank() + info["ep_size"] = mpu.get_expert_model_parallel_world_size() + info["etp_rank"] = mpu.get_expert_tensor_parallel_rank() + info["etp_size"] = mpu.get_expert_tensor_parallel_world_size() + info["edp_rank"] = mpu.get_expert_data_parallel_rank() + info["edp_size"] = mpu.get_expert_data_parallel_world_size() + info["tcp_rank"] = mpu.get_tensor_and_context_parallel_rank() + info["tcp_size"] = mpu.get_tensor_and_context_parallel_world_size() + info["etmp_rank"] = mpu.get_expert_tensor_and_model_parallel_rank() + info["etmp_size"] = mpu.get_expert_tensor_and_model_parallel_world_size() + info["tp_src_rank"] = mpu.get_tensor_model_parallel_src_rank() + info["mp_src_rank"] = mpu.get_model_parallel_src_rank() + info["dp_src_rank"] = mpu.get_data_parallel_src_rank() + except (ImportError, AttributeError, AssertionError): + info["megatron_error"] = True + + return info + + # -------------------------------------- http control server ------------------------------------------ diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index f5b492c2b..22cb8c25d 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -8,6 +8,8 @@ import torch import torch.distributed as dist from sglang.srt.debug_utils.dumper import ( + _collect_megatron_parallel_info, + _collect_sglang_parallel_info, _Dumper, _obj_to_dict, _torch_save, @@ -313,5 +315,54 @@ def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path: return matches[0] +class TestSaveValue: + def test_dump_output_format(self, tmp_path): + dumper = _make_test_dumper(tmp_path) + tensor = torch.randn(4, 4) + + dumper.dump("dict_test", tensor) + + path = _find_dump_file(tmp_path, rank=0, name="dict_test") + loaded = _load_dump(path) + assert torch.equal(loaded["value"], tensor) + assert loaded["meta"]["name"] == "dict_test" + assert loaded["meta"]["rank"] == 0 + + +class TestStaticMetadata: + def test_static_meta_contains_world_info(self): + dumper = _make_test_dumper(Path("/tmp")) + meta = dumper._static_meta + assert "world_rank" in meta + assert "world_size" in meta + assert meta["world_rank"] == 0 + assert meta["world_size"] == 1 + + def test_static_meta_caching(self): + dumper = _make_test_dumper(Path("/tmp")) + meta1 = dumper._static_meta + meta2 = dumper._static_meta + assert meta1 is meta2 + + def test_parallel_info_graceful_fallback(self): + sglang_info = _collect_sglang_parallel_info() + assert isinstance(sglang_info, dict) + + megatron_info = _collect_megatron_parallel_info() + assert isinstance(megatron_info, dict) + + def test_dump_includes_static_meta(self, tmp_path): + dumper = _make_test_dumper(tmp_path) + tensor = torch.randn(2, 2) + + dumper.dump("meta_test", tensor) + + path = _find_dump_file(tmp_path, rank=0, name="meta_test") + loaded = _load_dump(path) + meta = loaded["meta"] + assert "world_rank" in meta + assert "world_size" in meta + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))