Remove the legacy torchtitan.components.checkpoint import path

scripts/checkpoint_conversion/convert_from_hf.py

@@ -11,7 +11,7 @@
 import torch
 import torch.distributed.checkpoint as dcp
 from torch.distributed.checkpoint import HuggingFaceStorageReader
-from torchtitan.components.checkpoint import ModelWrapper
+from torchtitan.components.checkpointer import ModelWrapper
 
 
 @torch.inference_mode()

scripts/checkpoint_conversion/convert_to_hf.py

@@ -11,7 +11,7 @@
 import torch
 import torch.distributed.checkpoint as dcp
 from torch.distributed.checkpoint import HuggingFaceStorageWriter
-from torchtitan.components.checkpoint import ModelWrapper
+from torchtitan.components.checkpointer import ModelWrapper
 from torchtitan.config import TORCH_DTYPE_MAP
 
 

scripts/checkpoint_conversion/numerical_tests_example.py

@@ -8,7 +8,7 @@
 
 import torch.distributed.checkpoint as dcp
 import torch.nn.functional as F
-from torchtitan.components.checkpoint import ModelWrapper
+from torchtitan.components.checkpointer import ModelWrapper
 from torchtitan.config import ConfigManager
 from torchtitan.tools.logging import logger
 

scripts/checkpoint_conversion/numerical_tests_kimi.py

@@ -39,7 +39,7 @@
 import torch.nn.functional as F
 from PIL import Image
 
-from torchtitan.components.checkpoint import ModelWrapper
+from torchtitan.components.checkpointer import ModelWrapper
 from torchtitan.hf_datasets.multimodal.utils.image import (
     process_image,
     resize_to_patch_budget,

scripts/checkpoint_conversion/numerical_tests_qwen3_5.py

@@ -33,7 +33,7 @@
 
 torch._dynamo.config.disable = True
 
-from torchtitan.components.checkpoint import ModelWrapper
+from torchtitan.components.checkpointer import ModelWrapper
 from torchtitan.hf_datasets.multimodal.mm_collator import MultiModalCollator
 from torchtitan.hf_datasets.multimodal.utils.image import (
     process_image,

tests/unit_tests/test_checkpoint.py

@@ -165,7 +165,7 @@ def test_trainer_uses_checkpoint_interface_with_concrete_default(self):
         self.assertFalse(checkpoint.enable)
 
     def test_legacy_import_path(self):
-        from torchtitan.components.checkpoint import (
+        from torchtitan.components.checkpointer import (
             CheckpointManager as LegacyCheckpointManager,
             ModelWrapper as LegacyModelWrapper,
         )

tests/unit_tests/test_state_dict_keys.py

@@ -31,13 +31,13 @@
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     checkpoint_wrapper as ptd_checkpoint_wrapper,
 )
-
-from torchtitan.components.checkpoint import ModelWrapper
 from torchtitan.components.checkpoint_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

torchtitan/experiments/forge/engine.py

@@ -12,7 +12,7 @@
 import torch
 from torch.distributed.elastic.multiprocessing.errors import record
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import LossFunction
 from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer
 from torchtitan.config import Configurable, TORCH_DTYPE_MAP

torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py

@@ -25,7 +25,7 @@
 from tests.utils import hash_gradient, hash_model
 from torch.nn.attention.flex_attention import flex_attention
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import CrossEntropyLoss
 from torchtitan.components.tokenizer import HuggingFaceTokenizer
 from torchtitan.config import DebugConfig, ParallelismConfig, TrainingConfig

torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py

@@ -17,7 +17,7 @@
 from torch.nn.attention.flex_attention import flex_attention
 from torch.testing._internal.common_fsdp import FSDPTest
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.config import DebugConfig, ParallelismConfig, TrainingConfig
 from torchtitan.distributed import ParallelDims
 from torchtitan.experiments.graph_trainer.common_utils import (

torchtitan/experiments/rl/__init__.py

@@ -8,7 +8,7 @@
 Unified approach for running TorchTitan models with vLLM inference.
 
 To register TorchTitan models with vLLM:
-    from torchtitan.components.checkpoint import CheckpointManager
+    from torchtitan.components.checkpointer import CheckpointManager
     from torchtitan.experiments.rl.models.vllm_registry import register_to_vllm
 
     # Standalone inference (loads HF weights):

torchtitan/experiments/rl/actors/generator.py

@@ -28,7 +28,7 @@
     PortReceiver,
 )
 from torch.distributed.tensor import DTensor
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.config import CompileConfig, Configurable, DebugConfig, OverrideConfig
 from torchtitan.distributed.parallel_dims import unfold_dp_axes
 from torchtitan.distributed.utils import get_spmd_backend, set_batch_invariance

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 import CheckpointManager
 from torchtitan.components.checkpoint_utils import canonical_fqn
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import BaseLoss, ChunkedLossWrapper
 from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer
 from torchtitan.config import (

torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py

@@ -13,7 +13,7 @@
 
 import dataclasses
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import ChunkedLossWrapper
 from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer
 from torchtitan.config import (

torchtitan/experiments/rl/examples/dapo_math/config_registry.py

@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import ChunkedLossWrapper
 from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer
 from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig

torchtitan/experiments/rl/examples/search_r1/config_registry.py

@@ -18,7 +18,7 @@
 
 import dataclasses
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import ChunkedLossWrapper
 from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer
 from torchtitan.config import (

torchtitan/experiments/rl/generate.py

@@ -30,7 +30,7 @@
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.distributed.utils import set_batch_invariance
 from torchtitan.experiments.rl.examples.alphabet_sort import config_registry
 from torchtitan.experiments.rl.models.vllm_registry import (

torchtitan/experiments/rl/models/vllm_registry.py

@@ -30,7 +30,7 @@
 from dataclasses import dataclass
 from typing import Any, Literal
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.config import CompileConfig, OverrideConfig, ParallelismConfig
 from torchtitan.protocols.model_spec import ModelSpec
 

torchtitan/experiments/rl/models/vllm_wrapper.py

@@ -20,7 +20,7 @@
 import torch
 import torch.distributed as dist
 from torch.distributed.tensor import DTensor, Replicate, Shard
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.config import (
     apply_overrides,
     CompileConfig,

torchtitan/experiments/rl/tests/test_bitwise_parity.py

@@ -53,7 +53,7 @@
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.loss import compute_logprobs, IGNORE_INDEX
 from torchtitan.config import CommConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed import ParallelDims, utils as dist_utils

torchtitan/experiments/rl/tests/test_generator.py

@@ -29,7 +29,7 @@
 from vllm import SamplingParams
 from vllm.sampling_params import RequestOutputKind
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.config import CommConfig, DebugConfig
 from torchtitan.distributed import utils as dist_utils
 from torchtitan.experiments.rl.actors.generator import (

torchtitan/experiments/torchft/checkpoint.py

@@ -23,7 +23,7 @@
 import torch.distributed as dist
 import torch.nn as nn
 
-from torchtitan.components.checkpoint import (
+from torchtitan.components.checkpointer import (
     AsyncMode,
     CheckpointManager,
     DATALOADER,

torchtitan/experiments/transformers_modeling_backend/config_registry.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import (
     ConcatThenSplitPackingConfig,
     FirstFitPackingConfig,

torchtitan/models/deepseek_v3/config_registry.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import ConcatThenSplitPackingConfig, GrainDataLoader
 from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss
 from torchtitan.components.metrics import MetricsProcessor

torchtitan/models/flux/config_registry.py

@@ -6,7 +6,7 @@
 
 from dataclasses import replace
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import GrainDataLoader, SingleDatasetConfig
 from torchtitan.components.loss import MSELoss
 from torchtitan.components.metrics import MetricsProcessor

torchtitan/models/gpt_oss/config_registry.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import ConcatThenSplitPackingConfig, GrainDataLoader
 from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss
 from torchtitan.components.metrics import MetricsProcessor

torchtitan/models/kimi_k2_7/config_registry.py

@@ -9,7 +9,7 @@
 
 from torch.distributed.tensor import Shard
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import (
     ConcatThenSplitPackingConfig,
     GrainDataLoader,

torchtitan/models/llama3/config_registry.py

@@ -6,7 +6,7 @@
 
 from typing import cast
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import (
     ConcatThenSplitPackingConfig,
     FirstFitPackingConfig,

torchtitan/models/muse_glimmer/config_registry.py

@@ -6,7 +6,7 @@
 
 from dataclasses import replace
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import ConcatThenSplitPackingConfig, GrainDataLoader
 from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss
 from torchtitan.components.metrics import MetricsProcessor

torchtitan/models/qwen3/config_registry.py

@@ -6,7 +6,7 @@
 
 from typing import cast
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import (
     ConcatThenSplitPackingConfig,
     FirstFitPackingConfig,

torchtitan/models/qwen3_5/config_registry.py

@@ -6,7 +6,7 @@
 
 from dataclasses import replace
 
-from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.checkpointer import CheckpointManager
 from torchtitan.components.data import GrainDataLoader, SingleDatasetConfig
 from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss
 from torchtitan.components.metrics import MetricsProcessor