Llama 4 issue tracking

[tianyu-l]

High priority

• Grouped MM @tianyu-l
  • [Bug] Potential bugs in "_grouped_mm" in Llama4 MoE codes #1237
  • with Activation Checkpointing
    • gets stuck after a couple of iterations
  • with AdamW
    • gets stuck after a couple of iterations
  • with torch.compile @bdhirsh
    • may need to register torch._grouped_mm (and the triton kernel for aligning indices)
    • basic compile support for grouped_mm pytorch#153384
    • compile: turn off fullgraph=True to support llama4 #1182
• auxiliary-loss-free load balancing ([llama4] add auxiliary-loss-free load balancing to MoE token routing #1114)
  • remove (persistent) buffers in checkpoint
  • currently it's using default stream to do blocking communication when DP degree > 1, need to assess if it's OK
• selective activation checkpointing
  • currently we are checkpointing every other matmul, which is not adapted to MoE router gate / torch._grouped_mm ops (potential solution)
  • solved in Add option for selective op AC to filter mm shapes based on fqn #1380

Not high priority for now

• for-loop implementation of MoE
  • with DTensor TP: sharding propagation overhead due to dynamic shapes
    • need to lift cache hit criteria in DTensor sharding prop
    • may be needed by Loss Parallel for per-sequence loss as well
  • with torch.compile: branching on “unbacked” symbolic ints
    • static padding of DTensor may solve this

Not llama4 specific

• PP with num_microbatches < num_stages @H-Huang
  • tracked in [PP] Hang when num_microbatches < stages for Interleaved1F1B pytorch#151269
• FlexAttention
  • not compatible with SAC
    • tracked in Flex Attention is incompatible with selective AC pytorch#147879
    • fixed in Remove flex+sac restriction #1408
  • with full AC, loss becoming nan starting from the 2nd iteration
    • Flex Attention is incompatible with selective AC pytorch#147879 (comment)
  • with PP @fegin
    • fixed in Pass the original input to all PP stages #1130

[Wodlfvllf]

I wanna contribute. Can you assign this to me ?

[tianyu-l]

@Wodlfvllf there's a list of issues / items, do you have preference on what to work on?

[lkhphuc]

[ ]
with AdamW

* gets stuck after a couple of iterations

Have you tried this with cuda12.8? I was able to run the debug model for 1000 iterations on torch nightly with cuda 12.8, using AdamW.

Edit: Same for Full AC and Selective AC, int or op.

with Activation Checkpointing
gets stuck after a couple of iterations

Successfully run 1000 iters debug model with AdamW+(F/S)AC with cuda 12.8.

[lessw2020]

A deeper dive on the torch group gemm hanging issue finds that the root cause is if an expert is assigned zero tokens...this will cause a hang in the backwards pass. (big thanks to @rciocoiu for the initial isolation work and tracking this core issue down).

The choice of Adam vs AdamW does matter b/c with Adam, an expert has to get to zero tokens assigned before hanging, vs with AdamW, it only has to get to the min_align_m setting (default is 16) and thus that happens sooner, matching users reporting that it hangs faster with AdamW.

have opened a Pytorch issue on this here: pytorch/pytorch#152668 and there's a min repro case setup to hang or not hang by simply having a zero offset (tokens assigned) expert.
Will investigate if we can safely pad out any expert assignments that are zero with min_aligned_tokens as the likely workaround/fix. This may be also fixable directly in the kernel but I think the fastest fix is to autopad.

[ngimel]

I already raised this issue with cutlass, they indeed don't support K=0 today, they promise a fix soon though. The best thing would be to indeed pad the group to minimum supported tokens (8 for bf16, 16 for fp8)

[lessw2020]

The PR to fix the hang for groupedMM (torch._grouped_mm) has landed and that should resolve the issues regarding hanging with groupedMM, regardless of optimizer.
I ran 2K iters with no issue, with the expert load balancing disabled to force lots of zero token experts (zero token experts was the core issue wrt to the hang).
@tianyu-l
#1166