[RFC] Support sharding LM heads and adopting Online Softmax
Author: jinyouzhiCreated Jul 24, 2026Updated Sep 17, 2026
Labelsenhancement
Motivation
Online softmax and in-place computation could significantly save memory, especially for OPSD teacher TP sharding.
- Extend AutoTP capabilities by integrating Liger Kernel (sharding LM heads and adopting Online Softmax) #8104
Keypoints
- Support the sharding LM heads with
gather_output=Falseof ColumnParallelLinear - Enabling parallel CE loss feature (integrating Liger kernel)
- Profiling the GPU footprint to evaluate
- Unify the untied and tied path for LM heads
Revised Roadmap
| PR | Scope | Status |
|---|---|---|
| #8146 | Enable gather_output=True for column-parallel lm_head |
Merged |
| #8185 | Uneven sharding + universal checkpoint; remove uneven-dim rowwise fallback | Merged |
| PR-A (#8302) | Lift hardcoded name == "lm_head" routing; training defaults to column parallelism; reject unsupported row-parallel training |
Merged |
| PR-B (#8309) | gather_output=False + vocab-parallel cross entropy, pure PyTorch reference implementation |
Ready to merge |
| PR-C | VocabParallelEmbedding; unify tied and untied lm_head paths; fix #8290 |
To do |
| PR-D (#8540) | Optional Liger fused CE acceleration + memory profiling | WIP |
| PR-E (#8538) | Implement row-parallel lm_head training, including input slicing and backward dx all-reduce |
WIP |
Proposed PR Sequencing
#8146 ──► #8185 ──┬─► PR-A Lift hardcoded routing; independently mergeable
│
└─► PR-B gather_output=False + vocab-parallel CE
├─► PR-C VocabParallelEmbedding / tied-untied unification
└─► PR-D Liger integration + profiling
PR-A ──► PR-E Row-parallel lm_head trainingPR-A is small and low-risk, and has no code-level dependency on gather_output;
it can be submitted against master in parallel. PR-B is the core of this RFC —
we suggest landing the pure PyTorch implementation first so that numerical
correctness can be reviewed independently, and deferring the Liger kernel
integration to PR-D.
Source: deepspeedai/DeepSpeed