[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