Refactor dumper and change on_forward_pass_start API (#19065)
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user