[Was PR1912][Dev] feat(moe): Fine-grained activation offloading

[lhb8125]

What does this PR do ?

PR for main branch

Memory capacity are more and more important with the rising of extreme sparse MoE models like DeepSeek-V3 and Qwen3-235B. Fine-grained recomputing reduces the memory footprint at the cost of extra recomputation, while offloading could utilize the host-device bandwidth to achieve nearly zero-overhead.

The current CPU offloading strategy from TE is a layer-level strategy, which offloads the activations in a granularity of the transformer layer, which is coarse-level and hard to highlight the most prominent activations.

Fine-grained Activation Offloading targets at offloading the activation at the granularity of specific modules, so that we can calibrate the amount of offloading activation to maximize the training throughput.

Design Doc

Compared with the current cpu offloading strategy provided by TE, this PR has several advantages:

• support PP=1/PP/VPP;
• support MoE models;
• manually specify offloading the modules with a large memory footprint;
• Work with fine-grained recomputation to reduce the total activations as much as possible;

How does fine-grained offloading work with fine-grained recomputing?

• For modules with minor perf overhead like layernorm or moe_act, use recomputing to reduce memory footprint;
• For other modules, use offloading to reduce memory footprint;
• Make sure the offloading/reloading could be overlapped with computing;

Benchmark

DeepSeek-V3-proxy on H100

Setup

• Layer parameters are same as DeepSeek-V3 model
• Layer number is cut off to 14 layers
• Replace the fisrt 3 dense layers with 3 moe layers
• TP1PP4EP16VPP1CP1-MBS1GBS512, bf16 training
• Offload expert_fc1, moe_act, act_norm and mlp_norm

Results

	Throughput (TFlops)	Max reserved memory (MB)
Baseline	321	74306
Offload expert_fc1,moe_act,layernorm	315	61046

DeepSeek-V3 on GB200 (from @hongxiaob )

Setup

• Same model structure with DeepSeek-V3 but no mtp
• TP1PP8EP32CP1VPP4-MBS1GBS2048, mxfp8
• Offload moe_act

Results

	Throughput (TFlops)	Max reserved memory (MB)
Baseline	945	169094
Offload moe_act	930	151054

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[lhb8125]

/ok to test e853ad8

[lhb8125]

/ok to test 2770d2b

[lhb8125]

/ok to test 118e00a

[lhb8125]

/ok to test ee670e0

[lhb8125]

/ok to test a74f12c

[lhb8125]

/ok to test f128365

[lhb8125]

/ok to test e2c36dd