From 036a784a860404337910a8b38070f3a06d5999af Mon Sep 17 00:00:00 2001 From: aasim syed Date: Sun, 27 Sep 2026 06:27:34 +0530 Subject: [PATCH 1/2] Reuse existing Puzzletron subblock stats when resuming a run Restarting a Puzzletron run failed at step 5 because calculate_subblock_stats_for_puzzle_dir raised when subblock_stats.json already existed. The file is written only after all stats are computed, so reuse it instead, matching how activation scoring and one-block scoring already resume. The check now runs before loading the teacher. Fixes the restart failure reported in #1668. Signed-off-by: aasim syed --- .../subblock_stats/calc_subblock_stats.py | 12 +++--- .../puzzletron/test_calc_subblock_stats.py | 37 +++++++++++++++++++ 2 files changed, 43 insertions(+), 6 deletions(-) create mode 100644 tests/unit/torch/puzzletron/test_calc_subblock_stats.py diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 1d04cc01add..23f3907538a 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/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 written only 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) diff --git a/tests/unit/torch/puzzletron/test_calc_subblock_stats.py b/tests/unit/torch/puzzletron/test_calc_subblock_stats.py new file mode 100644 index 00000000000..b6dfdba15ff --- /dev/null +++ b/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}}]' From 36b9f5134a7179a804b05858e582e8ccc94ebd9d Mon Sep 17 00:00:00 2001 From: aasim syed Date: Sun, 27 Sep 2026 07:28:47 +0530 Subject: [PATCH 2/2] Publish Puzzletron subblock stats atomically Write subblock_stats.json to a temporary file and rename it into place, so an interrupted write cannot leave a partial file that a resumed run would then reuse. Signed-off-by: aasim syed --- .../torch/puzzletron/subblock_stats/calc_subblock_stats.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 23f3907538a..062acbcc90a 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -280,7 +280,7 @@ 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 written only once all stats are computed, so reuse it when resuming a run. + # 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 @@ -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)