Support dims annotation and enhance dump loader in dumper (#19276)

2026-02-25 09:41:48 +08:00
parent 8b1ab4aaf9
commit 02ca107b2c
9 changed files with 437 additions and 107 deletions
--- a/python/sglang/srt/debug_utils/comparator/dims.py
+++ b/python/sglang/srt/debug_utils/comparator/dims.py
@@ -0,0 +1,80 @@
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+
+class ParallelAxis(Enum):
+    TP = "tp"
+    CP = "cp"
+    EP = "ep"
+    SP = "sp"
+
+
+class Ordering(Enum):
+    ZIGZAG = "zigzag"
+    NATURAL = "natural"
+
+
+class Reduction(Enum):
+    PARTIAL = "partial"
+
+
+@dataclass(frozen=True)
+class DimSpec:
+    name: str
+    parallel: Optional[ParallelAxis] = None
+    ordering: Optional[Ordering] = None
+    reduction: Optional[Reduction] = None
+
+
+_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\((?P<modifiers>[^)]+)\))?$")
+
+_MODIFIER_FIELDS: list[tuple[type[Enum], str]] = [
+    (ParallelAxis, "parallel"),
+    (Ordering, "ordering"),
+    (Reduction, "reduction"),
+]
+
+_MODIFIER_LOOKUP: dict[str, tuple[str, Enum]] = {}
+for _enum_cls, _field in _MODIFIER_FIELDS:
+    for _member in _enum_cls:
+        _MODIFIER_LOOKUP[_member.value] = (_field, _member)
+
+
+def parse_dim(token: str) -> DimSpec:
+    match = _DIM_PATTERN.match(token)
+    if match is None:
+        raise ValueError(f"Invalid dim token: {token!r}")
+
+    name = match.group("name")
+    modifiers_str = match.group("modifiers")
+
+    if modifiers_str is None:
+        return DimSpec(name=name)
+
+    fields: dict[str, Enum] = {}
+    for part in (p.strip() for p in modifiers_str.split(",")):
+        if part not in _MODIFIER_LOOKUP:
+            raise ValueError(f"Unknown modifier {part!r} in dim spec: {token!r}")
+        field_name, enum_value = _MODIFIER_LOOKUP[part]
+        if field_name in fields:
+            raise ValueError(f"Multiple {field_name} values in dim token: {token!r}")
+        fields[field_name] = enum_value
+
+    return DimSpec(name=name, **fields)
+
+
+def parse_dims(dims_str: str) -> list[DimSpec]:
+    """Parse 'b s(cp,zigzag) h(tp) d' -> list[DimSpec]."""
+    if not dims_str.strip():
+        raise ValueError("dims string must not be empty")
+
+    result = [parse_dim(token) for token in dims_str.strip().split()]
+
+    names = [spec.name for spec in result]
+    if len(names) != len(set(names)):
+        duplicates = sorted({n for n in names if names.count(n) > 1})
+        raise ValueError(f"Duplicate dim names: {duplicates}")
+
+    return result
--- a/python/sglang/srt/debug_utils/comparator/entrypoint.py
+++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py
@@ -1,7 +1,9 @@
 import argparse
 from pathlib import Path
+from typing import Optional

 import polars as pl
+import torch

 from sglang.srt.debug_utils.comparator.output_types import (
    ComparisonRecord,
@@ -11,8 +13,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
    print_record,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparison import compare_tensors
-from sglang.srt.debug_utils.comparator.utils import load_object
-from sglang.srt.debug_utils.dump_loader import find_row, read_meta
+from sglang.srt.debug_utils.dump_loader import ValueWithMeta, find_row, read_meta


 def main() -> None:
@@ -70,8 +71,8 @@ def run(args: argparse.Namespace) -> None:

        path_baseline = Path(args.baseline_path) / row_baseline["filename"]

-        x_baseline = load_object(path_baseline)
-        x_target = load_object(path_target)
+        x_baseline = _load_tensor(path_baseline)
+        x_target = _load_tensor(path_target)

        if x_baseline is None or x_target is None:
            counts["skipped"] += 1
@@ -104,6 +105,13 @@ def run(args: argparse.Namespace) -> None:
    )


+def _load_tensor(path: Path) -> Optional[torch.Tensor]:
+    loaded = ValueWithMeta.load(path)
+    if not isinstance(loaded.value, torch.Tensor):
+        return None
+    return loaded.value
+
+
 def _parse_args() -> argparse.Namespace:
    # python -m sglang.srt.debug_utils.comparator --baseline-path ... --target-path ...
    parser = argparse.ArgumentParser()
--- a/python/sglang/srt/debug_utils/comparator/utils.py
+++ b/python/sglang/srt/debug_utils/comparator/utils.py
@@ -1,5 +1,4 @@
 import functools
-from pathlib import Path
 from typing import Optional, Tuple

 import torch
@@ -42,19 +41,3 @@ def calc_rel_diff(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    denominator = (x * x + y * y).sum()
    sim = 2 * (x * y).sum() / denominator
    return 1 - sim
-
-
-def load_object(path: Path) -> Optional[torch.Tensor]:
-    try:
-        x = torch.load(path, weights_only=False)
-    except Exception as e:
-        print(f"Skip load {path} since error {e}")
-        return None
-
-    if isinstance(x, dict) and "value" in x:
-        x = x["value"]
-
-    if not isinstance(x, torch.Tensor):
-        print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
-        return None
-    return x.cuda()
--- a/python/sglang/srt/debug_utils/dump_loader.py
+++ b/python/sglang/srt/debug_utils/dump_loader.py
@@ -1,11 +1,60 @@
 import functools
 import os
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, Tuple

 import polars as pl
 import torch

+_TYPED_FIELDS: list[tuple[str, type]] = [("rank", int)]
+
+
+def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
+    stem = Path(path).stem
+    result: Dict[str, Any] = {}
+    for kv in stem.split("___"):
+        if "=" in kv:
+            k, v = kv.split("=", 1)
+            result[k] = v
+    for field, converter in _TYPED_FIELDS:
+        if field in result:
+            result[field] = converter(result[field])
+    return result
+
+
+@dataclass
+class ValueWithMeta:
+    value: Any
+    meta: Dict[str, Any]
+
+    @staticmethod
+    def load(path: Path) -> "ValueWithMeta":
+        path = Path(path)
+        meta_from_filename = parse_meta_from_filename(path)
+
+        try:
+            raw = torch.load(path, weights_only=False, map_location="cpu")
+        except Exception as e:
+            print(f"Skip load {path} since error {e}")
+            return ValueWithMeta(
+                value=None, meta={**meta_from_filename, "filename": path.name}
+            )
+
+        value, meta_from_embedded = _unwrap_dict_format(raw)
+        return ValueWithMeta(
+            value=value,
+            meta={**meta_from_filename, **meta_from_embedded, "filename": path.name},
+        )
+
+
+def _unwrap_dict_format(obj: Any) -> Tuple[Any, Dict[str, Any]]:
+    if isinstance(obj, dict) and "value" in obj:
+        meta = obj.get("meta", {})
+        assert isinstance(meta, dict), f"Expected meta to be dict, got {type(meta)}"
+        return obj["value"], meta
+    return obj, {}
+

 class DumpLoader:
    def __init__(self):
@@ -50,10 +99,7 @@ def read_meta(directory):
    rows = []
    for p in directory.glob("*.pt"):
        try:
-            full_kwargs = {}
-            for kv in p.stem.split("___"):
-                k, v = kv.split("=")
-                full_kwargs[k] = v
+            full_kwargs = parse_meta_from_filename(p)
            rows.append(
                {
                    "filename": str(p.name),
@@ -83,26 +129,27 @@ def _add_duplicate_index(df: pl.DataFrame) -> pl.DataFrame:
    return df


-def find_row(df, conditions: Dict[str, Any]):
-    df_sub = df.filter(
-        functools.reduce(
-            lambda a, b: a & b,
-            [
-                (
-                    pl.col(col)
-                    == _cast_to_polars_dtype(conditions[col], df.schema[col])
-                    if conditions[col] is not None
-                    else pl.col(col).is_null()
-                )
-                for col in conditions.keys()
-                if col in df.columns
-            ],
+def filter_rows(df: pl.DataFrame, conditions: Dict[str, Any]) -> list[dict]:
+    filter_exprs = [
+        (
+            pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
+            if conditions[col] is not None
+            else pl.col(col).is_null()
        )
-    )
-    if len(df_sub) > 1:
-        print(f"find_row find ambiguous results: {df_sub=}")
+        for col in conditions
+        if col in df.columns
+    ]
+    if not filter_exprs:
+        return []
+    return df.filter(functools.reduce(lambda a, b: a & b, filter_exprs)).to_dicts()
+
+
+def find_row(df: pl.DataFrame, conditions: Dict[str, Any]):
+    rows = filter_rows(df, conditions)
+    if len(rows) > 1:
+        print(f"find_row find ambiguous results: {rows=}")
        return None
-    return df_sub.to_dicts()[0] if len(df_sub) > 0 else None
+    return rows[0] if rows else None


 def _cast_to_polars_dtype(value, target_dtype):
--- a/python/sglang/srt/debug_utils/dumper.py
+++ b/python/sglang/srt/debug_utils/dumper.py
@@ -223,7 +223,24 @@ class _Dumper:
        self._state.step += 1
        print(f"[Dumper] [{time.time()}] step={self._state.step}")

-    def dump(self, name: str, value, save: bool = True, **kwargs) -> None:
+    def dump(
+        self,
+        name: str,
+        value,
+        save: bool = True,
+        dims: Optional[str] = None,
+        dims_grad: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        value_meta: dict = {}
+        grad_meta: dict = {}
+        if dims is not None:
+            value_meta["dims"] = dims
+            grad_meta["dims"] = dims
+        if dims_grad is not None:
+            value_meta["dims_grad"] = dims_grad
+            grad_meta["dims"] = dims_grad
+
        self._dump_inner(
            name=name,
            value=value,
@@ -234,6 +251,8 @@ class _Dumper:
            enable_future_grad=self._config.enable_grad,
            value_tag="Dumper.Value",
            grad_tag="Dumper.Grad",
+            value_meta_only_fields=value_meta,
+            grad_meta_only_fields=grad_meta,
        )

    def dump_model(
@@ -336,6 +355,8 @@ class _Dumper:
        enable_future_grad: bool,
        value_tag: str,
        grad_tag: str,
+        value_meta_only_fields: Optional[dict] = None,
+        grad_meta_only_fields: Optional[dict] = None,
    ) -> None:
        self._http_manager  # noqa: B018

@@ -359,6 +380,7 @@ class _Dumper:
                tags=tags,
                value=value,
                save=save,
+                meta_only_fields=value_meta_only_fields or {},
            )

        if (
@@ -371,6 +393,7 @@ class _Dumper:
                tags={**tags, "name": f"grad__{name}"},
                value=g,
                save=save,
+                meta_only_fields=grad_meta_only_fields or {},
            )

        if enable_future_grad:
@@ -379,6 +402,7 @@ class _Dumper:
                tensor=value,
                extra_kwargs=extra_kwargs,
                save=save,
+                meta_only_fields=grad_meta_only_fields or {},
            )

    def _register_dump_grad_hook(
@@ -388,6 +412,7 @@ class _Dumper:
        tensor,
        extra_kwargs: dict,
        save: bool,
+        meta_only_fields: Optional[dict] = None,
    ) -> None:
        if not isinstance(tensor, torch.Tensor):
            return
@@ -396,6 +421,7 @@ class _Dumper:

        captured_step = self._state.step
        captured_tags = dict(name=f"grad__{name}", **deepcopy(extra_kwargs))
+        captured_meta_only = meta_only_fields or {}

        def grad_hook(grad: torch.Tensor) -> None:
            self._dump_single(
@@ -404,6 +430,7 @@ class _Dumper:
                value=grad,
                save=save,
                step=captured_step,
+                meta_only_fields=captured_meta_only,
            )

        tensor.register_hook(grad_hook)
@@ -416,6 +443,7 @@ class _Dumper:
        value,
        save: bool,
        step: Optional[int] = None,
+        meta_only_fields: Optional[dict] = None,
    ) -> None:
        self._ensure_exp_name()
        self._state.dump_index += 1
@@ -445,7 +473,11 @@ class _Dumper:
        if save and (self._config.enable_output_file or capturing):
            output_data = {
                "value": value,
-                "meta": dict(**full_kwargs, **self._static_meta),
+                "meta": dict(
+                    **full_kwargs,
+                    **self._static_meta,
+                    **(meta_only_fields or {}),
+                ),
            }

            if capturing:
--- a/test/registered/debug_utils/comparator/test_dims.py
+++ b/test/registered/debug_utils/comparator/test_dims.py
@@ -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__]))
--- a/test/registered/debug_utils/comparator/test_utils.py
+++ b/test/registered/debug_utils/comparator/test_utils.py
@@ -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__]))
--- a/test/registered/debug_utils/test_dump_loader.py
+++ b/test/registered/debug_utils/test_dump_loader.py
@@ -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__]))
--- a/test/registered/debug_utils/test_dumper.py
+++ b/test/registered/debug_utils/test_dumper.py
@@ -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__]))