Conversation
This was referenced Aug 19, 2026
meta-cla
Bot
added
the
CLA Signed
ivy-zhou
added a commit
that referenced
this pull request
Aug 20, 2026
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. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4184). * #4240 * #4230 * #4191 * #4190 * #4189 * #4188 * #4197 * #4187 * #4186 * #4185 * __->__ #4184
ivy-zhou
added a commit
that referenced
this pull request
Aug 20, 2026
…4185) 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. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/pytorch/torchtitan/pull/4185). * #4240 * #4230 * #4191 * #4190 * #4189 * #4188 * #4197 * #4187 * #4186 * __->__ #4185 * #4184
ivy-zhou
marked this pull request as ready for review
August 20, 2026 02:39
ivy-zhou
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 20, 2026 02:39
tianyu-l
approved these changes
Aug 20, 2026
Summary: Ports D115876261 to OSS. Needed by the DCPv2 512-GPU panel: the consolidation (Hugging Face export) cases cannot run without it. `_get_dtensor_shard_mesh_names` walked every placement on the mesh and read `placement.dim`, which only exists on `Shard` and `_StridedShard`. A MoE mesh that carries a `Replicate()` placement therefore raised `AttributeError` during HF export. Guard the comparison with an `isinstance` check so replicated placements are skipped instead of dereferenced. Adds `tests/unit_tests/test_state_dict_adapter.py` covering a DeepSeek V3 adapter export over a mesh that mixes `Replicate()` with sharded placements, which reproduces the failure without the fix. Test Plan: `pytest -q tests/unit_tests/test_state_dict_adapter.py` -> 1 passed in 3.31s. `pre-commit run --files torchtitan/models/utils.py tests/unit_tests/test_state_dict_adapter.py` passes all 13 hooks, including pyrefly. Verified the port is faithful: the fbsource patch restricted to `fbcode/pytorch/torchtitan` is byte-identical between the local draft and the published D115876261, and applies to the OSS tree with no fuzz after stripping the mirror prefix. `Shard` and `_StridedShard` were already imported in `torchtitan/models/utils.py`, and both symbols the new test imports (`deepseekv3_configs`, `DeepSeekV3StateDictAdapter`) exist in this tree.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Ports D115876261 to OSS. Needed by the DCPv2 512-GPU panel: the consolidation
(Hugging Face export) cases cannot run without it.
_get_dtensor_shard_mesh_nameswalked every placement on the mesh and readplacement.dim, which only exists onShardand_StridedShard. A MoE meshthat carries a
Replicate()placement therefore raisedAttributeErrorduringHF export. Guard the comparison with an
isinstancecheck so replicatedplacements are skipped instead of dereferenced.
Adds
tests/unit_tests/test_state_dict_adapter.pycovering a DeepSeek V3adapter export over a mesh that mixes
Replicate()with sharded placements,which reproduces the failure without the fix.
Test Plan:
pytest -q tests/unit_tests/test_state_dict_adapter.py-> 1 passed in 3.31s.pre-commit run --files torchtitan/models/utils.py tests/unit_tests/test_state_dict_adapter.pypasses all 13 hooks, includingpyrefly.
Verified the port is faithful: the fbsource patch restricted to
fbcode/pytorch/torchtitanis byte-identical between the local draft and thepublished D115876261, and applies to the OSS tree with no fuzz after stripping
the mirror prefix.
Shardand_StridedShardwere already imported intorchtitan/models/utils.py, and both symbols the new test imports(
deepseekv3_configs,DeepSeekV3StateDictAdapter) exist in this tree.