[Dev] feat(moe): Support placing MTP layers into standalone stages

[BestJuly]

What does this PR do ?

Support placing MTP to a standalone stage, e.g., MTP in the 2nd last VPP stage for better VPP balance.
PR for main branch: PR2136

Highlights

This document outlines the design of a new MTP (Multi-Token Prediction) standalone feature. This feature allows MTP layers to be placed into a standalone VPP (Virtual Pipeline Parallelism) stage, rather than being confined to the last VPP stage.

Key Benefits:

• Enhanced Flexibility in MTP VPP Placement
  • MTP layers can now be positioned in either the last or second-to-last VPP stage, offering greater control over pipeline configuration.
• Improved Load Balancing
  • The MTP standalone feature isolates transformer block-related computations, which have a similar computational cost to a normal transformer block. This split, with MTP loss calculation remaining in the second-to-last VPP stage, contributes to a more balanced workload.
• Minimal Impact on Existing Code Paths
  • The design aims for near-zero impact on current code. A minor change involves the final layernorm layer, which will now be placed at the end of the decoder stage. In scenarios where the last stage contains only MTP layers and loss, this layernorm will reside in the second-to-last layer alongside a decoder layer. While this differs from previous implementations (where the final layernorm was always in the last stage), bitwise correctness has been validated under various conditions, including forceful layernorm placement and disabled gradient clipping.

Considerations:

• Changes in P2P Communication Shapes
  • The need to pass original and MTP hidden states to subsequent VPP stages will alter p2p communication shapes. To mitigate the impact on core communication, the shape is communicated upfront when MTP standalone settings are detected in the pipeline layout.

Design idea

We split MTP computation into two parts

• Output head computation.
• Others.

The advantage is that we put the MTP loss computation as main, so for the untie embedding part, the MTP loss can use exactly the same code path of the lm loss as they are in the same vpp stage, which is simpler than putting MTP loss calculation in the same vpp stage as MR#2996 so that we can avoid bugs related to this. With the new design, the DSv3 TFLOPs will also be compatible with the #2996 one which included in the 0.13 EA branch.We divide MTP computation into two parts: output head computation and others. This approach offers the advantage of treating MTP loss computation as primary. Consequently, for the untied embedding, the MTP loss can utilize the exact same code path as the LM loss, as they are within the same VPP stage. This simplifies the process compared to placing MTP loss calculation in the same VPP stage as MR#2996, thereby mitigating potential bugs. Furthermore, this new design ensures DSv3 TFLOPs compatibility with the #2996 version included in the 0.13 EA branch.

Perf and results

Perf

• The throughput in DSv3 is on par or even higher than 0.13 EA (where we use previous design for this standalone feature).

Convergence

On full DSv3
(Real training mode) Almost overlapped with the main branch settings for the metrics which is related to convergence.

Deterministic correctness check
We validate the correctness in mainly two parts (when using different pp layout, can be bitwise aligned when manually disabling grad clipping)

• Deterministic alignment between main branch and this MR.
  • The same pipeline layout settings;
  • Different pipeline layout settings.
• Deterministic alignment within this MR to validate the compatibility of 1F1B overlap.

[BestJuly]

/ok to test 4a342ec

[BestJuly]

Pipeline 37554648 triggered in gitlab for functional test validation.