Memory Overhead and TPS Degradation with Increased Expert Parallelism (EP)

[AnWang-AI]

Bug description

When scaling Expert Parallelism (EP) from 1 to 4 on a system with 8×H20 GPUs (use_grouped_mm=True, num_experts=128, top_k=8) on the Llama4 model, we observe:

• Memory usage​ increases from ​84.57 GiB (EP=1)​​ to ​91.51 GiB (EP=4)​.
• ​Throughput (TPS)​​ decreases from ​2,125 (EP=1)​​ to ​1,121 (EP=4)​.

Versions

llama4 model:

"3bx128e": TransformerModelArgs(
    dim=2048,
    n_layers=48,
    n_heads=32,
    n_kv_heads=4,
    multiple_of=8,
    ffn_dim_multiplier=1.2657,
    rope_theta=500000,
    max_seq_len=4096,
    moe_args=MoEArgs(
        num_experts=128,
        top_k=8,
        use_grouped_mm=True,
    ),
    interleave_moe_layer_step=1,
),

toml file:

[job]
dump_folder = "./outputs"
description = "Llama 4 Scout training"

[profiling]
enable_profiling = false
save_traces_folder = "profile_trace"
profile_freq = 100

[metrics]
log_freq = 1
enable_tensorboard = false
save_tb_folder = "tb"

[model]
name = "llama4"
flavor = "3bx128e"
hf_assets_path = "./assets/hf/Llama-4-Scout-17B-16E"
# converters = ["float8"]

[optimizer]
name = "AdamW"
lr = 4e-3
eps = 1e-15

[lr_scheduler]
warmup_steps = 600
min_lr_factor = 0.1

[training]
local_batch_size = 1
seq_len = 8192
max_norm = 1.0  # grad norm clipping
steps = 3000
dataset = "c4"

[parallelism]
data_parallel_replicate_degree = 1
data_parallel_shard_degree = -1
tensor_parallel_degree = 1
enable_async_tensor_parallel = false
pipeline_parallel_degree = 1
context_parallel_degree = 1
expert_parallel_degree = 4
expert_tensor_parallel_degree = 1

[checkpoint]
enable_checkpoint = false
folder = "checkpoint"
interval = 500
last_save_model_only = true
export_dtype = "float32"
async_mode = "disabled" # ["disabled", "async", "async_with_pinned_mem"]

[activation_checkpoint]
mode = "full" # ["none", "selective", "full"]

[compile]
enable=false
components = ["model", "loss"]

[float8]
enable_fsdp_float8_all_gather = false
precompute_float8_dynamic_scale_for_fsdp = false
filter_fqns = ["output", "router.gate"]

[tianyu-l]

Thanks for reporting! This has been on my plate but I didn't get to debug it.
Do you think you could help debug with the memory profiling tool?
https://github.com/pytorch/torchtitan/blob/main/docs/debugging.md#enable-memory-profiling

cc @rakkit

[rakkit]

I can reproduce it, and in short, it seems to be a "BUG" of AC.

Generally, we dont expect EP can always be fast; it really depends. In your example here, we should expect EP=2/4/8 can give roughly the same throughput, rather than EP vs. no EP.

---

Here is my test on 4xH100, model is the same, but I reduced the n_layers to 24.

EP=1

step: 50 loss: 3.3422 grad_norm: 1.0499 memory: 70.25GiB(75.46%) tps: 8,881 tflops: 131.68 mfu: 13.31%

EP=2

step: 50 loss: 2.9068 grad_norm: 1.0388 memory: 81.15GiB(87.17%) tps: 6,209 tflops: 92.06 mfu: 9.31%

EP=4

step: 50 loss: 2.9771 grad_norm: 1.1847 memory: 68.21GiB(73.27%) tps: 7,010 tflops: 103.93 mfu: 10.51%

$\uparrow$ the loss is also not algined

So, as we can see, EP=1 and EP=4 actaully use roughly the same memory; the real problem is when EP != world size.
[from memory profile, the memory usage during forward is the same for EP=1/2/4, so the problem comes from backward]

Here i take the example of layer-21 to explain what's wrong there.
(EP=4)

(EP=2)

From the profiler, some weird things are happening for EP=2, the pre_backward hook of layer-20 triged the gather of layer-21, which is wrong.

At the same time, the pre_backward hook of layer-22 triggered the gather of layer-21, which is expected but not exactly, for some reason here the layer-21 and layer-21-moe are seprately gathered.

For some reason, pre_backward of layer-21 itself also triggered the gather of itself. Therefore, in the recomputation+backward of layer-21, there are multiple copies of Layer-21's weights memory + multiple copies of layers 20 and 22. That is why we see EP=2 takes more memory.

The reason EP=4 is faster than EP=2 is that the ALL2ALL in the example here takes roughly about same time. But in the EP=4 case, grouped_gemm is faster in forward and backward. Plus, the EP=2 needs to do two slow all_gather_copy_out for each layer due to the bug.

[rakkit]

and a quick test run w/ ac, w/ compile

    "testx128e": TransformerModelArgs(
        dim=1024,
        n_layers=12,
        n_heads=32,
        n_kv_heads=4,
        multiple_of=8,
        ffn_dim_multiplier=1.2657,
        rope_theta=500000,
        max_seq_len=4096,
        moe_args=MoEArgs(num_experts=128, top_k=8),
        interleave_moe_layer_step=1,
    ),

(EP=1)
step: 50 loss: 2.8659 grad_norm: 0.5586 memory: 40.42GiB(43.42%) tps: 29,928 tflops: 101.55 mfu: 10.27%

(EP=2)

step: 50 loss: 2.8621 grad_norm: 0.6026 memory: 40.74GiB(43.76%) tps: 27,027 tflops: 91.71 mfu: 9.27%

(EP=4)
step: 50 loss: 2.8596 grad_norm: 0.6600 memory: 40.97GiB(44.01%) tps: 28,211 tflops: 95.73 mfu: 9.68%

The loss and memory usage are aligned.

[tianyu-l]

Thank you very much for the initial investigation @rakkit !

for some reason here the layer-21 and layer-21-moe are seprately gathered.

I think this is expected, as MoE experts and TransformerBlock are sharded separately by FSDP
https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/llama4/infra/parallelize.py#L337-L355

@weifengpy I wonder if this is because the forward prefetch hooks are replayed during backward recompute?
https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/llama4/infra/parallelize.py#L373-L395

cc
@soulitzer on abnormal behavior between FSDP and AC
@wwwjn on why EP uses more memory than pure FSDP

[rakkit]

So it seems we need to prevent the AC wrapped MoE module from triggering the fwd/bwd (prefetch) hook for next layer

[soulitzer]

Hmm I don't see any call to AC in the profiler trace. I think if AC is applied below FSDP, AC shouldn't wrap any FSDP logic, so I wouldn't expect AC to cause any prefetch hooks to replay.