From 03309362675873389798371e626a20a113b3bda3 Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Wed, 19 Aug 2026 14:15:51 -0700 Subject: [PATCH] Remove the legacy torchtitan.components.checkpoint_utils import path Summary: `components/checkpoint_utils.py` became a re-export shim when the optimizer and checkpointer components were grouped into packages. Unlike the lr_scheduler shim removed in #4172, this one forwarded to two different destinations at once, which is what made it worth deleting rather than keeping: reading an import of `checkpoint_utils` told you nothing about whether the symbol was optimizer plumbing or checkpointer plumbing. Route each of the four importers to the module that actually defines the symbol. `canonical_fqn` lives in `checkpointer/utils.py`; `init_optim_state`, `get_flat_optim_state_dict`, and `load_flat_optim_state_dict` live in `optimizer/utils.py`. This also settles the naming objection fegin raised on #4140, that `canonical_fqn` does not belong under an optimizer-shaped name -- its importer in the rl trainer now names the checkpointer package directly. The three state-dict helpers are imported from `optimizer.utils` rather than re-exported through `optimizer/__init__.py`. They are low-level DCP plumbing with two callers between them, not part of the package's public surface, which stays `OptimizersContainer`, `LRSchedulersContainer`, `ParamGroupConfig`, and `default_adamw`. This is an import-path change only; no runtime behavior changes. Test Plan: `pytest tests/unit_tests/test_state_dict_keys.py tests/unit_tests/test_checkpoint.py tests/unit_tests/test_lr_scheduler.py tests/unit_tests/test_optimizer_param_groups.py tests/unit_tests/test_torch_checkpointing.py`: 77 passed, 4 subtests passed. `test_legacy_checkpoint_utils_imports`, which asserted the shim's symbols were identical to the submodule's, is dropped -- it cannot outlive the shim. `test_legacy_checkpoint_utils_can_be_imported_first` is kept but retargeted, as `test_state_dict_helpers_can_be_imported_first`. It guards a real property rather than the shim: the `optimizer` and `checkpointer` package `__init__` files import from each other, so importing either leaf `utils` module first in a fresh interpreter must not close an import cycle. It now subtests both leaf modules instead of the single shim. Also verified: - No `checkpoint_utils` references remain anywhere in the repo, across all file types, not only Python. - `import torchtitan.components.checkpoint_utils` now raises `ModuleNotFoundError`. - `ufmt` and `flake8 --config=.flake8` clean on the four changed files. --- tests/unit_tests/flex_shard/test_dist_muon.py | 2 +- tests/unit_tests/test_state_dict_keys.py | 47 ++++++++----------- torchtitan/components/checkpoint_utils.py | 21 --------- torchtitan/experiments/rl/actors/trainer.py | 2 +- torchtitan/experiments/torchft/optimizer.py | 2 +- 5 files changed, 23 insertions(+), 51 deletions(-) delete mode 100644 torchtitan/components/checkpoint_utils.py diff --git a/tests/unit_tests/flex_shard/test_dist_muon.py b/tests/unit_tests/flex_shard/test_dist_muon.py index 48b353397f..2c393f7ab0 100644 --- a/tests/unit_tests/flex_shard/test_dist_muon.py +++ b/tests/unit_tests/flex_shard/test_dist_muon.py @@ -15,7 +15,7 @@ DTensorTestBase, with_comms, ) -from torchtitan.components.checkpoint_utils import ( +from torchtitan.components.optimizer.utils import ( get_flat_optim_state_dict, init_optim_state, load_flat_optim_state_dict, diff --git a/tests/unit_tests/test_state_dict_keys.py b/tests/unit_tests/test_state_dict_keys.py index b77781c1c9..fc3697a978 100644 --- a/tests/unit_tests/test_state_dict_keys.py +++ b/tests/unit_tests/test_state_dict_keys.py @@ -31,14 +31,14 @@ from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( checkpoint_wrapper as ptd_checkpoint_wrapper, ) -from torchtitan.components.checkpoint_utils import ( + +from torchtitan.components.checkpointer import ModelWrapper +from torchtitan.components.optimizer import OptimizersContainer, ParamGroupConfig +from torchtitan.components.optimizer.utils import ( get_flat_optim_state_dict, init_optim_state, load_flat_optim_state_dict, ) - -from torchtitan.components.checkpointer import ModelWrapper -from torchtitan.components.optimizer import OptimizersContainer, ParamGroupConfig from torchtitan.models.llama3 import llama3_configs from torchtitan.models.llama3.model import Llama3Model @@ -102,30 +102,23 @@ def _debugmodel_optimizer_config() -> OptimizersContainer.Config: class TestStateDictKeys(unittest.TestCase): - def test_legacy_checkpoint_utils_can_be_imported_first(self): - subprocess.run( - [ - sys.executable, - "-c", - "from torchtitan.components.checkpoint_utils import " - "canonical_fqn, get_flat_optim_state_dict, init_optim_state, " + def test_state_dict_helpers_can_be_imported_first(self): + # These modules sit below the `optimizer` and `checkpointer` packages, + # whose __init__ files import from each other. Importing either leaf + # module first, in a fresh interpreter, must not close an import cycle. + for module, names in ( + ( + "torchtitan.components.optimizer.utils", + "get_flat_optim_state_dict, init_optim_state, " "load_flat_optim_state_dict", - ], - check=True, - ) - - def test_legacy_checkpoint_utils_imports(self): - from torchtitan.components.optimizer import utils as optimizer_utils - - self.assertIs(optimizer_utils.init_optim_state, init_optim_state) - self.assertIs( - optimizer_utils.get_flat_optim_state_dict, - get_flat_optim_state_dict, - ) - self.assertIs( - optimizer_utils.load_flat_optim_state_dict, - load_flat_optim_state_dict, - ) + ), + ("torchtitan.components.checkpointer.utils", "canonical_fqn"), + ): + with self.subTest(module=module): + subprocess.run( + [sys.executable, "-c", f"from {module} import {names}"], + check=True, + ) def setUp(self) -> None: # Ground-truth canonical keys come from the unwrapped model. diff --git a/torchtitan/components/checkpoint_utils.py b/torchtitan/components/checkpoint_utils.py deleted file mode 100644 index c8cf00d5b3..0000000000 --- a/torchtitan/components/checkpoint_utils.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Compatibility imports for relocated checkpoint and optimizer utilities.""" - -from .checkpointer.utils import canonical_fqn -from .optimizer.utils import ( - get_flat_optim_state_dict, - init_optim_state, - load_flat_optim_state_dict, -) - -__all__ = [ - "canonical_fqn", - "init_optim_state", - "get_flat_optim_state_dict", - "load_flat_optim_state_dict", -] diff --git a/torchtitan/experiments/rl/actors/trainer.py b/torchtitan/experiments/rl/actors/trainer.py index c7accfdb91..432f0be263 100644 --- a/torchtitan/experiments/rl/actors/trainer.py +++ b/torchtitan/experiments/rl/actors/trainer.py @@ -12,8 +12,8 @@ import torch import torchstore as ts from monarch.actor import Actor, concurrent_endpoint, current_rank -from torchtitan.components.checkpoint_utils import canonical_fqn from torchtitan.components.checkpointer import CheckpointManager +from torchtitan.components.checkpointer.utils import canonical_fqn from torchtitan.components.loss import BaseLoss, ChunkedLossWrapper from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer from torchtitan.config import ( diff --git a/torchtitan/experiments/torchft/optimizer.py b/torchtitan/experiments/torchft/optimizer.py index 28b771bb78..a80ddf7a99 100644 --- a/torchtitan/experiments/torchft/optimizer.py +++ b/torchtitan/experiments/torchft/optimizer.py @@ -10,8 +10,8 @@ import torch.nn as nn -from torchtitan.components.checkpoint_utils import init_optim_state from torchtitan.components.optimizer import OptimizersContainer +from torchtitan.components.optimizer.utils import init_optim_state if TYPE_CHECKING: from torchtitan.experiments.torchft.manager import TorchFTManager