Conversation
ivy-zhou
requested review from
IvanKobzarev,
SherlockNoMad,
aditvenk,
fegin,
felipemello1,
sanketpurandare,
tianyu-l,
wconstab,
wwwjn and
xmfan
as code owners
August 20, 2026 00:54
ivy-zhou
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 20, 2026 00:54
This was referenced Aug 20, 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
tianyu-l
approved these changes
Aug 20, 2026
tianyu-l
approved these changes
Aug 20, 2026
vkuzo
reviewed
Aug 20, 2026
| filter_fn=lambda mod, _fqn: isinstance(mod, GroupedExperts), | ||
| self._float8_op_config = Float8TrainingOpConfig() | ||
|
|
||
| def _grouped_mm(self, *, A, B_t, offs): |
Summary: Keep Float8 grouped-expert parameters as ordinary tensors and route the existing grouped-matmul seam through TorchAO stateless dynamic quantization. This avoids serializing TorchAO wrapper/config objects while preserving the FP8 computation path. Test Plan: - pytest -q tests/unit_tests/test_quantization.py with current TorchAO and 0.18.0 - pre-commit run --all-files - H100 eager and fullgraph-compiled forward/backward parity - 2-GPU FSDP2 + DCP save/load forward/backward parity
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
3 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: Keep Float8 grouped-expert parameters as ordinary tensors and route the existing grouped-matmul seam through TorchAO stateless dynamic quantization. This avoids serializing TorchAO wrapper/config objects while preserving the FP8 computation path.
Test Plan: