perf-cuda-graphs

To reduce GPU launch overhead during distributed training, capture and replay CUDA graphs in Megatron Bridge to minimize host-driver overhead for static-shape pretrain and fine-tune workloads. This applies to dense models and dropless MoE models.

Yes, CUDA graph capture supports dropless MoE training in Transformer Engine by targeting the attn, moe_router, and moe_preprocess scopes. You must use the local full_iteration capture or TE-scoped implementations while enforcing mandatory MoE scope composition constraints.

To enable TE-scoped CUDA graphs in Megatron Bridge, set model.cuda_graph_impl to transformer_engine and model.cuda_graph_scope to attn,moe_router,moe_preprocess. You must also enable model.use_te_rng_tracker and rng.te_rng_tracker to satisfy strict TE RNG tracker requirements.

You should not use CUDA graphs for MoE training when dealing with packed-sequence incompatibility or dynamic shapes. CUDA graphs require static shapes and specific NaN-check constraints for full_iteration capture, making dynamic workloads unsuitable.

The best way to validate CUDA graph performance tuning is to select the narrowest useful scope first, then compare replay timing against eager execution after warmup. Only widen the scope or combine with overlap features after confirming the initial performance gains.