[AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP#8168
[AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP#8168delock wants to merge 5 commits into
Conversation
AutoTP + ZeRO-3 silently produced incomplete checkpoints: both export paths handled only the ZeRO data-parallel dimension and dropped the tensor-parallel shards. - ds_to_universal.py: stage3 conversion recovers the (tp,dp) grid from checkpoint file names, extracts shards under the real tp_index, and reuses the stage<=2 TP-aware merge when tp_degree>1 (DP-only path preserved for tp_degree==1 -> no regression for plain ZeRO-3). - engine.py: _zero3_consolidated_16bit_state_dict nests GatherReplacedLayerParams inside GatheredParameters so save_16bit_model gathers both DP and TP; remove the blanket autotp+zero3 training block now that checkpoint consolidation is implemented. - stage3.py: load_hp_checkpoint_state resolves the TP shard before the ZeRO-DP partition, so universal checkpoint restore round-trips. Add end-to-end universal conversion tests and update existing tests for the refactored merge_tp_slices / extract_zero_shards_stage3 signatures. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
![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: 7b9dccb90c
ℹ️ 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".
| slice_shapes = { | ||
| name: shape | ||
| for sub_group_shapes in param_shapes | ||
| for name, shape in sub_group_shapes.items() if name in param_keys |
There was a problem hiding this comment.
For AutoTP ZeRO-3 checkpoints where the TP split produces unequal shard shapes (for example, a column-parallel dimension of 5 over TP=2, or a custom/GQA sub-parameter whose shards are not the same size), this builds slice_shapes from only param_shapes loaded from the first model-state file. _merge_zero_shards(..., slice_shape) then reshapes every TP rank's concatenated DP fragment to that first-rank shape, so any rank with a different shard size raises during conversion or is merged under the wrong shape; use the per-rank entries from param_shapes_grid when reconstructing each TP slice.
Useful? React with 👍 / 👎.
| # AutoTP + ZeRO stage 3 is supported for both inference and training. The | ||
| # checkpoint paths (16-bit consolidated export and universal checkpoint | ||
| # conversion) gather across both the ZeRO data-parallel and the tensor-parallel | ||
| # dimensions, so saved/converted weights are complete. |
There was a problem hiding this comment.
Since this change removes the ZeRO-3 guard and states AutoTP+ZeRO-3 training is supported, the public AutoTP docs are now contradictory: docs/code-docs/source/training.rst:395, docs/_pages/config-json.md:766, and docs/_tutorials/autotp-training.md:215 still tell users ZeRO Stage 3 is unsupported. Please update those docs with the new support and limitations so users do not avoid or misconfigure the newly enabled path.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
Address Codex review: the stage3 merge reshaped every TP slice to the shape loaded from the first model-state file, so uneven TP splits (a partition dim not divisible by tp_degree, or uneven GQA sub-params) would reshape a rank under the wrong shape. - _merge_zero_shards now takes a per-tp list of shapes and reshapes each TP slice to its own rank's shape. - merge_tp_slices threads per-tp shapes (instead of a single shape). - stage3 and stage<=2 callers build per-tp shape lists from each rank's PARAM_SHAPES / model-state file. Adds test_stage3_autotp_universal_conversion_uneven_tp (5-row weight over tp=2 -> 3+4 / 2+4 shards). Even-TP behavior is unchanged. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Address Codex review: training.rst, config-json.md, and autotp-training.md still stated ZeRO Stage 3 was unsupported, which is now contradictory. Update them to reflect that AutoTP + ZeRO-3 is supported (training + checkpoint save/load + universal checkpoint conversion), and note the remaining limitations (runtime TP divisibility for the TP math; cross-topology universal restore). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e1c059b47
ℹ️ 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".
|
|
||
| for state in ("fp32", "exp_avg", "exp_avg_sq"): | ||
| slices = _merge_zero_shards(slice_base_path, state, tp_degree, shape) | ||
| slices = _merge_zero_shards(slice_base_path, state, tp_degree, per_tp_shapes) |
There was a problem hiding this comment.
For parameters matching PARAMETER_WITH_SUB_PARAMS with uneven fused/GQA sizes, this per-TP merge still feeds into the branch below that re-views every slice to sum(sub_dims) // tp_degree and narrows each sub-param by sub_dim_size // tp_degree; when a sub-param dimension is not divisible by TP, a rank with the remainder shard either fails the .view() or loses those elements. Fresh evidence after the earlier comment is that the current branch still computes uniform partition_shape/part_sub_dim_size after this per-rank merge, despite the docs now advertising uneven fused/GQA checkpoint conversion.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
DeepSpeed's own AutoTP gather/partition infrastructure already blocks this case: _gather_logical_tensor (layers.py L1207) and _partition_logical_tensor (layers.py L1161) both raise ValueError when a sub-parameter size is not divisible by tp_size, so a model with non-divisible kV heads cannot be created under AutoTP in the first place. The per-tp-shape fix in the second commit addresses the real uneven scenario (a single parameter whose partition dim isn't divisible by TP, e.g. vocab size), which is the case that can actually occur. This sub-param path is a valid observation for a follow-up but doesn't block this PR.
| # fix back to normal flat dict, merge duplicates for tp>1 | ||
| slice_shapes = dict((k, v) for d in slice_shapes for k, v in d.items()) | ||
| slice_shapes_by_tp.append(dict((k, v) for d in mp_sd[PARAM_SHAPES] for k, v in d.items())) | ||
| slice_shapes = {name: [d[name] for d in slice_shapes_by_tp] for name in slice_shapes_by_tp[0]} |
There was a problem hiding this comment.
When converting ZeRO <= 2 checkpoints with pipeline parallelism (pp_degree > 1), mp_rank_files contains one model-state per TP×PP rank and each PP stage can have different PARAM_SHAPES; building slice_shapes from only the first file's keys and indexing every other stage dict by those names causes a KeyError for PP-local parameters, or omits parameters that only appear on later stages. The previous flatten/union logic handled the PP-local shape dictionaries, so this should group shapes by TP rank without requiring every PP file to contain the same names.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Thanks -- confirmed this was a regression from the per-TP refactor. Fixed in the latest commit: extracted a _group_per_tp_shapes helper that groups mp_rank_files by TP rank (they are tp-major) and unions PP stages within each TP rank, so PP-local params are preserved. Added test_group_per_tp_shapes_handles_pp_local_params covering the exact failure scenario (tp=2, pp=2, uneven TP shards, params only in pp0 / only in pp1).
9a879b5 to
6a5f66d
Compare
mp_rank_files are tp-major (tp0_pp0, tp0_pp1, ..., tp1_pp0, ...). Group by TP rank and union PP stages within each TP rank so PP-local parameters are not lost (fixes a regression from the per-TP refactor that indexed all (tp,pp) files flat and built the param-name set from the first file only). pp_degree==1 is unchanged. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
6a5f66d to
f3156b6
Compare
9ff49d9 to
ab2a220
Compare
ab2a220 to
4111888
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4111888b9d
ℹ️ 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".
| for tp in range(tp_degree): | ||
| tp_dict = {} | ||
| for pp in range(pp_degree): | ||
| tp_dict.update(slice_shapes_by_tp[tp * pp_degree + pp]) |
There was a problem hiding this comment.
For ZeRO <= 2 checkpoints with both PP and TP greater than one, this indexes mp_rank_files as TP-major, but DeepSpeedCheckpoint sorts those files and meg_2d_parallel_map.simple_init() maps index i to (pp=i // tp_degree, tp=i % tp_degree), making the order PP-major. Consequently each tp_dict combines different TP ranks from one pipeline stage; PP-local parameters then receive missing or incorrectly ranked shapes and _merge_zero_shards() can fail while reshaping or merge the wrong slices. Fresh evidence after the prior thread's claimed fix is the definitive production mapping in deepspeed/checkpoint/reshape_meg_2d.py:16-20; group entries using pp * tp_degree + tp (or the existing 2-D mapping) rather than tp * pp_degree + pp.
Useful? React with 👍 / 👎.
| "(no optimizer). Running it with an optimizer is disabled because TP-aware " | ||
| "checkpoint consolidation is not yet implemented and would silently produce " | ||
| "incomplete checkpoints.") | ||
| # AutoTP + ZeRO stage 3 is supported for both inference and training. The |
There was a problem hiding this comment.
The reviewed commit message has no Signed-off-by trailer, violating the repository's mandatory commit policy and potentially blocking DCO/CI acceptance. Recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
1 participant
AutoTP + ZeRO-3 silently produced incomplete checkpoints: both export paths handled only the ZeRO data-parallel dimension and dropped the tensor-parallel shards.
Add end-to-end universal conversion tests and update existing tests for the refactored merge_tp_slices / extract_zero_shards_stage3 signatures.