[RFC] Support sharding LM heads and adopting Online Softmax

[jinyouzhi]

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) DeepSpeed - Q3 Roadmap #8104

Keypoints

• Support the sharding LM heads with gather_output=False of 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	WIP
PR-C	VocabParallelEmbedding; unify tied and untied lm_head paths; fix #8290	To do
PR-D	Optional Liger fused CE acceleration + memory profiling	To do
PR-E	Implement row-parallel lm_head training, including input slicing and backward dx all-reduce	To do

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 training

PR-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.

[jinyouzhi]

For the current gather_output=True design, chunking and fusing the Linear and cross-entropy computation, similar to nn.LinearCrossEntropyLoss in PyTorch 2.13, could reduce peak memory by avoiding full-logits materialization. However, compared with vocabulary sharding, the benefit is limited because the full vocabulary weight remains replicated and the existing TP communication is not eliminated.