TE AITER groupedgemm optimizations in Megatron

[sudhu2k]

Motivation

This PR enhances the Mixture of Experts (MoE) implementation in Megatron-LM to support GPU-based token processing for grouped GEMM operations. The changes enable improved performance when using Transformer Engine's grouped GEMM functionality with the Triton backend, while maintaining backward compatibility with existing configurations and older Transformer Engine versions. Maintaining the m_splits as a GPU tensor prevents any further redundant device to host to device memcpys and in turn improves throughput.

Performance difference:

https://github.com/ROCm/frameworks-internal/issues/13792#issuecomment-3849739037

Technical Details

The implementation introduces GPU token handling across three main components:

1. Token Dispatcher Enhancement (commit 218891725)

   • Modified MoEAlltoAllTokenDispatcher in megatron/core/transformer/moe/token_dispatcher.py to track tokens_per_expert_gpu alongside the existing CPU-based tokens_per_expert
   • Updated TEGroupedMLP in megatron/core/transformer/moe/experts.py to accept tuple inputs containing both CPU and GPU token tensors
   • Enhanced forward methods in megatron/core/extensions/transformer_engine.py to process GPU token inputs

2. Compatibility Checking (commit 33485e7bc)

   • Added class_has_method_param() utility function to dynamically verify if Transformer Engine's GroupedLinear.forward supports the m_splits_tensor parameter
   • Implemented version-aware parameter passing with user warnings when using older Transformer Engine versions lacking the m_splits_tensor parameter
   • Ensures compatibility with Transformer Engine commit 2776c33 which introduced GPU tensor support

3. Conditional GPU Processing (commit e782b8cf2)

   • Added environment variable check for NVTE_USE_GROUPED_GEMM_TRITON to conditionally enable GPU token processing
   • Updated return signatures in MoEAlltoAllTokenDispatcher.forward() to return tuples containing both CPU and GPU tensors when the environment variable is set
   • Ensures GPU tensors are converted to torch.int32 for compatibility with Transformer Engine

Test Plan

• Tested with and without NVTE_USE_GROUPED_GEMM_TRITON=1 environment variable
• Verified backward compatibility with existing MoE training scripts
• Validated functionality with both older and newer versions of Transformer Engine
• Ensured proper tensor type conversions and parameter passing through the MoE pipeline

Test Result

The changes successfully enable GPU-based token processing when the feature is enabled, while maintaining full backward compatibility. Warning messages are properly displayed when using older Transformer Engine versions, guiding users to upgrade for optimal performance.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.