pytorch/pytorch

[Distributed Tools] MemoryTracker load does not restore operation count

オープン

#191,397 opened on 2026/07/28

 (1 件のコメント) (0 件のリアクション) (0 人の担当者)Python (28,884 件のフォーク)batch import
bot-triagedgood first issuemodule: distributed_tooloncall: distributed infratriaged

Repository metrics

Stars
 (102,440 個のスター)
PR merge metrics
 (平均マージ 1d 10h) (30d で 42 merged PRs)

説明

Problem

On current main (5d242823c6de2bbc21d636b2995a38340bd2a16c), MemoryTracker.save_stats() persists the memory traces but not _op_index, and load() restores the traces without reconstructing that count.

Relevant code: save_stats() and load(). summary() bounds its iteration using _op_index.

Reproduction

Save a tracker containing at least two operation entries, load it into a fresh MemoryTracker, and call summary().

The loaded trace entries are present, but _op_index remains 0, so the summary prints no operation deltas. In a minimal two-entry reproduction, the source tracker reports the operation while the loaded tracker does not.

Expected behavior

A save/load round trip should preserve the operation count and produce the same summary/plot data as the original tracker. This can be done by persisting _op_index or reconstructing it consistently from the loaded traces, with compatibility for existing stats files.

Impact

Offline analysis of saved memory traces silently omits the operator summary, one of the primary outputs of MemoryTracker.

Suggested regression test

Create a tracker with multiple indexed entries, save and reload it, then assert the restored operation count and captured summary() output match the original.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @pragupta @msaroufim @dcci @aditvenk @weifengpy @kapilsh

コントリビューターガイド