Remove the legacy torchtitan.components.checkpoint_utils import path

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,

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.

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 (

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