From 02ca107b2c072d6ef6dde6e76b5ac1d33306ad76 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 25 Feb 2026 09:41:48 +0800 Subject: [PATCH] Support dims annotation and enhance dump loader in dumper (#19276) --- .../sglang/srt/debug_utils/comparator/dims.py | 80 +++++++++++++ .../srt/debug_utils/comparator/entrypoint.py | 16 ++- .../srt/debug_utils/comparator/utils.py | 17 --- python/sglang/srt/debug_utils/dump_loader.py | 93 +++++++++++---- python/sglang/srt/debug_utils/dumper.py | 36 +++++- .../debug_utils/comparator/test_dims.py | 101 +++++++++++++++++ .../debug_utils/comparator/test_utils.py | 25 ---- .../debug_utils/test_dump_loader.py | 107 ++++++++++++------ test/registered/debug_utils/test_dumper.py | 69 +++++++++++ 9 files changed, 437 insertions(+), 107 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/dims.py create mode 100644 test/registered/debug_utils/comparator/test_dims.py diff --git a/python/sglang/srt/debug_utils/comparator/dims.py b/python/sglang/srt/debug_utils/comparator/dims.py new file mode 100644 index 000000000..55f6c3b78 --- /dev/null +++ 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[a-zA-Z_]\w*)(?:\((?P[^)]+)\))?$") + +_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 diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 196439277..00add6ddb 100644 --- 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() diff --git a/python/sglang/srt/debug_utils/comparator/utils.py b/python/sglang/srt/debug_utils/comparator/utils.py index 63a310fa7..67b9b219b 100644 --- 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() diff --git a/python/sglang/srt/debug_utils/dump_loader.py b/python/sglang/srt/debug_utils/dump_loader.py index f334ca74e..989586780 100644 --- 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): diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index da718eef9..0a081cd05 100644 --- 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: diff --git a/test/registered/debug_utils/comparator/test_dims.py b/test/registered/debug_utils/comparator/test_dims.py new file mode 100644 index 000000000..8e5e1a322 --- /dev/null +++ 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__])) diff --git a/test/registered/debug_utils/comparator/test_utils.py b/test/registered/debug_utils/comparator/test_utils.py index 824d645d0..9d36c116a 100644 --- 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__])) diff --git a/test/registered/debug_utils/test_dump_loader.py b/test/registered/debug_utils/test_dump_loader.py index 9c9d768a2..6187561f8 100644 --- 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__])) diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 47e1b5bbf..6b21f2e33 100644 --- 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__]))