NVIDIA / NVIDIA/Megatron-LM

Offload MoE expert weights to pinned host memory to enable training at smaller EP

Open
#6,491 11 comments 0 reactions 1 assignee Claimed by @lhb8125 View on GitHub
community-request enhancement waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

Propose MoE expert offloading, which keeps expert weights on pinned CPU RAM and stream them to GPU chunk-by-chunk during forward and backward, overlapping each chunk's H2D copy with the previous chunk's GEMM. This is an internal feature for the training of [Apertus](https://www.apertus-ai.org/) model and we would like to open a PR (#6494) for it.

### Motivation

On clusters with limited inter-node all-to-all bandwidth, expert parallelism does not scale well when it comes to inter-node communication. On clusters that use GH200 with Slingshot inter-node connection (i.e., [alps](https://www.cscs.ch/computers/alps), [Isambard](https://docs.isambard.ac.uk/specs/)), EP > 4 is mostly dominated by A2A cost. This brings a problem to train an MoE model: if the model requires cross-node EP to fit into the GPU memory, the training efficiency will be low.

Take the training of Qwen3-30B-A3B on GH200 as an example. Following the Megatron-bridge recipe of EP16-TP1-PP1, the training on 64 GH200 is fully EP communication dominated:

Image

Reducing EP size is only possible if memory permits. The existing memory-reduction features in Megatron do not fully address the problem, or they are more useful in other scenarios.
- `--optimizer-cpu-offload` moves optimizer state, with ZeRO-1 sharding or FSDP the optimizer states do not form memory bottleneck.
- `--fine-grained-activation-offloading` moves activations, but only when VPP + 1F1B is used which require warm-up steps will the activation memory form bottleneck.

The hardware makes moving the expert weights off the GPU practical: **GH200 has ~450 GB/s H2D bandwidth**, enough to stream expert weights from host memory and hide the transfer behind the expert GEMMs.

### Design

Keep routed expert weights in pinned host memory and stream them to GPU chunk-by-chunk during forward and backward, overlapping each chunk's H2D copy with the previous chunk's GEMM.
Image

Whether this is possible is a bandwidth-vs-FLOPs question. For a chunk of `C` experts with `M` tokens per expert, the time required for computation and H2D on GH200 is:

```
T_load = C · (H · 2h_e) · 2B / 450 GB/s
T_gemm = C · 2 · M · H · 2h_e / 989 TFLOP/s
overlap efficiency = T_gemm / T_load = M · 450 / 989e3
```

**The overlap depends only on tokens per expert.** At 100% MBU and MFU, break-even is `M ≈ 2200`; at a realistic 90% MBU / 60% MFU it is `M ≈ 1465`. It can be verified with an unit test (16 experts, M=1465, expert shape=(7168, 4096), bf16):

```
[H2D] FC1 2.232 ms | 421 GB/s | MBU 93.6%
FC2 1.119 ms | 420 GB/s | MBU 93.3%
[GroupedGEMM] FC1 2.263 ms | 608 TFLOPS | MFU 61.5%
FC2 1.285 ms | 536 TFLOPS | MFU 54.2%
Overlap efficiency: FC1 101.4% | FC2 114.8%
```

In a balanced-routing regime `M = mbs · seq · (EP/TP) · (N_a/N_e)`, so the requirement reduces to a constraint on the MoE model activation ratio. **This makes it straightforward to say in advance whether a given model and parallel layout will benefit from offloading expert weights**.

### Implementation

Proposed implementation (an `OffloadingExpertsMLP` expert module plus the supporting plumbing):

- **`OffloadingExpertsMLP`** — expert module whose weights live in pinned host memory, with a custom autograd function that pipelines per-chunk H2D staging against grouped GEMMs on separate CUDA streams. Double buffering by default, so chunk `i+1` loads while chunk `i` computes.
- **`_ParamAndGradBuffer`** — allocate the `param_data` buffer in pinned host memory for buffers made up entirely of offloaded expert params. Main gradients stay on GPU, deliberately: they are written every microbatch, and moving them would add H2D/D2H traffic on the critical path.
- **`DistributedOptimizer` / `Float16OptimizerWithFloat16Params`** — for a host-resident parameter, allocate the fp32 master weight on GPU so the optimizer step itself is unchanged.
- **`_ParamAndGradBucketGroup.start_param_sync()`** — NCCL cannot gather into host memory, so host-resident buckets are gathered through a bounded GPU staging buffer (stage the local shard up, all-gather, copy back).

Happy to open the PR against this issue. The branch contains unit tests covering numerical equivalence of the forward, dgrad, and wgrad. Would also welcome guidance on whether it might be more reasonable to integrate into TransformerEngine as the
`OffloadingExpertsMLP` is a separate auto grad function.

**Additional context**

Scope of the proposed PR is **expert weight offloading only**. Master weight and optimizer state offloading are natural follow-ups but are not included.

Current constraints:

- expert tensor parallelism must be 1 (`expert_tensor_parallel_size > 1` is rejected)
- requires `gradient_accumulation_fusion`
- requires TransformerEngine for the grouped GEMM
- the parameter all-gather for host-resident buckets is synchronous; it does not yet overlap with forward

Verified to work with virtual pipeline parallelism, `grad-reduce-overlap`, `param-gather-overlap`, activation recomputation.

An optional batched H2D copy path (`batched_h2d_async`) reduces per-chunk CPU launch overhead when available; the implementation falls back to per-expert `copy_` without it.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.