Skip to content

trunk/62bda070d37c0c292c066a9d0885bdcc851cee8f

@ccyyy1023 ccyyy1023 tagged this 22 Aug 04:06
Fixes #191397

## Summary

- reconstruct `MemoryTracker._op_index` from the loaded allocation records
- keep the existing pickle format unchanged so previously saved stats files remain compatible
- add a CPU-compatible round-trip test that loads stats into a fresh tracker and compares the restored operation count and summary output

## Root cause

`save_stats()` persists the recorded memory dictionaries, but `load()` did not restore the operation count used by `summary()`. A fresh tracker therefore retained `_op_index == 0` even though its loaded records were present.

## Testing

Using Python 3.12 and `torch==2.14.0.dev20260728+cpu`, with the modified module copied into the nightly wheel environment:

- `python -m pytest -q test/distributed/_tools/test_memory_tracker.py`
  - `1 passed, 1 skipped`
  - the existing accelerator test was skipped because the environment used a CPU-only wheel

This was a wheel-based validation of the Python changes rather than a full source build; repository CI will provide source-integration coverage.

## AI assistance disclosure

This change was developed with assistance from OpenAI Codex. I reviewed the code and test changes and ran the test reported above.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/191453
Approved by: https://github.com/d4l3k
Assets 2
Loading