[Puzzletron] Reuse existing subblock stats when resuming a run

modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py

@@ -279,6 +279,12 @@ def calculate_subblock_stats_for_puzzle_dir(
         ]
 
     master_puzzle_dir = Path(master_puzzle_dir)
+    subblock_stats_file = master_puzzle_dir / subblock_stats_filename
+    # The file is published atomically once all stats are computed, so reuse it when resuming a run.
+    if subblock_stats_file.exists() and not merge_with_existing_stats:
+        mprint(f"Subblock stats file {subblock_stats_file} already exists, skipping...")
+        return
+
     teacher_dir = (
         Path(teacher_dir) if teacher_dir is not None else master_puzzle_dir / "ckpts" / "teacher"
     )
@@ -288,12 +294,6 @@ def calculate_subblock_stats_for_puzzle_dir(
     lm_config = descriptor.get_language_model_config(model_config)
     subblock_configs = _load_subblock_configs(master_puzzle_dir, ffn_hidden_sizes)
 
-    subblock_stats_file = master_puzzle_dir / subblock_stats_filename
-    if subblock_stats_file.exists() and not merge_with_existing_stats:
-        raise ValueError(
-            f"Subblock stats file {subblock_stats_file} already exists and `merge_with_existing_stats` was set to False."
-        )
-
     if subblock_stats_file.exists():
         with open(subblock_stats_file) as f:
             subblock_stats = json.load(f)
@@ -361,7 +361,9 @@ def calculate_subblock_stats_for_puzzle_dir(
 
         subblock_stats.append(curr_subblock_stats)
 
-    json_dump(subblock_stats, subblock_stats_file)
+    tmp_stats_file = subblock_stats_file.with_name(subblock_stats_file.name + ".tmp")
+    json_dump(subblock_stats, tmp_stats_file)
+    tmp_stats_file.replace(subblock_stats_file)
 
     mprint(subblock_stats_file)
 

tests/unit/torch/puzzletron/test_calc_subblock_stats.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from omegaconf import OmegaConf
+
+from modelopt.torch.puzzletron.subblock_stats.calc_subblock_stats import (
+    calculate_subblock_stats_for_puzzle_dir,
+)
+
+
+def test_existing_subblock_stats_are_reused_on_resume(tmp_path):
+    stats_file = tmp_path / "subblock_stats.json"
+    stats_file.write_text('[{"args": {"batch_size": 1}}]')
+    teacher_dir = tmp_path / "ckpts" / "teacher"  # absent: resuming must not reload the teacher
+
+    calculate_subblock_stats_for_puzzle_dir(
+        OmegaConf.create({}),
+        master_puzzle_dir=tmp_path,
+        teacher_dir=teacher_dir,
+        descriptor=None,
+        model_hidden_sizes=OmegaConf.create([]),
+        ffn_hidden_sizes=OmegaConf.create([]),
+    )
+
+    assert stats_file.read_text() == '[{"args": {"batch_size": 1}}]'