Megatron

Large MoE Performance: The Three Walls After Sparsity Sparsity made MoE cheap on paper. At production scale it made training harder than dense: total parameters grow with E, per-token FLOPs grow with k, and the gap between those two numbers is exactly where systems break. The useful framing is not “optimize the MoE kernel.” It is the one NVIDIA’s Megatron-Core MoE report uses (arXiv:2603.07685): Memory, Communication, and Compute Efficiency are three coupled walls. Push on one and pressure shows up in another. ByteDance’s MegaScale-MoE (arXiv:2505.11432) proves the same thesis from the other direction — on 1,440 Hoppers, communication was ~44% of forward time before their redesign, and fixing parallelism + overlap delivered 1.88× over Megatron-LM. ...

Hiding Tensor-Parallel Collectives: AG/RS Overlap in Megatron Once Megatron SP is enabled, tensor-parallel communication often appears as AllGather and ReduceScatter rather than a single AllReduce. That is a memory win because activations can stay sequence-sharded between tensor-parallel regions. It also creates a scheduling question. Can the collectives be hidden under GEMM work? That question sits directly on top of Megatron’s tensor/sequence-parallel training path (arXiv:2104.04473, arXiv:2205.05198). This post is about that scheduling question. It complements Megatron tensor parallelism and Megatron SP. The goal is not to memorize every flag in Megatron or Transformer Engine. The goal is to recognize where the dependency graph allows overlap and where it does not. ...

Sequence Parallelism IV: Megatron Context Parallel and Load-Balanced Rings Megatron Context Parallel, usually shortened to CP, brings long-context attention into Megatron’s hybrid-parallel world. It keeps the familiar tensor, pipeline, and data-parallel axes, then adds a context axis that shards the sequence. Inside that context group, attention behaves like a ring, borrowing the same blockwise K/V circulation idea as Ring Attention (arXiv:2310.01889) while staying compatible with Megatron-style tensor and pipeline parallelism (arXiv:2104.04473). ...

MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation DeepSpeed-Megatron MoE is easy to misread if you only inspect one initialization function. Megatron builds the usual TP, PP, and DP topology. DeepSpeed then wraps the model, finds MoE modules, and gives those modules expert-parallel groups. The topology is split across the training framework and the MoE layer implementation. This post connects the pieces: deepspeed.initialize(), EP groups, expert-DP groups, MoELayer, and the difference between a DeepSpeed MoE wrapper and a Megatron-integrated SwitchMLP-style module. For routing, capacity, and All-to-All fundamentals, start with MoE Parallelism Principles. ...

Sequence Parallelism I: Megatron SP Cuts Activation Memory Along the Sequence Tensor parallelism is usually introduced as a way to make matrix multiplications fit. That is true, but it hides a second problem. After the weights are split, many activations are still replicated on every tensor-parallel rank. For short contexts this is tolerable. For long contexts it becomes one of the reasons training throughput collapses into activation checkpointing. Megatron sequence parallelism, usually shortened to Megatron SP, is a targeted fix from Reducing Activation Recomputation in Large Transformer Models (arXiv:2205.05198). It does not replace tensor parallelism. It keeps Megatron’s column-parallel and row-parallel linear layers, then shards the sequence-local regions that tensor parallelism had left replicated. The trick is small enough to miss and important enough to change the memory budget of a whole Transformer block. ...