Handle recompute and verify closeness in dumper (#19564)

This commit is contained in:
fzyzcjy
2026-02-28 18:07:44 +08:00
committed by GitHub
parent 63a4778542
commit 40facdb28c
11 changed files with 557 additions and 13 deletions

View File

@@ -24,12 +24,18 @@ def normalize_parallel_info(meta: dict) -> dict[ParallelAxis, AxisInfo]:
info = value
if info is None:
return {}
info = {}
result: dict[ParallelAxis, AxisInfo] = {}
for axis in ParallelAxis:
axis_rank = info.get(f"{axis.value}_rank")
axis_size = info.get(f"{axis.value}_size")
# Recompute pseudo-axis lives at top-level meta, not inside parallel_info
if axis_rank is None:
axis_rank = meta.get(f"{axis.value}_rank")
axis_size = meta.get(f"{axis.value}_size")
if axis_rank is not None and axis_size is not None and axis_size > 1:
result[axis] = AxisInfo(
axis_rank=axis_rank,

View File

@@ -20,6 +20,7 @@ class ParallelAxis(Enum):
CP = "cp"
EP = "ep"
SP = "sp"
RECOMPUTE_PSEUDO = "recompute_pseudo"
class Ordering(Enum):

View File

@@ -124,7 +124,7 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
def _compute_skip_keys(args, *, has_token_aligner_plan: bool):
skip_keys: set[str] = {"dump_index", "filename"}
if args.grouping == "logical":
skip_keys |= {"rank"}
skip_keys |= {"rank", "recompute_status"}
if has_token_aligner_plan:
skip_keys |= {"step"}
return skip_keys

View File

@@ -2,12 +2,12 @@ import functools
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
from typing import Any, Callable, Dict, Optional, Tuple
import polars as pl
import torch
_TYPED_FIELDS: list[tuple[str, type]] = [("rank", int)]
LOAD_FAILED: object = object()
LOAD_FAILED: object = object()
@@ -19,9 +19,9 @@ def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
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])
for field_name, converter in _TYPED_FIELDS:
if field_name in result:
result[field_name] = converter(result[field_name])
return result
@@ -177,4 +177,9 @@ def read_tokenizer_path(directory: Path) -> Optional[str]:
return None
_TYPED_FIELDS: list[tuple[str, Callable[[str], Any]]] = [
("rank", int),
]
dump_loader = DumpLoader()

View File

@@ -1,3 +1,4 @@
import enum
import functools
import json
import os
@@ -415,7 +416,14 @@ class _Dumper:
if not self._config.enable:
return
tags = dict(name=name, **extra_kwargs, **self._state.global_ctx)
recompute_status = _detect_recompute_status()
tags = dict(
name=name,
recompute_status=recompute_status.value,
**extra_kwargs,
**self._state.global_ctx,
)
if (f := self._config.filter) is not None and re.search(
f, _format_tags(tags)
) is None:
@@ -424,6 +432,7 @@ class _Dumper:
if not (enable_value or enable_curr_grad or enable_future_grad):
return
recompute_meta = recompute_status.to_pseudo_parallel_meta()
value = _materialize_value(value)
if enable_value:
@@ -432,7 +441,7 @@ class _Dumper:
tags=tags,
value=value,
save=save,
meta_only_fields=value_meta_only_fields or {},
meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta},
)
if (
@@ -445,7 +454,7 @@ class _Dumper:
tags={**tags, "name": f"grad__{name}"},
value=g,
save=save,
meta_only_fields=grad_meta_only_fields or {},
meta_only_fields={**(grad_meta_only_fields or {}), **recompute_meta},
)
if enable_future_grad:
@@ -472,7 +481,10 @@ class _Dumper:
return
captured_step = self._state.step
captured_tags = dict(name=f"grad__{name}", **deepcopy(extra_kwargs))
captured_tags = dict(
name=f"grad__{name}",
**deepcopy(extra_kwargs),
)
captured_meta_only = meta_only_fields or {}
def grad_hook(grad: torch.Tensor) -> None:
@@ -1142,6 +1154,20 @@ def _get_local_ip_by_remote() -> Optional[str]:
# -------------------------------------- framework plugins ------------------------------------------
class _RecomputeStatus(enum.Enum):
DISABLED = "disabled"
ORIGINAL = "original" # inside checkpoint, original forward
RECOMPUTE = "recompute" # inside checkpoint, recompute forward
def to_pseudo_parallel_meta(self) -> dict[str, Any]:
if self == _RecomputeStatus.DISABLED:
return {}
return {
"recompute_pseudo_rank": 1 if self == _RecomputeStatus.RECOMPUTE else 0,
"recompute_pseudo_size": 2,
}
class _FrameworkPlugin(ABC):
@property
@abstractmethod
@@ -1168,6 +1194,9 @@ class _FrameworkPlugin(ABC):
def get_tokenizer_path(self) -> Optional[str]:
return None
def detect_recompute_status(self) -> _RecomputeStatus:
return _RecomputeStatus.DISABLED
class _SGLangPlugin(_FrameworkPlugin):
_available = True
@@ -1353,10 +1382,32 @@ class _MegatronPlugin(_FrameworkPlugin):
{"input_ids", "position_ids", "cu_seqlens_q", "cu_seqlens_kv", "qkv_format"}
)
def detect_recompute_status(self) -> _RecomputeStatus:
if not self._available:
return _RecomputeStatus.DISABLED
try:
from megatron.core.tensor_parallel.random import is_checkpointing
if not is_checkpointing():
return _RecomputeStatus.DISABLED
if torch.is_grad_enabled():
return _RecomputeStatus.RECOMPUTE
return _RecomputeStatus.ORIGINAL
except (ImportError, AttributeError):
return _RecomputeStatus.DISABLED
_plugins: list[_FrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()]
def _detect_recompute_status() -> _RecomputeStatus:
for plugin in _plugins:
info = plugin.detect_recompute_status()
if info != _RecomputeStatus.DISABLED:
return info
return _RecomputeStatus.DISABLED
# -------------------------------------- singleton ------------------------------------------

View File

@@ -506,6 +506,138 @@ class TestVerifyReplicatedGroup:
assert len(warnings) == 1
assert warnings[0].differing_index == 1
def test_recompute_pseudo_mismatch_warns(self) -> None:
"""_verify_replicated_group produces warning for RECOMPUTE_PSEUDO axis mismatch."""
tensor_a = torch.ones(4)
tensor_b = torch.ones(4) + 0.1
with warning_sink.context() as warnings:
_verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.RECOMPUTE_PSEUDO,
group_index=0,
)
assert len(warnings) == 1
assert warnings[0].axis == "recompute_pseudo"
assert warnings[0].group_index == 0
assert warnings[0].differing_index == 1
assert warnings[0].baseline_index == 0
assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
class TestThdCpConcat:
def test_single_seq(self) -> None:
"""Single seq THD unshard: 2 ranks → per-seq concat."""
rank0 = torch.tensor([1, 2, 3]).refine_names("t")
rank1 = torch.tensor([4, 5, 6]).refine_names("t")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
expected = torch.tensor([1, 2, 3, 4, 5, 6])
assert torch.equal(result[0].rename(None), expected)
def test_multi_seq(self) -> None:
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
seq_a_r0 = torch.arange(0, 50)
seq_b_r0 = torch.arange(100, 132)
pad_r0 = torch.full((46,), -1)
rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
seq_a_r1 = torch.arange(50, 100)
seq_b_r1 = torch.arange(132, 164)
pad_r1 = torch.full((46,), -2)
rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
# seqB: r0(32) + r1(32) = 64 tokens
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
# pad: r0(46) + r1(46) = 92 tokens
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
def test_with_hidden_dim(self) -> None:
"""THD unshard with trailing hidden dim: shape [T, H]."""
torch.manual_seed(42)
hidden: int = 4
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
seq_a_r0 = torch.randn(3, hidden)
seq_b_r0 = torch.randn(2, hidden)
rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
seq_a_r1 = torch.randn(3, hidden)
seq_b_r1 = torch.randn(2, hidden)
rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert unsharded.shape == (10, hidden)
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
def test_with_leading_batch_dim(self) -> None:
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
torch.manual_seed(42)
batch: int = 2
hidden: int = 4
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
seq_a_r0 = torch.randn(batch, 3, hidden)
seq_b_r0 = torch.randn(batch, 2, hidden)
rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
seq_a_r1 = torch.randn(batch, 3, hidden)
seq_b_r1 = torch.randn(batch, 2, hidden)
rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert unsharded.shape == (batch, 10, hidden)
# seqA: r0(3) + r1(3) = 6 tokens per batch
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
# seqB: r0(2) + r1(2) = 4 tokens per batch
assert torch.equal(
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
)
class TestThdCpConcat:
def test_single_seq(self) -> None:

View File

@@ -81,6 +81,19 @@ class TestNormalizeParallelInfo:
}
assert normalize_parallel_info(meta) == {}
def test_recompute_pseudo_from_top_level_meta(self) -> None:
"""recompute_pseudo_rank/size at top-level meta is extracted alongside TP."""
meta = {
"recompute_pseudo_rank": 1,
"recompute_pseudo_size": 2,
"sglang_parallel_info": {"tp_rank": 0, "tp_size": 2},
}
result = normalize_parallel_info(meta)
assert result == {
ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2),
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
}
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

View File

@@ -421,6 +421,20 @@ class TestReplicatedAxes:
with pytest.raises(ValueError, match="missing parallel_info"):
compute_unsharder_plan(dim_specs, parallel_infos)
def test_recompute_pseudo_replicated(self) -> None:
"""RECOMPUTE_PSEUDO with no dim annotation → replicated → PickParams."""
dim_specs = parse_dims("h d")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
{ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.RECOMPUTE_PSEUDO
assert isinstance(plans[0].params, PickParams)
assert plans[0].groups == [[0, 1]]
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

View File

@@ -13,13 +13,14 @@ from sglang.srt.debug_utils.comparator.output_types import (
ConfigRecord,
GeneralWarning,
NonTensorRecord,
ReplicatedMismatchWarning,
SkipRecord,
SummaryRecord,
WarningRecord,
_OutputRecord,
parse_record_json,
)
from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper, _RecomputeStatus
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=30, suite="default", nightly=True)
@@ -881,6 +882,142 @@ class TestEntrypointGroupingLogical:
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
def test_recompute_pseudo_replicated_verification(self, tmp_path, capsys):
"""Recompute pseudo-axis with identical original/recompute tensors → passed."""
torch.manual_seed(42)
tensor = torch.randn(4, 8)
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
for side_dir in [baseline_dir, target_dir]:
_create_recompute_rank_dump(
side_dir,
rank=0,
name="hidden",
original_tensor=tensor,
recompute_tensor=tensor.clone(),
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
def test_recompute_pseudo_mismatch_warning(self, tmp_path, capsys):
"""Recompute pseudo-axis with differing original/recompute → ReplicatedMismatchWarning."""
torch.manual_seed(42)
tensor = torch.randn(4, 8)
mismatched_tensor = tensor + torch.randn(4, 8) * 10.0
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
for side_dir in [baseline_dir, target_dir]:
_create_recompute_rank_dump(
side_dir,
rank=0,
name="hidden",
original_tensor=tensor,
recompute_tensor=mismatched_tensor,
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
recompute_warnings = [
w
for w in comparisons[0].warnings
if isinstance(w, ReplicatedMismatchWarning) and w.axis == "recompute_pseudo"
]
assert len(recompute_warnings) > 0
class TestEntrypointAxisSwapper:
"""Test cross-framework dim reordering through the full entrypoint pipeline."""
def test_axis_swap_different_dim_order(self, tmp_path, capsys):
"""Baseline dims 'b h d' vs target dims 'b d h': axis swapper rearranges baseline to match."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
_create_rank_dump(
baseline_dir,
rank=0,
name="hidden",
tensor=full_tensor,
dims="b h d",
)
_create_rank_dump(
target_dir,
rank=0,
name="hidden",
tensor=full_tensor.permute(0, 2, 1).contiguous(),
dims="b d h",
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
assert comp.baseline.shape == [4, 16, 8]
assert comp.target.shape == [4, 16, 8]
def test_axis_swap_with_tp_unshard(self, tmp_path, capsys):
"""Baseline TP=2 with dims 'b h(tp) d' vs target TP=2 with dims 'b d h(tp)': unshard + axis swap."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
_create_tp_sharded_dumps(
baseline_dir,
full_tensor=full_tensor,
name="hidden",
tp_size=2,
shard_dim=1,
dims_str="b h(tp) d",
)
_create_tp_sharded_dumps(
target_dir,
full_tensor=full_tensor.permute(0, 2, 1).contiguous(),
name="hidden",
tp_size=2,
shard_dim=2,
dims_str="b d h(tp)",
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
class TestEntrypointAxisSwapper:
"""Test cross-framework dim reordering through the full entrypoint pipeline."""
@@ -1826,6 +1963,53 @@ def _create_tp_sharded_dumps(
return directory / _FIXED_EXP_NAME
def _create_recompute_rank_dump(
directory: Path,
*,
rank: int,
name: str,
original_tensor: torch.Tensor,
recompute_tensor: torch.Tensor,
dims: str = "h d",
) -> Path:
"""Create a dump with both original and recompute forward passes via monkeypatched dumper.
The dumper naturally produces recompute_pseudo_rank=0 for original and =1 for recompute,
plus recompute_pseudo_size=2.
"""
with pytest.MonkeyPatch.context() as mp:
mp.setattr(_dumper_module, "_get_rank", lambda: rank)
dumper = _Dumper(
config=DumperConfig(
enable=True,
dir=str(directory),
exp_name=_FIXED_EXP_NAME,
)
)
dumper.__dict__["_static_meta"] = {"world_rank": rank, "world_size": 1}
# dump original forward
mp.setattr(
_dumper_module,
"_detect_recompute_status",
lambda: _RecomputeStatus.ORIGINAL,
)
dumper.dump(name, original_tensor, dims=dims)
# dump recompute forward
mp.setattr(
_dumper_module,
"_detect_recompute_status",
lambda: _RecomputeStatus.RECOMPUTE,
)
dumper.dump(name, recompute_tensor, dims=dims)
dumper.step()
return directory / _FIXED_EXP_NAME
def _zigzag_split_seq(seq_natural: torch.Tensor, *, cp_size: int) -> list[torch.Tensor]:
"""Split a natural-order seq into per-rank zigzag segments."""
num_chunks: int = cp_size * 2

View File

@@ -5,10 +5,12 @@ import pytest
import torch
from sglang.srt.debug_utils.dump_loader import (
LOAD_FAILED,
ValueWithMeta,
_add_duplicate_index,
_cast_to_polars_dtype,
find_row,
parse_meta_from_filename,
read_meta,
)
from sglang.test.ci.ci_register import register_cpu_ci
@@ -92,9 +94,35 @@ class TestValueWithMeta:
path.write_text("not a valid pt file")
loaded = ValueWithMeta.load(path)
assert loaded.value is None
assert loaded.value is LOAD_FAILED
assert loaded.meta["name"] == "bad"
class TestRecomputeStatusParsing:
def test_parse_recompute_status_from_filename(self) -> None:
from pathlib import Path
meta_disabled = parse_meta_from_filename(
Path(
"step=0___rank=0___dump_index=1___name=x___recompute_status=disabled.pt"
)
)
assert meta_disabled["recompute_status"] == "disabled"
meta_recompute = parse_meta_from_filename(
Path(
"step=0___rank=0___dump_index=1___name=x___recompute_status=recompute.pt"
)
)
assert meta_recompute["recompute_status"] == "recompute"
meta_original = parse_meta_from_filename(
Path(
"step=0___rank=0___dump_index=1___name=x___recompute_status=original.pt"
)
)
assert meta_original["recompute_status"] == "original"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

View File

@@ -16,6 +16,7 @@ from sglang.srt.debug_utils.dumper import (
DumperConfig,
_collective_with_timeout,
_deepcopy_or_clone,
_detect_recompute_status,
_Dumper,
_format_tags,
_get_default_exp_name,
@@ -23,6 +24,7 @@ from sglang.srt.debug_utils.dumper import (
_materialize_value,
_MegatronPlugin,
_obj_to_dict,
_RecomputeStatus,
_register_forward_hook_or_replace_fn,
_SGLangPlugin,
_torch_save,
@@ -2400,5 +2402,113 @@ class TestCtxDecorator:
d.ctx()
class TestRecomputeStatus:
def test_disabled_by_default(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path)
tensor = torch.randn(3, 3)
d.dump("test_tensor", tensor)
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["recompute_status=disabled"])
def test_recompute_status_in_embedded_meta(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path)
tensor = torch.randn(3, 3)
d.dump("test_tensor", tensor)
path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
raw = _load_dump(path)
assert raw["meta"]["recompute_status"] == "disabled"
def test_recompute_status_recompute(self, tmp_path: Path, monkeypatch) -> None:
import sglang.srt.debug_utils.dumper as dumper_mod
monkeypatch.setattr(
dumper_mod, "_detect_recompute_status", lambda: _RecomputeStatus.RECOMPUTE
)
d = _make_test_dumper(tmp_path)
tensor = torch.randn(3, 3)
d.dump("test_tensor", tensor)
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["recompute_status=recompute"])
path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
raw = _load_dump(path)
assert raw["meta"]["recompute_status"] == "recompute"
assert raw["meta"]["recompute_pseudo_rank"] == 1
assert raw["meta"]["recompute_pseudo_size"] == 2
def test_recompute_status_original(self, tmp_path: Path, monkeypatch) -> None:
import sglang.srt.debug_utils.dumper as dumper_mod
monkeypatch.setattr(
dumper_mod,
"_detect_recompute_status",
lambda: _RecomputeStatus.ORIGINAL,
)
d = _make_test_dumper(tmp_path)
tensor = torch.randn(3, 3)
d.dump("test_tensor", tensor)
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["recompute_status=original"])
path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
raw = _load_dump(path)
assert raw["meta"]["recompute_status"] == "original"
assert raw["meta"]["recompute_pseudo_rank"] == 0
assert raw["meta"]["recompute_pseudo_size"] == 2
def test_disabled_no_recompute_pseudo_fields(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path)
tensor = torch.randn(3, 3)
d.dump("test_tensor", tensor)
path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
raw = _load_dump(path)
assert "recompute_pseudo_rank" not in raw["meta"]
assert "recompute_pseudo_size" not in raw["meta"]
def test_grad_hook_has_no_recompute_status(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path, enable_grad=True)
x = torch.randn(3, 3, requires_grad=True)
y = (x * 2).sum()
d.dump("test_tensor", x)
y.backward()
grad_files = [f for f in _get_filenames(tmp_path) if "grad__test_tensor" in f]
assert len(grad_files) == 1
assert "recompute_status" not in grad_files[0]
def test_non_intrusive_hooks_have_recompute_status(self, tmp_path: Path) -> None:
class Simple(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.linear(x)
model = Simple()
d = _make_test_dumper(tmp_path, non_intrusive_mode="all")
d.register_non_intrusive_dumper(model)
with d.capture_output() as captured:
model(torch.randn(2, 4))
for key, data in captured.items():
assert (
"recompute_status" in data["meta"]
), f"missing recompute_status in {key}"
assert data["meta"]["recompute_status"] == "disabled"
def test_detect_recompute_status_default(self) -> None:
assert _detect_recompute_status() == _RecomputeStatus.DISABLED
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))