What does this PR do ?

PR for dev 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

DeepSeek-V3 on GB200 (from @hongxiaob )

Setup

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

Results