Refactor dumper and change on_forward_pass_start API (#19065)

This commit is contained in:
fzyzcjy
2026-02-22 16:03:27 +08:00
committed by GitHub
parent eccee4c48e
commit 5eccc3cff9
8 changed files with 184 additions and 190 deletions

View File

@@ -16,14 +16,12 @@ from sglang.srt.debug_utils.dumper import get_truncated_value
def main(args):
df_target = read_meta(args.target_path)
df_target = df_target.filter(
(pl.col("forward_pass_id") >= args.start_id)
& (pl.col("forward_pass_id") <= args.end_id)
(pl.col("step") >= args.start_id) & (pl.col("step") <= args.end_id)
)
if args.filter:
df_target = df_target.filter(pl.col("filename").str.contains(args.filter))
assert all(
c in df_target.columns
for c in ["rank", "forward_pass_id", "dump_index", "name"]
c in df_target.columns for c in ["rank", "step", "dump_index", "name"]
)
df_baseline = read_meta(args.baseline_path)
@@ -37,14 +35,14 @@ def main(args):
path_target = Path(args.target_path) / row["filename"]
if location_info_of_target_pass_id is not None:
location_info = location_info_of_target_pass_id.get(row["forward_pass_id"])
location_info = location_info_of_target_pass_id.get(row["step"])
if location_info is None:
continue
baseline_forward_pass_id = location_info.baseline_forward_pass_id
baseline_step = location_info.baseline_step
baseline_token_slice = location_info.baseline_token_slice
else:
baseline_forward_pass_id = (
row["forward_pass_id"] - args.start_id + args.baseline_start_id
baseline_step = (
row["step"] - args.start_id + args.baseline_start_id
)
baseline_token_slice = None
@@ -61,11 +59,11 @@ def main(args):
row_baseline = find_row(
df_baseline,
conditions=dict(
forward_pass_id=baseline_forward_pass_id,
step=baseline_step,
**{
k: v
for k, v in row.items()
if k not in ["forward_pass_id", "dump_index", "filename"]
if k not in ["step", "dump_index", "filename"]
},
),
)
@@ -297,7 +295,7 @@ def _comparison_preprocessor(x_baseline, x_target, name):
@dataclass
class LocationInfo:
baseline_forward_pass_id: int
baseline_step: int
baseline_token_slice: slice

View File

@@ -25,8 +25,8 @@ class DumpLoader:
from sglang.srt.debug_utils.dumper import dumper
forward_pass_id = dumper._forward_pass_id
conditions = dict(name=name, forward_pass_id=forward_pass_id, **kwargs)
step = dumper._step
conditions = dict(name=name, step=step, **kwargs)
row = find_row(self._df, conditions=conditions)
assert (
row is not None
@@ -65,7 +65,7 @@ def read_meta(directory):
df = pl.DataFrame(rows)
df = df.with_columns(
pl.col("forward_pass_id").cast(int),
pl.col("step").cast(int),
pl.col("rank").cast(int),
pl.col("dump_index").cast(int),
)

View File

@@ -105,7 +105,8 @@ class _DumperConfig(_FrozenConfig):
@classmethod
def _env_prefix(cls) -> str:
return "SGLANG_DUMPER_"
# 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"]]]:
@@ -125,8 +126,8 @@ class _Dumper:
"""Utility to dump tensors, which can be useful when comparison checking models.
Example usage:
dumper.on_forward_pass_start()
dumper.dump("layer_start__hidden_states", hidden_states, layer_id=self.layer_id)
dumper.step()
Import from non-SGLang system:
```
@@ -136,10 +137,10 @@ class _Dumper:
```
Then run the program:
`SGLANG_DUMPER_ENABLE=1 python ...`
`DUMPER_ENABLE=1 python ...`
Auto-cleanup old dumps before first write:
`SGLANG_DUMPER_CLEANUP_PREVIOUS=1 python ...`
`DUMPER_CLEANUP_PREVIOUS=1 python ...`
Alternatively, disable at startup and configure via HTTP:
1. `python ...`
@@ -158,15 +159,16 @@ class _Dumper:
self._cleanup_previous_handled = not config.cleanup_previous
self._dump_index = 0
self._forward_pass_id = 0
self._step = 0
self._global_ctx: dict = {}
self._captured_output_data: Optional[dict] = None
self._rpc_broadcast: "_RpcBroadcastBase" = _LocalOnlyBroadcast(self)
def on_forward_pass_start(self):
"""This should be called on all ranks."""
# ------------------------------- public :: core ---------------------------------
def step(self):
"""This should be called on all ranks at the end of each iteration."""
# Even if SGLANG_DUMPER_ENABLE=0, users may want to use HTTP endpoint to enable it
self._ensure_http_server()
if not self._config.enable:
@@ -175,109 +177,8 @@ class _Dumper:
# Users may want to `dump` only on some ranks, thus determine name here
self._ensure_partial_name()
self._forward_pass_id += 1
print(
f"[Dumper] [{time.time()}] on_forward_pass_start id={self._forward_pass_id}"
)
def _ensure_http_server(self):
if self._http_server_handled:
return
self._http_server_handled = True
http_port = self._config.server_port_parsed
if http_port is None:
return
rpc_broadcast = _create_zmq_rpc_broadcast(
self,
base_port=get_int_env_var("SGLANG_DUMPER_ZMQ_BASE_PORT", 16800),
timeout_seconds=self._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}")
def _ensure_partial_name(self):
if self._config.partial_name is None:
name = _get_partial_name(timeout_seconds=self._config.collective_timeout)
self.configure(partial_name=name)
print(f"[Dumper] Choose partial_name={name}")
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._global_ctx = {
k: v for k, v in (self._global_ctx | kwargs).items() if v is not None
}
def reset(self) -> None:
self._dump_index = 0
self._forward_pass_id = 0
self._global_ctx = {}
@contextmanager
def capture_output(self):
assert self._captured_output_data is None
self._captured_output_data = {}
try:
yield self._captured_output_data
finally:
self._captured_output_data = None
def get_state(self) -> dict:
return {
"config": asdict(self._config),
"dump_index": self._dump_index,
"forward_pass_id": self._forward_pass_id,
}
def _handle_http_control_request(
self, *, method: str, body: dict[str, Any]
) -> list[dict]:
return self._rpc_broadcast._handle_http_control_request_inner(
method=method, body=body
)
def _handle_http_control_request_inner(
self, *, method: str, body: dict[str, Any]
) -> dict:
if method == "get_state":
return self.get_state()
elif method == "configure":
self.configure(**body)
return {}
elif method == "reset":
self.reset()
return {}
else:
raise ValueError(f"Unknown dumper control method: {method!r}")
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 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)
self._step += 1
print(f"[Dumper] [{time.time()}] step={self._step}")
def dump(self, name: str, value, save: bool = True, **kwargs) -> None:
self._dump_inner(
@@ -312,6 +213,78 @@ class _Dumper:
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._global_ctx = {
k: v for k, v in (self._global_ctx | kwargs).items() if v is not None
}
# ------------------------------- 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:
self._dump_index = 0
self._step = 0
self._global_ctx = {}
@contextmanager
def capture_output(self):
assert self._captured_output_data is None
self._captured_output_data = {}
try:
yield self._captured_output_data
finally:
self._captured_output_data = None
def get_state(self) -> dict:
return {
"config": asdict(self._config),
"dump_index": self._dump_index,
"step": self._step,
}
# ------------------------- public :: only used internally -----------------------------
def _handle_http_control_request(
self, *, method: str, body: dict[str, Any]
) -> list[dict]:
return self._rpc_broadcast._handle_http_control_request_inner(
method=method, body=body
)
def _handle_http_control_request_inner(
self, *, method: str, body: dict[str, Any]
) -> dict:
if method == "get_state":
return self.get_state()
elif method == "configure":
self.configure(**body)
return {}
elif method == "reset":
self.reset()
return {}
else:
raise ValueError(f"Unknown dumper control method: {method!r}")
# ------------------------- private :: related to dump -----------------------------
def _dump_inner(
self,
*,
@@ -339,9 +312,6 @@ class _Dumper:
if not (enable_value or enable_curr_grad or enable_future_grad):
return
if self._forward_pass_id < 1:
print("Dump without on_forward_pass_start()")
value = _materialize_value(value)
if enable_value:
@@ -385,7 +355,7 @@ class _Dumper:
if not tensor.requires_grad:
return
captured_forward_pass_id = self._forward_pass_id
captured_step = self._step
captured_tags = dict(name=f"grad__{name}", **deepcopy(extra_kwargs))
def grad_hook(grad: torch.Tensor) -> None:
@@ -394,7 +364,7 @@ class _Dumper:
tags=captured_tags,
value=grad,
save=save,
forward_pass_id=captured_forward_pass_id,
step=captured_step,
)
tensor.register_hook(grad_hook)
@@ -406,18 +376,14 @@ class _Dumper:
tags: dict,
value,
save: bool,
forward_pass_id: Optional[int] = None,
step: Optional[int] = None,
) -> None:
self._ensure_partial_name()
self._dump_index += 1
rank = _get_rank()
full_kwargs = dict(
forward_pass_id=(
forward_pass_id
if forward_pass_id is not None
else self._forward_pass_id
),
step=(step if step is not None else self._step),
rank=rank,
dump_index=self._dump_index,
**tags,
@@ -458,10 +424,49 @@ class _Dumper:
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()
# Even if DUMPER_ENABLE=0, users may want to use HTTP endpoint to enable it
def _ensure_http_server(self):
if self._http_server_handled:
return
self._http_server_handled = True
http_port = self._config.server_port_parsed
if http_port is None:
return
rpc_broadcast = _create_zmq_rpc_broadcast(
self,
base_port=get_int_env_var("DUMPER_ZMQ_BASE_PORT", 16800),
timeout_seconds=self._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}")
def _ensure_partial_name(self):
if self._config.partial_name is None:
name = _get_partial_name(timeout_seconds=self._config.collective_timeout)
self.configure(partial_name=name)
print(f"[Dumper] Choose partial_name={name}")
# -------------------------------------- util fn ------------------------------------------
def _torch_save(value, path: str):
try:

View File

@@ -630,7 +630,7 @@ async def set_internal_state(obj: SetInternalStateReq, request: Request):
# Do not import `dumper.py` to avoid dependency
if os.environ.get("SGLANG_DUMPER_SERVER_PORT") == "reuse":
if os.environ.get("DUMPER_SERVER_PORT") == "reuse":
@app.api_route("/dumper/{method}", methods=["POST"])
@auth_level(AuthLevel.ADMIN_OPTIONAL)

View File

@@ -12,7 +12,7 @@ def temp_set_env(*, allow_sglang: bool = False, **env_vars: Any):
By default, SGLANG_*/SGL_* keys are rejected — use ``Envs`` descriptors
for those. Pass ``allow_sglang=True`` only for special env vars that
intentionally bypass ``environ.py`` (e.g. ``SGLANG_DUMPER_*``).
intentionally bypass ``environ.py``.
"""
if not allow_sglang:
for key in env_vars:

View File

@@ -104,18 +104,18 @@ class TestEndToEnd(CustomTestCase):
enable_http_server=False,
)
)
dumper.on_forward_pass_start()
dumper.dump("tensor_a", tensor)
dumper.on_forward_pass_start()
dumper.step()
dumper.dump("tensor_b", tensor * 2)
dumper.step()
dump_dirs.append(Path(d) / f"sglang_dump_{dumper._config.partial_name}")
args = Namespace(
baseline_path=str(dump_dirs[0]),
target_path=str(dump_dirs[1]),
start_id=1,
end_id=2,
baseline_start_id=1,
start_id=0,
end_id=1,
baseline_start_id=0,
diff_threshold=1e-3,
filter=None,
)

View File

@@ -17,16 +17,14 @@ class TestDumpLoader(CustomTestCase):
with tempfile.TemporaryDirectory() as tmpdir:
for fn in [
"forward_pass_id=1___rank=0___dump_index=1___name=a.pt",
"forward_pass_id=2___rank=0___dump_index=2___name=b.pt",
"step=1___rank=0___dump_index=1___name=a.pt",
"step=2___rank=0___dump_index=2___name=b.pt",
]:
torch.save(torch.randn(5), Path(tmpdir) / fn)
df = read_meta(tmpdir)
self.assertEqual(len(df), 2)
self.assertTrue(
all(c in df.columns for c in ["forward_pass_id", "rank", "name"])
)
self.assertTrue(all(c in df.columns for c in ["step", "rank", "name"]))
def test_find_row(self):
from sglang.srt.debug_utils.dump_loader import find_row

View File

@@ -57,21 +57,21 @@ class TestDumperConfig:
assert _DumperConfig.from_env() == _DumperConfig()
def test_from_env_bool(self):
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_ENABLE="1"):
with temp_set_env(DUMPER_ENABLE="1"):
assert _DumperConfig.from_env().enable is True
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_ENABLE="false"):
with temp_set_env(DUMPER_ENABLE="false"):
assert _DumperConfig.from_env().enable is False
def test_from_env_str(self):
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_FILTER="layer_id=0"):
with temp_set_env(DUMPER_FILTER="layer_id=0"):
assert _DumperConfig.from_env().filter == "layer_id=0"
def test_from_env_dir(self):
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_DIR="/my/dir"):
with temp_set_env(DUMPER_DIR="/my/dir"):
assert _DumperConfig.from_env().dir == "/my/dir"
def test_from_env_int(self):
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_COLLECTIVE_TIMEOUT="120"):
with temp_set_env(DUMPER_COLLECTIVE_TIMEOUT="120"):
assert _DumperConfig.from_env().collective_timeout == 120
def test_configure_overrides(self):
@@ -92,7 +92,7 @@ class TestDumperConfig:
_DumperConfig(filter=123)
def test_configure_default_skips_when_env_set(self):
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_FILTER="from_env"):
with temp_set_env(DUMPER_FILTER="from_env"):
d = _Dumper(config=_DumperConfig.from_env())
d.configure_default(filter="from_code")
assert d._config.filter == "from_env"
@@ -195,9 +195,8 @@ class TestCollectiveTimeout:
class TestDumperDistributed:
def test_basic(self, tmp_path):
with temp_set_env(
allow_sglang=True,
SGLANG_DUMPER_ENABLE="1",
SGLANG_DUMPER_DIR=str(tmp_path),
DUMPER_ENABLE="1",
DUMPER_DIR=str(tmp_path),
):
run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path))
@@ -205,21 +204,21 @@ class TestDumperDistributed:
def _test_basic_func(rank, tmpdir):
tensor = torch.randn(10, 10, device=f"cuda:{rank}")
dumper.on_forward_pass_start()
dumper.dump("tensor_a", tensor, arg=100)
dumper.step()
dumper.on_forward_pass_start()
dumper.set_ctx(ctx_arg=200)
dumper.dump("tensor_b", tensor)
dumper.set_ctx(ctx_arg=None)
dumper.step()
dumper.on_forward_pass_start()
dumper.configure(filter=r"^$")
dumper.dump("tensor_skip", tensor)
dumper.configure(filter=None)
dumper.step()
dumper.on_forward_pass_start()
dumper.dump_dict("obj", {"a": torch.randn(3, device=f"cuda:{rank}"), "b": 42})
dumper.step()
dist.barrier()
filenames = _get_filenames(tmpdir)
@@ -230,7 +229,7 @@ class TestDumperDistributed:
)
def test_collective_timeout(self):
with temp_set_env(allow_sglang=True, SGLANG_DUMPER_ENABLE="1"):
with temp_set_env(DUMPER_ENABLE="1"):
run_distributed_test(self._test_collective_timeout_func)
@staticmethod
@@ -246,7 +245,7 @@ class TestDumperDistributed:
with _capture_stdout() as captured:
if rank != 0:
time.sleep(6)
dumper.on_forward_pass_start()
dumper.step()
output = captured.getvalue()
print(f"Rank {rank} captured output: {output!r}")
@@ -257,9 +256,8 @@ class TestDumperDistributed:
def test_file_content_correctness(self, tmp_path):
with temp_set_env(
allow_sglang=True,
SGLANG_DUMPER_ENABLE="1",
SGLANG_DUMPER_DIR=str(tmp_path),
DUMPER_ENABLE="1",
DUMPER_DIR=str(tmp_path),
):
run_distributed_test(self._test_file_content_func, tmpdir=str(tmp_path))
@@ -267,8 +265,8 @@ class TestDumperDistributed:
def _test_file_content_func(rank, tmpdir):
tensor = torch.arange(12, device=f"cuda:{rank}").reshape(3, 4).float()
dumper.on_forward_pass_start()
dumper.dump("content_check", tensor)
dumper.step()
dist.barrier()
path = _find_dump_file(tmpdir, rank=rank, name="content_check")
@@ -283,19 +281,18 @@ class TestDumperDistributed:
class TestDumperFileWriteControl:
def test_filter(self, tmp_path):
with temp_set_env(
allow_sglang=True,
SGLANG_DUMPER_ENABLE="1",
SGLANG_DUMPER_DIR=str(tmp_path),
SGLANG_DUMPER_FILTER="name=keep",
DUMPER_ENABLE="1",
DUMPER_DIR=str(tmp_path),
DUMPER_FILTER="name=keep",
):
run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))
@staticmethod
def _test_filter_func(rank, tmpdir):
dumper.on_forward_pass_start()
dumper.dump("keep_this", torch.randn(5, device=f"cuda:{rank}"))
dumper.dump("skip_this", torch.randn(5, device=f"cuda:{rank}"))
dumper.dump("not_keep_this", torch.randn(5, device=f"cuda:{rank}"))
dumper.step()
dist.barrier()
filenames = _get_filenames(tmpdir)
@@ -307,16 +304,15 @@ class TestDumperFileWriteControl:
def test_save_false(self, tmp_path):
with temp_set_env(
allow_sglang=True,
SGLANG_DUMPER_ENABLE="1",
SGLANG_DUMPER_DIR=str(tmp_path),
DUMPER_ENABLE="1",
DUMPER_DIR=str(tmp_path),
):
run_distributed_test(self._test_save_false_func, tmpdir=str(tmp_path))
@staticmethod
def _test_save_false_func(rank, tmpdir):
dumper.on_forward_pass_start()
dumper.dump("no_save_tensor", torch.randn(5, device=f"cuda:{rank}"), save=False)
dumper.step()
dist.barrier()
assert len(_get_filenames(tmpdir)) == 0
@@ -422,7 +418,7 @@ class TestDumpDictFormat:
meta = raw["meta"]
assert meta["name"] == "fmt_test"
assert meta["custom_key"] == "hello"
assert "forward_pass_id" in meta
assert "step" in meta
assert "rank" in meta
assert "dump_index" in meta
@@ -448,9 +444,7 @@ def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
enable_http_server=False,
**overrides,
)
d = _Dumper(config=config)
d.on_forward_pass_start()
return d
return _Dumper(config=config)
def _get_filenames(tmpdir):
@@ -584,18 +578,18 @@ class TestDumpGrad:
not_exist=["grad__"],
)
def test_dump_grad_captures_forward_pass_id(self, tmp_path):
def test_dump_grad_captures_step(self, tmp_path):
d = _make_test_dumper(tmp_path, enable_grad=True)
d._forward_pass_id = 42
d._step = 42
x = torch.randn(3, 3, requires_grad=True)
y = (x * 2).sum()
d.dump("id_test", x)
d._forward_pass_id = 999
d._step = 999
y.backward()
grad_file = _find_dump_file(tmp_path, name="grad__id_test")
assert "forward_pass_id=42" in grad_file.name
assert "step=42" in grad_file.name
def test_dump_grad_file_content(self, tmp_path):
d = _make_test_dumper(tmp_path, enable_grad=True)
@@ -807,7 +801,7 @@ class TestReset:
d.reset()
assert d._dump_index == 0
assert d._forward_pass_id == 0
assert d._step == 0
assert d._global_ctx == {}
def test_dump_works_after_reset(self, tmp_path):
@@ -815,7 +809,6 @@ class TestReset:
d.dump("pre", torch.randn(3, 3))
d.reset()
d.on_forward_pass_start()
d.dump("post", torch.randn(3, 3))
filenames = _get_filenames(tmp_path)
@@ -847,7 +840,7 @@ class TestDumperHttp:
thread.join(timeout=10)
else:
base_url = DEFAULT_URL_FOR_TEST
env = {**os.environ, "SGLANG_DUMPER_SERVER_PORT": "reuse"}
env = {**os.environ, "DUMPER_SERVER_PORT": "reuse"}
proc = popen_launch_server(
"Qwen/Qwen3-0.6B",
base_url,
@@ -863,7 +856,7 @@ class TestDumperHttp:
@staticmethod
def _standalone_mode_worker(rank, http_port: int, stop_event):
dumper.configure(enable=False, server_port=str(http_port))
dumper.on_forward_pass_start()
dumper.step()
stop_event.wait()
@staticmethod
@@ -927,7 +920,7 @@ class TestDumperHttp:
self._post(dumper_http_url, "reset")
states = self._post(dumper_http_url, "get_state")
self._assert_all_ranks(states, "dump_index", 0)
self._assert_all_ranks(states, "forward_pass_id", 0)
self._assert_all_ranks(states, "step", 0)
def test_get_state(self, dumper_http_url: str):
self._post(dumper_http_url, "configure", enable=True, filter="layer_id=[0-3]")
@@ -936,7 +929,7 @@ class TestDumperHttp:
self._assert_all_ranks(states, "config.filter", "layer_id=[0-3]")
for state in states:
assert "dump_index" in state
assert "forward_pass_id" in state
assert "step" in state
def test_all_ranks_consistent(self, dumper_http_url: str):
self._post(dumper_http_url, "configure", enable=True, dir="/tmp/multi")