chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Support multi sharding group on the same dimension in dump comparator (#19601)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:36:48 +08:00
committed by GitHub
parent 46960e65cf
commit ea6ff7b01f
14 changed files with 469 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/test_dumper.py
+3 -3
View File
@@ -2348,7 +2348,7 @@ class TestDumperDims:
)
tensor = torch.randn(4, 8, requires_grad=True)
dumper.dump("hidden", tensor, dims="b h(tp)", dims_grad="b h(tp,partial)")
dumper.dump("hidden", tensor, dims="b h(tp)", dims_grad="b h(tp:partial)")
dumper.step()
tensor.backward(torch.ones_like(tensor))
@@ -2362,10 +2362,10 @@ class TestDumperDims:
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)"
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)"
assert grad_data["meta"]["dims"] == "b h(tp:partial)"
def test_dims_grad_inherits(self, tmp_path) -> None:
dumper = _Dumper(