From 32e482ae92cf8d287b1df102f67a10f5634ce18c 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 import path Summary: `components/checkpoint.py` is the last of the three re-export shims left behind when the checkpointer was grouped into a package, after the lr_scheduler shim in #4172 and checkpoint_utils in the preceding change. With it gone there are no compatibility shims left under `torchtitan/components/`. This one is a pure module rename at the callsite. `checkpointer/__init__.py` already re-exports exactly the same eight symbols the shim forwarded -- `AsyncMode`, `CheckpointManager`, `ModelWrapper`, and the `MODEL` / `OPTIMIZER` / `LR_SCHEDULER` / `DATALOADER` / `TRAIN_STATE` key constants -- so every importer changes only the module it names, with the imported names and their grouping untouched. Verified by parsing each importer and checking every imported symbol against the package's `__all__` before touching anything; nothing referenced a symbol the package does not expose, and no callsite used the plain `import torchtitan.components.checkpoint` form. Thirty modules are updated, spanning the checkpoint-conversion scripts, the forge, torchft, graph_trainer and rl experiments, and the unit tests. A thirty-first file, `experiments/rl/__init__.py`, carries the import inside its module docstring as a usage example; that is updated too, so the documented path matches the working one. This is an import-path change only; no runtime behavior changes. Test Plan: Full `pytest tests/unit_tests` (excluding `test_rope.py`, which cannot be collected without the optional `fla` package): 622 passed, 18 failed. The 18 are the same set that fails on unmodified main in this environment -- missing optional dependencies (`transformers`, `fla`) and environment-specific kernel/compile failures (helion rope, inductor lora). Also verified: - No `components.checkpoint` references remain anywhere in the repo, across all file types, and no `Compatibility imports` shim remains under `torchtitan/components/`. - `import torchtitan.components.checkpoint` now raises `ModuleNotFoundError`. - All eight symbols import cleanly from `torchtitan.components.checkpointer`. - All 31 changed files byte-compile, and the affected non-test modules (`trainer`, `forge.engine`, `torchft.checkpoint`, `torchft.optimizer`, both checkpoint-conversion scripts) import cleanly. `experiments.rl` fails only on the absent optional `vllm` package, unrelated to this change. - `ufmt` and `flake8 --config=.flake8` clean on all 31 files. --- .../checkpoint_conversion/convert_from_hf.py | 2 +- .../checkpoint_conversion/convert_to_hf.py | 2 +- .../numerical_tests_example.py | 2 +- .../numerical_tests_kimi.py | 2 +- .../numerical_tests_qwen3_5.py | 2 +- tests/unit_tests/test_checkpoint.py | 2 +- tests/unit_tests/test_state_dict_keys.py | 4 +-- torchtitan/components/checkpoint.py | 28 ------------------- torchtitan/experiments/forge/engine.py | 2 +- .../tests/test_bitwise_deterministic.py | 2 +- .../tests/test_graph_pp_passes.py | 2 +- torchtitan/experiments/rl/__init__.py | 2 +- torchtitan/experiments/rl/actors/generator.py | 2 +- torchtitan/experiments/rl/actors/trainer.py | 2 +- .../examples/alphabet_sort/config_registry.py | 2 +- .../rl/examples/dapo_math/config_registry.py | 2 +- .../rl/examples/search_r1/config_registry.py | 2 +- torchtitan/experiments/rl/generate.py | 2 +- .../experiments/rl/models/vllm_registry.py | 2 +- .../experiments/rl/models/vllm_wrapper.py | 2 +- .../rl/tests/test_bitwise_parity.py | 2 +- .../experiments/rl/tests/test_generator.py | 2 +- torchtitan/experiments/torchft/checkpoint.py | 2 +- .../config_registry.py | 2 +- .../models/deepseek_v3/config_registry.py | 2 +- torchtitan/models/flux/config_registry.py | 2 +- torchtitan/models/gpt_oss/config_registry.py | 2 +- .../models/kimi_k2_7/config_registry.py | 2 +- torchtitan/models/llama3/config_registry.py | 2 +- .../models/muse_glimmer/config_registry.py | 2 +- torchtitan/models/qwen3/config_registry.py | 2 +- torchtitan/models/qwen3_5/config_registry.py | 2 +- 32 files changed, 32 insertions(+), 60 deletions(-) delete mode 100644 torchtitan/components/checkpoint.py diff --git a/scripts/checkpoint_conversion/convert_from_hf.py b/scripts/checkpoint_conversion/convert_from_hf.py index ab65f806ef..2997d1e27a 100644 --- a/scripts/checkpoint_conversion/convert_from_hf.py +++ b/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() diff --git a/scripts/checkpoint_conversion/convert_to_hf.py b/scripts/checkpoint_conversion/convert_to_hf.py index 9285f67a72..827c938a20 100644 --- a/scripts/checkpoint_conversion/convert_to_hf.py +++ b/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 diff --git a/scripts/checkpoint_conversion/numerical_tests_example.py b/scripts/checkpoint_conversion/numerical_tests_example.py index d6a23d2da9..debd52c5f2 100644 --- a/scripts/checkpoint_conversion/numerical_tests_example.py +++ b/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 diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi.py b/scripts/checkpoint_conversion/numerical_tests_kimi.py index ba4471144c..515e8ff98d 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi.py +++ b/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, diff --git a/scripts/checkpoint_conversion/numerical_tests_qwen3_5.py b/scripts/checkpoint_conversion/numerical_tests_qwen3_5.py index d120461a40..298eff6258 100644 --- a/scripts/checkpoint_conversion/numerical_tests_qwen3_5.py +++ b/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, diff --git a/tests/unit_tests/test_checkpoint.py b/tests/unit_tests/test_checkpoint.py index 96c1194d42..0c686bea02 100644 --- a/tests/unit_tests/test_checkpoint.py +++ b/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, ) diff --git a/tests/unit_tests/test_state_dict_keys.py b/tests/unit_tests/test_state_dict_keys.py index daf640fcd0..b77781c1c9 100644 --- a/tests/unit_tests/test_state_dict_keys.py +++ b/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 diff --git a/torchtitan/components/checkpoint.py b/torchtitan/components/checkpoint.py deleted file mode 100644 index d05e6ac84b..0000000000 --- a/torchtitan/components/checkpoint.py +++ /dev/null @@ -1,28 +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 the relocated DCP checkpoint manager.""" - -from .checkpointer.base import ( - DATALOADER, - LR_SCHEDULER, - MODEL, - ModelWrapper, - OPTIMIZER, - TRAIN_STATE, -) -from .checkpointer.dcp import AsyncMode, CheckpointManager - -__all__ = [ - "AsyncMode", - "CheckpointManager", - "DATALOADER", - "LR_SCHEDULER", - "MODEL", - "ModelWrapper", - "OPTIMIZER", - "TRAIN_STATE", -] diff --git a/torchtitan/experiments/forge/engine.py b/torchtitan/experiments/forge/engine.py index f7b14f42e3..5ca80b865b 100644 --- a/torchtitan/experiments/forge/engine.py +++ b/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 diff --git a/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py b/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py index 7339a23789..3de4877ed7 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py +++ b/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 diff --git a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py index 44a465ee54..6fd75127a7 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py +++ b/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 ( diff --git a/torchtitan/experiments/rl/__init__.py b/torchtitan/experiments/rl/__init__.py index 6643ca8dec..eddfeed3e0 100644 --- a/torchtitan/experiments/rl/__init__.py +++ b/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): diff --git a/torchtitan/experiments/rl/actors/generator.py b/torchtitan/experiments/rl/actors/generator.py index 9e00f19631..3c21b941a7 100644 --- a/torchtitan/experiments/rl/actors/generator.py +++ b/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 diff --git a/torchtitan/experiments/rl/actors/trainer.py b/torchtitan/experiments/rl/actors/trainer.py index 770d936ad7..c7accfdb91 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 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 ( diff --git a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py index 016e1f854d..58b6d57402 100644 --- a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py +++ b/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 ( diff --git a/torchtitan/experiments/rl/examples/dapo_math/config_registry.py b/torchtitan/experiments/rl/examples/dapo_math/config_registry.py index 6b23f1b022..a93cee5b2f 100644 --- a/torchtitan/experiments/rl/examples/dapo_math/config_registry.py +++ b/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 diff --git a/torchtitan/experiments/rl/examples/search_r1/config_registry.py b/torchtitan/experiments/rl/examples/search_r1/config_registry.py index eee8d093ee..90a7e1de7c 100644 --- a/torchtitan/experiments/rl/examples/search_r1/config_registry.py +++ b/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 ( diff --git a/torchtitan/experiments/rl/generate.py b/torchtitan/experiments/rl/generate.py index bd797303a0..ad6a97103f 100755 --- a/torchtitan/experiments/rl/generate.py +++ b/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 ( diff --git a/torchtitan/experiments/rl/models/vllm_registry.py b/torchtitan/experiments/rl/models/vllm_registry.py index 2bbc05c9f7..5b6e0ae4e0 100644 --- a/torchtitan/experiments/rl/models/vllm_registry.py +++ b/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 diff --git a/torchtitan/experiments/rl/models/vllm_wrapper.py b/torchtitan/experiments/rl/models/vllm_wrapper.py index 16f6b14cc0..686b4a5798 100644 --- a/torchtitan/experiments/rl/models/vllm_wrapper.py +++ b/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, diff --git a/torchtitan/experiments/rl/tests/test_bitwise_parity.py b/torchtitan/experiments/rl/tests/test_bitwise_parity.py index 6f27dc3866..4cdf9c4e7a 100644 --- a/torchtitan/experiments/rl/tests/test_bitwise_parity.py +++ b/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 diff --git a/torchtitan/experiments/rl/tests/test_generator.py b/torchtitan/experiments/rl/tests/test_generator.py index da03229de2..ad3c1fd35f 100644 --- a/torchtitan/experiments/rl/tests/test_generator.py +++ b/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 ( diff --git a/torchtitan/experiments/torchft/checkpoint.py b/torchtitan/experiments/torchft/checkpoint.py index a0982b3004..3b021d5c9c 100644 --- a/torchtitan/experiments/torchft/checkpoint.py +++ b/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, diff --git a/torchtitan/experiments/transformers_modeling_backend/config_registry.py b/torchtitan/experiments/transformers_modeling_backend/config_registry.py index eb7ba12aa5..01f4047fc6 100644 --- a/torchtitan/experiments/transformers_modeling_backend/config_registry.py +++ b/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, diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 94e71ecd4d..8711b8a253 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/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 diff --git a/torchtitan/models/flux/config_registry.py b/torchtitan/models/flux/config_registry.py index 55a3af2ca5..a9ae89f1f4 100644 --- a/torchtitan/models/flux/config_registry.py +++ b/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 diff --git a/torchtitan/models/gpt_oss/config_registry.py b/torchtitan/models/gpt_oss/config_registry.py index fb9528d971..7b9c8403b6 100644 --- a/torchtitan/models/gpt_oss/config_registry.py +++ b/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 diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index ac406205be..f0046535c4 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/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, diff --git a/torchtitan/models/llama3/config_registry.py b/torchtitan/models/llama3/config_registry.py index 4ad0e33632..90639954da 100644 --- a/torchtitan/models/llama3/config_registry.py +++ b/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, diff --git a/torchtitan/models/muse_glimmer/config_registry.py b/torchtitan/models/muse_glimmer/config_registry.py index 234807a4b5..57e0ec817d 100644 --- a/torchtitan/models/muse_glimmer/config_registry.py +++ b/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 diff --git a/torchtitan/models/qwen3/config_registry.py b/torchtitan/models/qwen3/config_registry.py index e6dfd7cc8f..56cd8be816 100644 --- a/torchtitan/models/qwen3/config_registry.py +++ b/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, diff --git a/torchtitan/models/qwen3_5/config_registry.py b/torchtitan/models/qwen3_5/config_registry.py index 60ddd7de65..6460316b67 100644 --- a/torchtitan/models/qwen3_5/config_registry.py +++ b/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