Support dims annotation and enhance dump loader in dumper (#19276)
This commit is contained in:
fzyzcjy
101
test/registered/debug_utils/comparator/test_dims.py
Normal file
101
test/registered/debug_utils/comparator/test_dims.py
Normal file
@@ -0,0 +1,101 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims import (
|
||||
DimSpec,
|
||||
Ordering,
|
||||
ParallelAxis,
|
||||
Reduction,
|
||||
parse_dim,
|
||||
parse_dims,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestParseDim:
|
||||
def test_plain_name(self) -> None:
|
||||
assert parse_dim("b") == DimSpec(name="b")
|
||||
|
||||
def test_parallel_axis(self) -> None:
|
||||
assert parse_dim("h(tp)") == DimSpec(name="h", parallel=ParallelAxis.TP)
|
||||
|
||||
def test_all_parallel_axes(self) -> None:
|
||||
assert parse_dim("a(tp)").parallel == ParallelAxis.TP
|
||||
assert parse_dim("a(cp)").parallel == ParallelAxis.CP
|
||||
assert parse_dim("a(ep)").parallel == ParallelAxis.EP
|
||||
assert parse_dim("a(sp)").parallel == ParallelAxis.SP
|
||||
|
||||
def test_ordering(self) -> None:
|
||||
assert parse_dim("s(cp,zigzag)").ordering == Ordering.ZIGZAG
|
||||
assert parse_dim("s(cp,natural)").ordering == Ordering.NATURAL
|
||||
|
||||
def test_reduction(self) -> None:
|
||||
assert parse_dim("h(tp,partial)").reduction == Reduction.PARTIAL
|
||||
|
||||
def test_all_modifiers(self) -> None:
|
||||
assert parse_dim("s(cp,zigzag,partial)") == DimSpec(
|
||||
name="s",
|
||||
parallel=ParallelAxis.CP,
|
||||
ordering=Ordering.ZIGZAG,
|
||||
reduction=Reduction.PARTIAL,
|
||||
)
|
||||
|
||||
def test_invalid_token_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("h()")
|
||||
with pytest.raises(ValueError, match="Invalid dim token"):
|
||||
parse_dim("h(tp(x))")
|
||||
|
||||
def test_unknown_modifier_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown modifier"):
|
||||
parse_dim("h(xyz)")
|
||||
with pytest.raises(ValueError, match="Unknown modifier"):
|
||||
parse_dim("h(tp,foobar)")
|
||||
|
||||
def test_multiple_ordering_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Multiple ordering"):
|
||||
parse_dim("s(cp,zigzag,natural)")
|
||||
|
||||
def test_multiple_reduction_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Multiple reduction"):
|
||||
parse_dim("h(tp,partial,partial)")
|
||||
|
||||
|
||||
class TestParseDims:
|
||||
def test_multi_dims(self) -> None:
|
||||
assert parse_dims("b s h d") == [
|
||||
DimSpec(name="b"),
|
||||
DimSpec(name="s"),
|
||||
DimSpec(name="h"),
|
||||
DimSpec(name="d"),
|
||||
]
|
||||
|
||||
def test_single_dim(self) -> None:
|
||||
assert parse_dims("t") == [DimSpec(name="t")]
|
||||
|
||||
def test_mixed_annotated(self) -> None:
|
||||
assert parse_dims("b s(cp,zigzag) h(tp) d") == [
|
||||
DimSpec(name="b"),
|
||||
DimSpec(name="s", parallel=ParallelAxis.CP, ordering=Ordering.ZIGZAG),
|
||||
DimSpec(name="h", parallel=ParallelAxis.TP),
|
||||
DimSpec(name="d"),
|
||||
]
|
||||
|
||||
def test_empty_string_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
parse_dims("")
|
||||
|
||||
def test_whitespace_only_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
parse_dims(" ")
|
||||
|
||||
def test_duplicate_name_raises(self) -> None:
|
||||
with pytest.raises(ValueError, match="Duplicate"):
|
||||
parse_dims("h h")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -1,5 +1,4 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -8,7 +7,6 @@ from sglang.srt.debug_utils.comparator.utils import (
|
||||
argmax_coord,
|
||||
calc_rel_diff,
|
||||
compute_smaller_dtype,
|
||||
load_object,
|
||||
try_unify_shape,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -92,28 +90,5 @@ class TestComputeSmallerDtype:
|
||||
assert compute_smaller_dtype(torch.int32, torch.int64) is None
|
||||
|
||||
|
||||
class TestLoadObject:
|
||||
def test_load_tensor(self, tmp_path):
|
||||
path = tmp_path / "tensor.pt"
|
||||
torch.save(torch.randn(5, 5), path)
|
||||
assert load_object(path).shape == (5, 5)
|
||||
|
||||
def test_load_dict_with_value_key(self, tmp_path):
|
||||
path = tmp_path / "wrapped.pt"
|
||||
tensor = torch.randn(3, 3)
|
||||
torch.save({"value": tensor}, path)
|
||||
result = load_object(path)
|
||||
assert result is not None
|
||||
assert result.shape == (3, 3)
|
||||
|
||||
def test_non_tensor_returns_none(self, tmp_path):
|
||||
path = tmp_path / "tensor.pt"
|
||||
torch.save({"dict": 1}, path)
|
||||
assert load_object(path) is None
|
||||
|
||||
def test_nonexistent_returns_none(self):
|
||||
assert load_object(Path("/nonexistent.pt")) is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
@@ -1,50 +1,58 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
import polars as pl
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.dump_loader import (
|
||||
ValueWithMeta,
|
||||
_add_duplicate_index,
|
||||
_cast_to_polars_dtype,
|
||||
find_row,
|
||||
read_meta,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=30, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestDumpLoader(CustomTestCase):
|
||||
def test_read_meta(self):
|
||||
from sglang.srt.debug_utils.dump_loader import read_meta
|
||||
class TestReadMeta:
|
||||
def test_basic(self, tmp_path):
|
||||
for fn in [
|
||||
"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), tmp_path / fn)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
for fn in [
|
||||
"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(str(tmp_path))
|
||||
assert len(df) == 2
|
||||
assert all(c in df.columns for c in ["step", "rank", "name"])
|
||||
|
||||
df = read_meta(tmpdir)
|
||||
self.assertEqual(len(df), 2)
|
||||
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
|
||||
|
||||
class TestFindRow:
|
||||
def test_single_match(self):
|
||||
df = pl.DataFrame({"id": [1, 2], "name": ["a", "b"], "file": ["f1", "f2"]})
|
||||
self.assertEqual(find_row(df, {"id": 2})["file"], "f2")
|
||||
self.assertIsNone(find_row(df, {"id": 999}))
|
||||
assert find_row(df, {"id": 2})["file"] == "f2"
|
||||
|
||||
df_dup = pl.DataFrame({"id": [1, 1], "file": ["f1", "f2"]})
|
||||
self.assertIsNone(find_row(df_dup, {"id": 1}))
|
||||
def test_no_match(self):
|
||||
df = pl.DataFrame({"id": [1, 2], "name": ["a", "b"], "file": ["f1", "f2"]})
|
||||
assert find_row(df, {"id": 999}) is None
|
||||
|
||||
def test_cast_to_polars_dtype(self):
|
||||
from sglang.srt.debug_utils.dump_loader import _cast_to_polars_dtype
|
||||
def test_ambiguous(self):
|
||||
df = pl.DataFrame({"id": [1, 1], "file": ["f1", "f2"]})
|
||||
assert find_row(df, {"id": 1}) is None
|
||||
|
||||
self.assertEqual(_cast_to_polars_dtype("42", pl.Int64), 42)
|
||||
self.assertEqual(_cast_to_polars_dtype("3.14", pl.Float64), 3.14)
|
||||
|
||||
def test_add_duplicate_index(self):
|
||||
from sglang.srt.debug_utils.dump_loader import _add_duplicate_index
|
||||
class TestCastToPolars:
|
||||
def test_int(self):
|
||||
assert _cast_to_polars_dtype("42", pl.Int64) == 42
|
||||
|
||||
def test_float(self):
|
||||
assert _cast_to_polars_dtype("3.14", pl.Float64) == pytest.approx(3.14)
|
||||
|
||||
|
||||
class TestAddDuplicateIndex:
|
||||
def test_basic(self):
|
||||
df = pl.DataFrame(
|
||||
{
|
||||
"name": ["a", "a", "b"],
|
||||
@@ -53,13 +61,40 @@ class TestDumpLoader(CustomTestCase):
|
||||
}
|
||||
)
|
||||
result = _add_duplicate_index(df)
|
||||
self.assertEqual(
|
||||
result.filter(pl.col("name") == "a")
|
||||
.sort("dump_index")["duplicate_index"]
|
||||
.to_list(),
|
||||
[0, 1],
|
||||
)
|
||||
assert result.filter(pl.col("name") == "a").sort("dump_index")[
|
||||
"duplicate_index"
|
||||
].to_list() == [0, 1]
|
||||
|
||||
|
||||
class TestValueWithMeta:
|
||||
def test_load_dict_format(self, tmp_path) -> None:
|
||||
path = tmp_path / "step=0___rank=0___dump_index=1___name=hidden.pt"
|
||||
tensor = torch.randn(4, 8)
|
||||
torch.save({"value": tensor, "meta": {"custom": "field"}}, path)
|
||||
|
||||
loaded = ValueWithMeta.load(path)
|
||||
assert torch.allclose(loaded.value, tensor)
|
||||
assert loaded.meta["custom"] == "field"
|
||||
assert loaded.meta["name"] == "hidden"
|
||||
assert loaded.meta["rank"] == 0
|
||||
|
||||
def test_load_bare_tensor(self, tmp_path) -> None:
|
||||
path = tmp_path / "step=0___rank=0___dump_index=1___name=bare.pt"
|
||||
tensor = torch.randn(3, 3)
|
||||
torch.save(tensor, path)
|
||||
|
||||
loaded = ValueWithMeta.load(path)
|
||||
assert torch.allclose(loaded.value, tensor)
|
||||
assert loaded.meta["name"] == "bare"
|
||||
|
||||
def test_load_corrupted_file(self, tmp_path) -> None:
|
||||
path = tmp_path / "step=0___rank=0___dump_index=1___name=bad.pt"
|
||||
path.write_text("not a valid pt file")
|
||||
|
||||
loaded = ValueWithMeta.load(path)
|
||||
assert loaded.value is None
|
||||
assert loaded.meta["name"] == "bad"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
@@ -2256,5 +2256,74 @@ class TestNonIntrusiveKwargsModel(_NonIntrusiveTestBase):
|
||||
assert captured["qkv_format"]["value"] == "thd"
|
||||
|
||||
|
||||
class TestDumperDims:
|
||||
def test_dims_in_meta_not_filename(self, tmp_path) -> None:
|
||||
dumper = _make_test_dumper(tmp_path)
|
||||
tensor = torch.randn(4, 8)
|
||||
dumper.dump("hidden", tensor, dims="b h(tp)")
|
||||
dumper.step()
|
||||
|
||||
exp_dir = tmp_path / dumper._config.exp_name
|
||||
pt_files = list(exp_dir.glob("*.pt"))
|
||||
assert len(pt_files) == 1
|
||||
|
||||
assert "dims" not in pt_files[0].stem
|
||||
|
||||
data = torch.load(pt_files[0], weights_only=False)
|
||||
assert "dims" in data["meta"]
|
||||
assert data["meta"]["dims"] == "b h(tp)"
|
||||
|
||||
def test_dims_grad_override(self, tmp_path) -> None:
|
||||
dumper = _Dumper(
|
||||
config=DumperConfig(
|
||||
enable=True,
|
||||
dir=str(tmp_path),
|
||||
enable_http_server=False,
|
||||
enable_grad=True,
|
||||
)
|
||||
)
|
||||
|
||||
tensor = torch.randn(4, 8, requires_grad=True)
|
||||
dumper.dump("hidden", tensor, dims="b h(tp)", dims_grad="b h(tp,partial)")
|
||||
dumper.step()
|
||||
|
||||
tensor.backward(torch.ones_like(tensor))
|
||||
|
||||
exp_dir = tmp_path / dumper._config.exp_name
|
||||
pt_files = sorted(exp_dir.glob("*.pt"))
|
||||
assert len(pt_files) == 2
|
||||
|
||||
value_file = [f for f in pt_files if "grad__" not in f.stem][0]
|
||||
grad_file = [f for f in pt_files if "grad__" in f.stem][0]
|
||||
|
||||
value_data = torch.load(value_file, weights_only=False)
|
||||
assert value_data["meta"]["dims"] == "b h(tp)"
|
||||
assert value_data["meta"]["dims_grad"] == "b h(tp,partial)"
|
||||
|
||||
grad_data = torch.load(grad_file, weights_only=False)
|
||||
assert grad_data["meta"]["dims"] == "b h(tp,partial)"
|
||||
|
||||
def test_dims_grad_inherits(self, tmp_path) -> None:
|
||||
dumper = _Dumper(
|
||||
config=DumperConfig(
|
||||
enable=True,
|
||||
dir=str(tmp_path),
|
||||
enable_http_server=False,
|
||||
enable_grad=True,
|
||||
)
|
||||
)
|
||||
|
||||
tensor = torch.randn(4, 8, requires_grad=True)
|
||||
dumper.dump("hidden", tensor, dims="b h(tp)")
|
||||
dumper.step()
|
||||
|
||||
tensor.backward(torch.ones_like(tensor))
|
||||
|
||||
exp_dir = tmp_path / dumper._config.exp_name
|
||||
grad_file = [f for f in exp_dir.glob("*.pt") if "grad__" in f.stem][0]
|
||||
grad_data = torch.load(grad_file, weights_only=False)
|
||||
assert grad_data["meta"]["dims"] == "b h(tp)"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
Reference in New Issue
Block a user