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

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