[AutoTP] Enable HF `colwise_gather_output` to support `lm_head` replace by jinyouzhi · Pull Request #8146 · deepspeedai/DeepSpeed

jinyouzhi · 2026-07-15T20:53:06Z

Changes

Enable HF colwise_gather_output in tp_plan through AutoTP, allowing vocab parallelism for an untied lm_head.

Known Limitations

Only supports an untied lm_head. Tied lm_head falls back to the legacy implementation.
Uneven tensor parallelism is not supported in this PR.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 386d2bf3bf

ℹ️ 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".

…iguration(`colwise_gather_output`/`colwise_rep`) - Introduced `gather_output` option in `TPLayerSpec` for column-parallel layers. - Added validation to prevent using gathered output with tied embeddings. - Updated `LinearLayer` to handle gathered output during forward pass. - Enhanced documentation to reflect new gathered column parallelism capabilities. - Added tests for gathered column parallelism scenarios and configurations. Signed-off-by: iLeGend <824040212@qq.com>

…e tp_plan extraction logic Signed-off-by: iLeGend <824040212@qq.com>

… for fallback behavior Signed-off-by: iLeGend <824040212@qq.com>

Signed-off-by: iLeGend <824040212@qq.com>

PKUWZP

Thanks for submitting the PR, a couple of suggestions:

I don't see this PR leverages online softmax and in-place computation, this is the common optimization that saves GPU memory significantly.
We need to run some profiling to measure the GPU footprint.

Happy to work closely together to have a separate PR for this.

jinyouzhi · 2026-07-23T08:39:01Z

Thanks for submitting the PR, a couple of suggestions:

I don't see this PR leverages online softmax and in-place computation, this is the common optimization that saves GPU memory significantly.

We need to run some profiling to measure the GPU footprint.

Happy to work closely together to have a separate PR for this.

Thank you for your suggestions! You‘re absolutely right, this PR is far from the target. I've refined the PR description, and I'll submit a separate PR for those optimizations.

Signed-off-by: iLeGend <824040212@qq.com>

delock · 2026-07-24T08:00:45Z

Thanks for submitting the PR, a couple of suggestions:

I don't see this PR leverages online softmax and in-place computation, this is the common optimization that saves GPU memory significantly.

We need to run some profiling to measure the GPU footprint.

Happy to work closely together to have a separate PR for this.

Suggest profile GPU foot print with the separate PR. Online softmax and in-place computate could save memory for OPSD teacher TP sharding.

jinyouzhi · 2026-07-24T09:01:02Z

saves GPU memory significantly.

Great! list the keypoints into an issue to track #8173

jinyouzhi · 2026-07-27T15:58:42Z

Hi @PKUWZP, Could you have a look?

PKUWZP

LGTM.

…an (deepspeedai#8187) ## Summary Attach AutoTP universal checkpoint metadata before returning from a successfully converted HuggingFace `tp_plan`. Previously, this path replaced modules and returned without preserving the metadata required by `ds_to_universal`, unlike custom-config and heuristic AutoTP paths. ## Testing Add coverage for: - `colwise` - `rowwise` - replicated row-parallel bias - `colwise_gather_output` (Due to deepspeedai#8146 merged) --------- Signed-off-by: iLeGend <824040212@qq.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

…eepspeedai#8185) Follow up deepspeedai#8146. ## Summary This pull request introduces support for uneven sub-parameter sharding in DeepSpeed's universal checkpoint conversion, updating the universal checkpoint format to version 0.4. The changes ensure that partitioned parameters with sub-parameters of varying sizes are correctly handled during checkpoint conversion, merging, and restoration. Additionally, the PR adds validation to prevent conversion of unsupported checkpoint layouts and improves error handling and metadata validation. Key updates by theme: **Universal Checkpoint Format and Metadata:** - Bumped the universal checkpoint version to 0.4 and introduced the `SUB_PARAM_SHARD_WIDTHS` field to record per-rank widths for each sub-parameter, enabling correct handling of uneven sub-parameter layouts. (`deepspeed/checkpoint/constants.py`, `deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eL61-R64) [[2]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eR93-R97) [[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R35) [[4]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R302-R303) [[5]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L335-R441) - Added `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` to checkpoint metadata and implemented validation to prevent conversion if unsupported parameter patterns are present. (`deepspeed/checkpoint/constants.py`, `deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eL61-R64) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R47) [[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R950-R957) [[4]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R1134-R1139) **Parameter Merging and Sharding Logic:** - Enhanced the merging logic to correctly handle missing fragments for ranks with no data in uneven parameter sharding, ensuring proper alignment of slices and placeholder insertion. (`deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R239-R248) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R273-R285) - Refactored the sub-parameter merging code to use the new shard widths metadata, supporting both legacy (even) and new (uneven) layouts. Added logic to reconstruct logical shapes with placeholder dimensions. (`deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L302-R327) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L335-R441) [[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R463-R482) **Validation and Error Handling:** - Added early validation for unsupported AutoTP conversions to fail fast before expensive extraction steps. (`deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R950-R957) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R1134-R1139) - Improved shape consistency checks for pipeline-parallel parameters to ensure all replicas agree on shape. (`deepspeed/checkpoint/ds_to_universal.py`) **Restoration Logic:** - Updated the restoration logic to use the new shard widths metadata, ensuring correct reconstruction of sub-parameters during model loading. (`deepspeed/checkpoint/universal_checkpoint.py`) [[1]](diffhunk://#diff-22d7b9e3b6eac1dc6e989cd8582946321c36f4fc80b527648eeb2d77a8fd3ee2L13-R13) [[2]](diffhunk://#diff-22d7b9e3b6eac1dc6e989cd8582946321c36f4fc80b527648eeb2d77a8fd3ee2R34-R73) These changes collectively improve the robustness and flexibility of DeepSpeed's checkpoint conversion, especially for advanced tensor parallelism scenarios. ## Testing * Added coverage for uneven vocabulary, GQA projections, checkpoint conversion/restore, and PP + TP tied parameters. * loss curve deepspeedai/DeepSpeedExamples#1008 ## Limitations tp_size > num_kv need to further be optimized --------- Signed-off-by: iLeGend <824040212@qq.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ma,Guokai <guokai.ma@intel.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

jinyouzhi requested review from hwchen2017, loadams, tjruwase and tohtana as code owners July 15, 2026 20:53

jinyouzhi force-pushed the lm_head_colwise branch from 386d2bf to 67285d6 Compare July 15, 2026 20:54

chatgpt-codex-connector Bot reviewed Jul 15, 2026

View reviewed changes

Comment thread deepspeed/module_inject/auto_tp.py

jinyouzhi force-pushed the lm_head_colwise branch 3 times, most recently from 2ab4f47 to c63d2d4 Compare July 16, 2026 13:56

jinyouzhi added 5 commits July 16, 2026 21:58

Enhance AutoTP logging for lm_head diagnostics and improve HuggingFac…

4a45694

…e tp_plan extraction logic Signed-off-by: iLeGend <824040212@qq.com>

Refactor AutoTP gathered column tie handling and update related tests…

0e2a0c2

… for fallback behavior Signed-off-by: iLeGend <824040212@qq.com>

fix the AttributeError for _gathered_column_tie_fallbacks_configured

02b214c

Signed-off-by: iLeGend <824040212@qq.com>

handle uneven gathered lm_head shards

c63d2d4

Signed-off-by: iLeGend <824040212@qq.com>

PKUWZP self-requested a review July 16, 2026 19:56

delock reviewed Jul 17, 2026

View reviewed changes

Comment thread deepspeed/module_inject/auto_tp.py

PKUWZP requested changes Jul 17, 2026

View reviewed changes

delock reviewed Jul 17, 2026

View reviewed changes

Comment thread deepspeed/module_inject/autotp_config.py Outdated

delock reviewed Jul 17, 2026

View reviewed changes

Comment thread deepspeed/module_inject/auto_tp.py

delock mentioned this pull request Jul 22, 2026

Add AutoTP=2 smoke test config and script deepspeedai/DeepSpeedExamples#1004

Merged

jinyouzhi added 4 commits July 23, 2026 16:45

fix pattern

bca340e

Signed-off-by: iLeGend <824040212@qq.com>

fix code style

da9d0db

Signed-off-by: iLeGend <824040212@qq.com>

Merge remote-tracking branch 'origin/master' into HEAD

fdfce90

add ut for uneven vocab slice

afc886f

Signed-off-by: iLeGend <824040212@qq.com>

delock enabled auto-merge July 24, 2026 08:01

jinyouzhi mentioned this pull request Jul 24, 2026

AutoTP for training skips lm_head ? #7407

Open

Merge branch 'master' into lm_head_colwise

5f982d6

delock requested a review from PKUWZP July 25, 2026 08:04

Merge branch 'master' into lm_head_colwise

e377fef

jinyouzhi mentioned this pull request Jul 28, 2026

[AutoTP] Complete uneven sharding and universal checkpoint support #8185

Merged

jinyouzhi changed the title ~~[AutoTP] Enalbe HF colwise_gather_output to support lm_head replace~~ [AutoTP] Enable HF colwise_gather_output to support lm_head replace Jul 28, 2026

delock disabled auto-merge July 28, 2026 08:02

delock enabled auto-merge July 28, 2026 08:03

jinyouzhi mentioned this pull request Jul 28, 2026

[RFC] Support sharding LM heads and adopting Online Softmax #8173

Open

Merge branch 'master' into lm_head_colwise

90edcf2

PKUWZP approved these changes Jul 29, 2026

View reviewed changes

delock added this pull request to the merge queue Jul 29, 2026

Merged via the queue into deepspeedai:master with commit 48d54e0 Jul 29, 2026
13 checks passed

jinyouzhi mentioned this pull request Jul 29, 2026

[AutoTP] Preserve universal checkpoint metadata for HuggingFace tp_plan #8187

Merged

delock mentioned this pull request Aug 25, 2026

[AutoTP] Fix training lm_head routing #8302

Open

Navigation Menu

Uh oh!

[AutoTP] Enable HF `colwise_gather_output` to support `lm_head` replace - #8146

[AutoTP] Enable HF `colwise_gather_output` to support `lm_head` replace#8146
delock merged 12 commits into
deepspeedai:masterfrom
jinyouzhi:lm_head_colwise

jinyouzhi commented Jul 15, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

PKUWZP left a comment

Uh oh!

Uh oh!

Uh oh!

jinyouzhi commented Jul 23, 2026

Uh oh!

delock commented Jul 24, 2026

Uh oh!

jinyouzhi commented Jul 24, 2026

Uh oh!

jinyouzhi commented Jul 27, 2026

Uh oh!

PKUWZP left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

jinyouzhi commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Known Limitations

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

PKUWZP left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jinyouzhi commented Jul 23, 2026

Uh oh!

delock commented Jul 24, 2026

Uh oh!

jinyouzhi commented Jul 24, 2026

Uh oh!

jinyouzhi commented Jul 27, 2026

Uh oh!

PKUWZP left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jinyouzhi commented Jul 15, 2026 •

edited

Loading