Add AMD ROCm / MI300X optimized MoE framework for Megatron-LM

[Copilot]

Megatron-LM had no AMD/ROCm support. This PR introduces megatron/core/amd_rocm/ and integrates AMD MI300X hardware optimizations into the MoE training path. All changes are zero-impact on NVIDIA hardware (conditional imports + runtime detection).

Key changes

New module: megatron/core/amd_rocm/

File	Purpose
device_utils.py	ROCm/HIP detection (is_rocm_available, is_amd_gpu, is_mi300x), HBM/XGMI property queries; covers MI200 (gfx90a) through MI300X (gfx942)
rocm_config.py	ROCmMoEConfig dataclass — XGMI-aware dispatch, HBM capacity check, hipBLASLt grouped GEMM, load-balancing strategy
moe_dispatcher.py	ROCmMoEAlltoAllDispatcher extending MoEAlltoAllTokenDispatcher with: (1) hierarchical XGMI A2A — intra-node fast XGMI first, then inter-node RCCL/IB; (2) HBM capacity check at init; (3) ROCM_USE_GROUPED_GEMM=1 hint for hipBLASLt batched expert GEMM

Modified files

• transformer_config.py: Five new moe_rocm_* fields — moe_use_rocm_dispatcher, moe_rocm_xgmi_aware_dispatch, moe_rocm_intra_node_ep_size, moe_rocm_use_hipblaslt_grouped_gemm, moe_rocm_hbm_capacity_check. All default to off; ignored on NVIDIA.
• moe_layer.py: Selects ROCmMoEAlltoAllDispatcher when moe_use_rocm_dispatcher=True and AMD hardware is detected; silently falls back to MoEAlltoAllTokenDispatcher otherwise.
• initialize.py: Logs MI300X capabilities (HBM, XGMI BW, CU count, ROCm version) on rank-0 at startup.

Usage

config = TransformerConfig(
    num_moe_experts=64,
    expert_model_parallel_size=8,
    moe_token_dispatcher_type="alltoall",
    moe_use_rocm_dispatcher=True,           # activates ROCmMoEAlltoAllDispatcher on AMD
    moe_rocm_xgmi_aware_dispatch=True,      # 2-phase A2A: intra-node XGMI then inter-node
    moe_rocm_intra_node_ep_size=8,          # auto-detected if omitted
    moe_rocm_use_hipblaslt_grouped_gemm=True,
)

On non-AMD hardware moe_use_rocm_dispatcher=True is silently ignored.

New artifacts

• tests/test_utils/recipes/mi300x/{gpt,moe}.yaml — MI300X CI recipes (mirroring H100 suite + MI300X-specific large expert-pool tests)
• tests/unit_tests/transformer/moe/amd_rocm/test_amd_moe.py — 20+ unit tests covering device detection, config validation, HBM capacity check, hierarchical A2A, and TransformerConfig AMD fields
• docs/amd_rocm/ — Chinese and English design docs explaining the architecture and rationale

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

---

🔒 GitHub Advanced Security automatically protects Copilot coding agent pull requests. You can protect all pull requests by enabling Advanced Security for your repositories. Learn more about Advanced Security.