diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 1d04cc01add..062acbcc90a 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 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) 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}}]'