Add Feature Universal Checkpoint for AutoTP#7908
Add Feature Universal Checkpoint for AutoTP#7908nathon-lee wants to merge 15 commits intodeepspeedai:masterfrom
Conversation
This reverts commit ff88670. Co-authored-by: nathon-lee <248585198+nathon-lee@users.noreply.github.com>
Revert "fix: update 1 file reformatted." (ff88670)
This reverts commit b90aee5.
Revert accidental Muon optimizer code re-introduction from copilot PRs
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
nathon-lee
requested review from
hwchen2017,
loadams,
tjruwase and
tohtana
as code owners
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5ab684d950
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| conversion_meta = _get_param_uc_conversion_meta(param) | ||
| if not conversion_meta: | ||
| continue |
There was a problem hiding this comment.
AutoTP shards nn.Embedding weights in _slice_embedding() along hidden dimension (dim=1), but this collector skips any parameter without ds_autotp_universal_checkpoint_meta, so embedding parameters never enter UNIVERSAL_CHECKPOINT_INFO; during conversion they fall back to default merge behavior (cat_dim=0 in ds_to_universal.merge_tp_slices), which reconstructs embeddings along the wrong axis and breaks restores when TP degree changes.
Useful? React with 👍 / 👎.
| self._set_param_uc_meta(self.weight, | ||
| partition_type='column', | ||
| partition_dim=0, | ||
| logical_shape=original_weight_shape, | ||
| output_shape=(original_out_dim, ), | ||
| original_shape=original_weight_shape) |
There was a problem hiding this comment.
fused_LinearLayer partitions weights with prepare_tp_fused_qkvw() (model-specific Q/K/V reordering), but it inherits this generic LinearLayer metadata path, which records only plain column partitioning and no sub-parameter schema; converter logic then treats these tensors as simple concat-able shards, producing interleaved QKV layouts that are not portable across different TP sizes.
Useful? React with 👍 / 👎.
nathon-lee
changed the title
PawnOfDelock
left a comment
PawnOfDelock
left a comment
There was a problem hiding this comment.
📝 Review for PR #7908: AutoTP Universal Checkpoint Support
✅ Overall Impression
This is a well-structured and well-documented PR that adds critical functionality for AutoTP universal checkpoint support. The implementation is clean, tests are comprehensive, and CI passes.
🔍 Code Review
1. Design & Architecture 🎯
- ✅ Good separation of concerns between restore-time and conversion-time metadata
- ✅ The dual-view metadata design (top-level for restore, nested
conversionfor model-level aggregation) is elegant - ✅ Backward compatibility is properly preserved (non-AutoTP paths unchanged)
2. Implementation Details 💻
universal_checkpoint.py:
- ✅
_resolve_autotp_partitionhandles various partition scenarios correctly (row/column, sub-params, replicated) - ✅ Clean integration with existing
load_hp_checkpoint_state- minimal changes to existing logic - ✅ Proper shape normalization and error handling
layers.py:
- ✅ Consistent
_mark_uc_metadataimplementation across all TP layer types - ✅
collect_autotp_universal_checkpoint_infoproperly aggregates parameter-level metadata into model-level schema - ✅ Regex pattern generation is clean and correct
Optimizers (bf16_optimizer.py, stage_1_and_2.py):
- ✅
_enable_universal_checkpointproperly caches UC info from parameters - ✅ State dict integration is consistent across both optimizer types
engine.py:
- ✅ Properly collects and attaches UC info to model after AutoTP partitioning
- ✅ Checkpoint saving includes UC info
3. Testing 🧪
- ✅ Test coverage is excellent - both unit and integration tests
- ✅ Tests cover key scenarios:
- Row/column parallel weights
- Subparam partitioning
- Replicated biases
- Metadata aggregation
- Optimizer state handling
- ✅ Mocking is well-implemented
4. Documentation 📖
- ✅ Function docstrings are clear and helpful
- ✅ Code comments explain non-obvious logic
- ✅ PR description is comprehensive with motivation, implementation details, and testing instructions
🔮 Suggestions & Questions
-
Edge Cases:
- Consider adding a test for empty
sub_param_sizesscenario in_resolve_autotp_partition - Could add a test for mismatched
logical_shapevsoutput_shapeexpectations
- Consider adding a test for empty
-
Performance:
- The metadata collection happens at partition time (good)
- Consider if
collect_autotp_universal_checkpoint_infoshould be cached if called multiple times
-
Error Handling:
- What happens if
AUTOTP_UC_META_KEYexists but is malformed? Could add validation
- What happens if
-
Future Compatibility:
- The schema design is flexible for future additions - good job
⚠️ Minor Nitpicks
- In
universal_checkpoint.py, function signature:
def _resolve_autotp_partition(self, ckpt_dict, full_hp_param, tp_rank, tp_world_size):- The
selfparameter is passed but_resolve_autotp_partitionis not a method - consider making it clearer
- In
layers.py, could regex pattern generation be a utility function to avoid repetition?
These are very minor and don't block approval.
✅ CI Status
- ✅ All CI checks pass
- ✅ Tests run successfully
- ✅ DCO check passes
🎯 Recommendation
LGTM! 🚀
This PR is ready for merge. The implementation is solid, tests are comprehensive, and it properly integrates with existing checkpoint infrastructure.
Suggested follow-ups (post-merge):
- Add e2e integration tests with actual model training/saving/loading
- Document the UC metadata schema in user-facing docs
- Consider adding a migration guide for existing AutoTP checkpoints
| from .constants import (FP32_WEIGHT_KEY, PARAM, VOCAB_TENSOR, CAT_DIM, PARAM_N_SUB_PARAMS, SUB_PARAM_SHAPE) | ||
|
|
||
|
|
||
| AUTOTP_UC_META_KEY = 'ds_autotp_universal_checkpoint_meta' |
There was a problem hiding this comment.
This is not the only place the reference the string 'ds_autotp_universal_checkpoint_meta'. The other place is in layers.py with a different name DS_AUTOTP_UC_META. The name is not consistent. (DS_AUTOTP_UC_META looks better). If the string needs to be the same all over the place, should have a single definition (i.e. in this file) then import from other places (i.e. from deepspeed.checkpoint.universal_checkpoint.py import DS_AUTOTP_UC_META). There shouldn't be a second place with the content of the string, including tests.
| return getattr(param, AUTOTP_UC_META_KEY, None) | ||
|
|
||
|
|
||
| def _resolve_autotp_partition(self, ckpt_dict, full_hp_param, tp_rank, tp_world_size): |
There was a problem hiding this comment.
The parameter name self is not a proper name for a none class member function (This is also true for load_hp_checkpoint_state which is already there). If self mean current parameter, can we changes its name to current_param or a proper name? Thanks!
| is_bias = meta.get('is_bias', False) | ||
| replicated = meta.get('replicated', False) | ||
|
|
||
| if replicated or partition_dim is None: |
There was a problem hiding this comment.
Is there a scenario that one of them is True and another one is False? Probably better to have the following instead:
if replicated:
assert(partition_dim is None)
...
| if replicated or partition_dim is None: | ||
| slice_tensor = full_hp_param | ||
| else: | ||
| target_shape = output_shape if is_bias else logical_shape |
There was a problem hiding this comment.
This line implies that output_shape would be different from logical_shape when is_bias is True. Will this really happen?
| if target_shape is None: | ||
| return None | ||
|
|
||
| full_view = full_hp_param.view(target_shape) |
There was a problem hiding this comment.
full_view is not used in this if statement, better move it down below before the next if statement.
| return None | ||
|
|
||
| full_view = full_hp_param.view(target_shape) | ||
| if sub_param_sizes is None and sub_param_shape is not None: |
There was a problem hiding this comment.
Is there a case that sub_param_sizes cannot be deduced from sub_param_shape, and must save sub_param_size but not sub_param_shape when saving the meta data? I feel one of them is redudant.
deepspeed/module_inject/layers.py
Outdated
|
|
||
| for module_name, module in model.named_modules(): | ||
| marker = getattr(module, '_mark_uc_metadata', None) | ||
| if callable(marker): |
There was a problem hiding this comment.
Is it necessary to check whether marker is callable? I didn't see this possibility unless its an implementation error.
There was a problem hiding this comment.
Good point. _mark_uc_metadata is intended to be a callable hook; using callable() could hide an implementation error. I’ll remove the callable() guard and simply check for None so incorrect types fail fast.
| @@ -0,0 +1,241 @@ | |||
| import types | |||
There was a problem hiding this comment.
Thanks for the PR! The implementation looks good overall. Could you please add some documentation and usage descriptions?
There was a problem hiding this comment.
Thanks @inkcherry! Sure — I’ll add documentation and usage descriptions and update the PR shortly.
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
|
Hi Delock, thanks for the suggestion. I’ve centralized the AutoTP universal-checkpoint metadata attribute name into deepspeed/checkpoint/constants.py as DS_AUTOTP_UC_META, and updated both universal_checkpoint.py and module_inject/layers.py to import and use it (so we no longer hardcode the string in multiple places). Please let me know if you’d prefer I split the other _resolve_autotp_partition changes into a separate PR to keep this one focused. |
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
|
|
||
|
|
||
| def _write_tp_states(base_dir, param_name, tp_idx, fp32_tensor): | ||
| # merge_tp_slices 会尝试合并这三个 state，所以测试必须把它们都写出来 |
There was a problem hiding this comment.
This might not matter but I would suggest to keep comments in plain English for consistency.
| partition_dim=0, | ||
| name="packed") | ||
|
|
||
| weight_meta = getattr(layer.weight, "ds_autotp_universal_checkpoint_meta") |
There was a problem hiding this comment.
I would suggest to keep these string imported rather than hard coded.
| chunks = [sub.chunk(2, dim=0)[0] for sub in full_hp_param.view(3, 2, 4)] | ||
| expected = torch.cat(chunks, dim=0).flatten() | ||
| assert torch.equal(slice_flat, expected) | ||
|
|
There was a problem hiding this comment.
Is it possible to add a test when size of each subparam are not equal to each other? This could cover the case of GQA
sub_param_sizes 描述的是多个子参数在拼接维度上各自的大小不等的情况。
典型场景是 QKV fused linear，比如一个模型的 Q/K/V head 数量不同：
Q: 32 heads × 128 dim = 4096
K: 8 heads × 128 dim = 1024
V: 8 heads × 128 dim = 1024
拼在一起后 sub_param_shape = (4096, 1024, 1024)，总共 6144。
做 TP 切分时，如果简单地把 6144 均分成 N 份，会切到 Q/K/V 的边界中间。正确做法是按每个子参数各自均分：
TP=2 时:
Q: 4096/2 = 2048
K: 1024/2 = 512
V: 1024/2 = 512
每个 rank 拿到 3072
sub_param_sizes 就是告诉 restore 逻辑每个子参数有多大，这样才能按子参数分别切而不是整体均分。
在 _resolve_autotp_partition 中，当 sub_param_sizes 不为 None 时，代码会对每个子参数单独做 chunk/split，然后按 tp_rank 取对应的片段，再 concat 回来。而当 sub_param_sizes 为 None
时（所有子参数等大，或者根本没有子参数），直接整体 chunk 就行。
所以测试中缺的 case 就是：子参数大小不等时，restore 是否正确按各自大小切分。现有测试 2 的 sub_param_shape = ((2,2,2), 4) 三个子参数等大（都是 2），没有覆盖不等大的情况。
|
Hi @nathon-lee , thanks for the timely response! No need to split. I have given my comments. Overall it LGTM. I agree with @inkcherry that documentation better be updated, specifically,
|
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com> fix: update some logic for _resolve_autotp_partition Signed-off-by: nathon-lee <leejianwoo@gmail.com> fix: update some logic for test_load_hp_checkpoint_state_prefers_autotp_metadata Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Signed-off-by: nathon-lee <leejianwoo@gmail.com> docs: update universal checkpointing and AutoTP checkpoint docs Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Hi DeepSpeed team — thanks for your time reviewing this PR.
Summary
Add Universal Checkpoint (UC) metadata support for DeepSpeed AutoTP to enable saving and resuming from Universal Checkpoints.
Motivation
AutoTP partitions parameters across TP ranks. To make checkpoints portable and restorable, we need a stable UC metadata representation that can be collected at save time and consumed at restore time.
What’s in this PR
Testing
pytest -q tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.pypytest -q tests/unit/checkpoint/test_autotp_universal_checkpoint.pyRequest for feedback
Could you please take a look at the UC metadata schema and let me know if you’d prefer any changes to naming, field placement, or compatibility expectations? I’m happy to iterate quickly based on your guidance.
References