Section 15: Optimization Attempts and Lessons

We also tested approaches that we did not adopt in the training recipes. They targeted saved activations, exposed communication, conversion kernels, host launch overhead, and the number of separate kernels in the MoE path. In several cases, reducing one cost increased another: offload exceeded the host-link bandwidth budget, overlap slowed concurrent kernels, and preserving the natural MXFP8 layout made Wgrad slower. Other attempts required more integration work than their prototypes justified. CUDA Graph replay needed pipeline-aware activation storage, while the megakernel needed a complete distributed training runtime.

In its supported BF16 rowwise configuration, two-batch overlap gave a small gain that did not justify the added coordination. Table 31 summarizes the decisions.

15.1 CPU Activation Offload Was Limited by PCIe

CPU activation offload exchanges GPU memory for host-link traffic. After a forward layer produces tensors that autograd will need later, the runtime copies selected tensors from device to host (D2H), releases their GPU storage, and copies them back from host to device (H2D) before the corresponding backward operation. Every offloaded byte therefore crosses the host link twice over one training step. Overlap can hide these copies only if the link drains tensors at least as quickly as the forward pass produces them and returns them within the backward prefetch window. Activation offload follows a line of work exemplified by vDNN (Rhu et al. 2016); the decision here concerns its service-rate balance on the measured Olmo workload.

Let \(A\) be the selected activation bytes produced during one layer interval, \(t_f\) the corresponding forward production interval, and \(t_b\) the usable backward prefetch interval. If \(B_{\mathrm{D2H}}\) and \(B_{\mathrm{H2D}}\) are the sustained pinned-memory bandwidths, necessary service-rate conditions are

\[ B_{\mathrm{D2H}} \geq \frac{A}{t_f}, \qquad B_{\mathrm{H2D}} \geq \frac{A}{t_b}. \tag{61} \]

When \(A/t_f\) exceeds \(B_{\mathrm{D2H}}\), each layer produces a larger copy backlog than PCIe can retire. Moving the copy earlier cannot remove that backlog, even on a separate CUDA stream; it only changes where the eventual wait appears. We write the required-to-available D2H bandwidth ratio as \(\rho=(A/t_f)/B_{\mathrm{D2H}}\).

Table 32 works the inequality for the routed block of the four-B300 benchmark of Section 13, using that profile’s GPU-projected layer timings and an activation inventory implied by the block’s autograd graph at its shape (\(T=16{,}384\) tokens per rank microbatch, \(d=h=4096\), top-4 routing, one shared expert, BF16). The inventory counts valid routed rows only. As an aggregate consistency check, the inventory summed over the seven routed blocks, the dense block, and the output logits is within a few percent of the measured 40.8 GiB transient peak. The logits are not offload candidates, so the table’s whole-microbatch traffic excludes them.

A routed block produces its 3.96 GB of candidate activations in 10.3 ms, a 384 GB/s stream, against a 63 GB/s link, so \(\rho\approx6\) with the link entirely to itself and above 7 at the copy rate actually measured on a host of the same link generation. Even spread over the entire 295 ms microbatch, the traffic is still 104 GB/s in each direction, and the measured copies lose a further 15–25% of their rate when both directions run at once. A node’s eight GPUs offloading together would push about 3 TB/s into host memory, more than a two-socket DDR5 host can absorb regardless of link generation; a PCIe Gen6 host link would halve \(\rho\) but leave both the per-GPU forward-rate mismatch and the host-memory limit.

Faster compute further narrows the offload window unless host bandwidth increases with it. The activation-production rate in Table 32 is Blackwell’s, since the block timings come from B300 GPUs, while the link is the same PCIe Gen5 x16 that served Hopper hosts. At Hopper’s roughly \(2.3\times\) lower BF16 peak (Table 21), scaling the block time inversely with peak arithmetic throughput would put \(\rho\) near 3, still above the available bandwidth. This is an estimate, not a Hopper measurement.

A selective policy can offload at most approximately a fraction \(f\leq 1/\rho\) of the candidate bytes without building a sustained D2H backlog. At \(\rho\approx6\) that is about one sixth, roughly 0.65 GB per block or about 5 GB over a microbatch: less than a fifth of what per-block recompute removes at the same operating point (Figure 62), for an additional subsystem with its own failure modes. That reduction is too small to justify the subsystem in the current recipe.

15.2 Wave Overlap Was Not a Reliable Win

Wave scheduling divides one MoE invocation into several route groups. While experts compute one wave, another CUDA stream can dispatch a later wave or combine an earlier one. This creates an opportunity to hide communication, but it also replaces one larger grouped GEMM with several smaller grouped GEMMs and runs communication beside compute on the same GPU. Fine-grained overlap systems such as COMET and Lancet show why this direction is attractive (Zhang et al. 2025; C. Jiang et al. 2024). Its benefit depends on how much both communication and compute slow down when run together.

The controlled measurements in Section 5.9 show why stream concurrency alone did not guarantee a win. In an isolated two-node EP-MP\({}=16\) experiment, rowwise NVSHMEM dispatch or combine slowed the concurrent expert GEMM by about 23–27% across the tested communication-block sweep. The communication region also stretched. A wide 128-block launch still shortened the pair’s makespan by 22–25% relative to sequential execution, but even a one-block communication kernel did not protect the GEMM; it mainly made communication take longer. A DeepEP v2 sweep found the wave path slower than the no-wave path at the tested communication settings of 4, 6, 16, 32, and 64 SMs on its own two-node EP-MP\({}=16\) shape; without a wave-sequential control, that loss cannot be split between smaller GEMMs and concurrent-kernel interference.

The measurements do not identify the interference mechanism. The rowwise path executes GPU warps that issue one-sided transfers and stream barriers, so candidate causes include SM issue pressure, NVSHMEM progress work, L2 or memory-fabric contention, and barrier overhead. Wave partitioning independently reduces GEMM efficiency. Other measured shapes include both wins and losses, so the result does not show that communication–compute overlap is never useful. Concurrent kernels can still take longer to finish (Lee et al. 2025). For the tested setup, phase-sequential execution is the appropriate reference; the current evidence does not establish that overlap improves the complete forward and backward path.

15.3 Two-Batch Overlap

Two-batch overlap (TBO) applies the overlap idea at a larger granularity. It splits the local microbatch into two half-batch lanes and interleaves their work across adjacent MoE blocks (Figure 63). While one lane runs attention or expert GEMMs, rowwise dispatch or combine for the other lane can remain pending on the EP communication stream. We implemented this BF16 rowwise schedule and tried it during development, but decided not to use it in our current training recipes. The implementation requires BF16 rowwise NVSHMEM EP and rejects rowwise MXFP8, per-block recompute, context parallelism, and the DeepEP v2 path, so the MXFP8, per-layer-recompute Ultra recipes of Section 14 could not use it; the adoption decision concerns the BF16 rowwise recipes. DeepEP documents a contrasting legacy TBO path designed around low-interference communication (DeepSeek-AI 2026).

Figure 63Interactive · conceptual schedule, steps are not measured durations

Step through the schedule: whose communication runs beside whose compute?

Step

3 / 10

Follow

Both lanesLane 0Lane 1

Step 3 · compute Experts 0 · EP comm. Dispatch 1

Step 3 of 10 · Block 0

Compute streamExperts 0lane 0 · block 0 · after Dispatch 0

EP communication streamDispatch 1lane 1 · block 0 · after Attention 1

Dispatch 1 for lane 1 can remain pending on the EP communication stream while lane 0 runs expert GEMMs.

8 of 10 steps pair one lane’s communication with the other lane’s compute.

• Lane 0 (blue in the report)
• Lane 1 (green)
• Data dependency within a lane

Redrawn from the report’s diagram with the same ops, lanes and arrows. Each step (a column, or a row on narrow screens) is one position in the diagram’s left-to-right order, and box sizes carry no time. The measured result for this schedule is Figure 64.

Figure 63. Two-batch overlap interleaves two half-microbatch lanes across adjacent MoE blocks. Blue and green denote lanes 0 and 1. The schedule interleaves attention, expert computation, and expert-parallel dispatch and combine so that communication from one lane can run alongside computation from the other. The arrows show data dependencies; the visual overlap alone does not imply a wall-time speedup.

Figure 64 reports a matched end-to-end comparison using the four-GPU configuration and timing protocol of Section 13: TBO improves throughput by one percent and trims peak active memory by about 1.5 GiB. The figure also draws what ideal overlap would gain. EP transport occupies a seventh of the normal run’s GPU span, so even a scheme that hid every millisecond of it without slowing the compute beside it would gain at most seventeen percent, or about fifteen percent after excluding the transport that already overlaps compute in normal training; TBO achieves a one-percent gain. This is a small positive result for one operating point, not evidence that TBO is generally faster.

Figure 64Interactive · measured, one operating point · four B300 GPUs, rowwise EP

One percent measured against seventeen percent possible: where does the overlap go?

View

Change vs normalChangeTen-step windowsWindowsInside the stepProfile

• TBO, mean of three
• Each ten-step window
• Ideal overlap

Normal rowwise training54.68Ktokens/s/GPU · 103.03 GiB peak active

Two-batch overlap55.22Ktokens/s/GPU · 101.57 GiB peak active

Measured gain · ideal+1.00% of +17.4%three-repetition means; ideal hides all EP transport

Plotted from the campaign’s CSVs, the same rows as the original figure (steady_state_repetitions.csv; profile rows from mode_summary.csv, rank 0, mean of three captured train_batch ranges). Percent changes and the ideal line are computed as in the report’s plotting script; the configuration and timing protocol are those of Figure 62. One operating point; not evidence that TBO is generally faster.

Figure 64. Two-batch overlap improves throughput by 1%, compared with a 17% ideal-overlap estimate. Filled bars show TBO’s change of the three-repetition mean relative to normal rowwise training, with the individual ten-step windows as dots. The dashed outline is an ideal-overlap upper bound: EP transport occupies 14.8% of the normal run’s GPU span on rank 0, so hiding all of it with no slowdown of the concurrent compute would raise throughput by at most 17%; about 40 ms of that transport already overlaps compute in normal training, so a tighter bound is about 15%. Normal training sustains 54.68K tokens/s/GPU at 103.03 GiB peak active memory; TBO sustains 55.22K at 101.57 GiB. Same four-GPU B300 configuration, matched model, batch, routing, and timing protocol as Figure 62.

Concurrent kernels take longer under TBO. Olmo’s rowwise NVSHMEM communication is not an SM-free direct memory access (DMA) transfer. Its route-level PUT and GET kernels use SM scheduling and memory-system resources also needed by expert GEMMs. In the profile, TBO overlaps 503.6 ms of communication and compute, equal to 72.2% of TBO’s own EP transport interval. At the same time, the summed EP-transport duration grows by 98.9%, the GEMM duration by 20.4%, and attention duration by 16.5%. These sums combine contention with the lower efficiency of two smaller lane kernels and can overlap with one another; they are not an additive wall-time decomposition. Taking the union of the intervals, the projected GPU span falls by 0.9%, consistent with the independent 1.00% throughput gain. Thus contention is real, but it does not quite erase the overlap benefit in this configuration. NCCL All-to-All can have a different contention profile, but it also retains the packing and unpacking permutations that rowwise EP removes.

Splitting the microbatch roughly doubles the host scheduling work. Each half-batch lane has its own attention, dispatch, expert, combine, and synchronization stages. The profile records 89.8% more kernel launches and 126.4% more CUDA API calls than normal execution. Kernel-idle gaps of at least 10 \(\mu\)s also become 2.09\(\times\) as frequent, although total idle time within the projected range rises by only 9.7 ms. The CPU remains far enough ahead at this rank-local microbatch size, so launch work does not erase the small win. At smaller microbatches or on a faster kernel path, the same extra submission work moves the system closer to a launch-limited regime. CUDA Graph replay could lower this cost, but the current stack has no stable graph path that captures this schedule end to end.

TBO adds coordination to the pipeline schedule. The implementation carries two live lanes, separate communication buffers, and stream-ordering dependencies across neighboring blocks. Pipeline parallelism (PP) already imposes its own microbatch order and point-to-point dependencies, so combining the two schedules requires additional coordination rather than a local MoE-layer change. The code accepts TBO with PP, but we measured only a no-PP configuration, and its one-percent gain did not justify adopting the additional schedule-dependent coordination. A different microbatch size, transport, interconnect, or pipeline schedule can change the balance between hidden communication and the costs of contention, smaller kernels, and extra launches.

15.4 The MegaMoE Kernel Became a Runtime Project

To control this interference directly, a persistent MoE kernel can assign explicit roles to communication, scheduling, and tensor-core work. MegaMoE, our BF16 prototype of such a persistent MoE megakernel in the style of DeepGEMM’s Mega MoE (Section 5.9), explored this direction by bringing route dispatch, persistent expert scheduling, the two expert linear layers, the activation, and weighted combine under one Olmo-controlled CUDA implementation.

The prototype needed a complete distributed execution path. A usable training backend still needed a production tensor-core and data-movement scheduler, cross-rank signaling, reusable peer workspaces, returned expert outputs, backward for activations and both expert weights, and inter-node transport. These pieces are coupled: a fast local forward body cannot replace the rowwise path if backward or remote movement returns to a separate, incompatible schedule. The prototype grew across Python, C++, CUDA, and distributed-runtime code before it had a production training backward.

We removed that path and kept rowwise EP as the end-to-end baseline. The prototype’s tensor-core body ran as a serialized debug version, so its timing does not estimate the performance of the intended persistent implementation.

15.5 Wgrad Needed a Different Scale Layout

The rowwise MXFP8 forward path naturally represents activations and output gradients with scales \(s_A[m,k_b]\) and \(s_{dY}[m,n_b]\), where \(m\) is a routed-token row and \(k_b\) and \(n_b\) are feature blocks. Expert weight-gradient computation is

\[ dW_e = A_e^{\mathsf T} dY_e, \tag{62} \]

so its reduction dimension is the token index \(m\). The natural scale factors therefore change along the Wgrad reduction. Blackwell’s native block-scaled matrix instructions expect scale factors arranged according to their supported GEMM reduction layout; the already-available rowwise representation is not that layout.

The exact-natural prototype handled the mismatch by dequantizing tiles into BF16 shared memory and using BF16 tensor-core arithmetic. Table 33 shows that it was substantially slower than two global dequantization kernels followed by the tuned BF16 grouped Wgrad. Despite doing a separate global conversion, that path benefits from the faster GEMM implementation.

The current path therefore dequantizes the saved MXFP8 activation and output gradient, then calls BF16 grouped GEMM for Wgrad. This does not imply that native MXFP8 Wgrad is impossible. Such a path would deliberately requantize or repack the operands into a Wgrad-oriented scale layout before calling a native block-scaled grouped GEMM. Its cost includes that conversion as part of the complete backward region. The prototype measurements favor the converted layout and tuned BF16 GEMM over directly consuming the forward layout.

15.6 Static CUDA Graphs Need Pipeline-Aware Storage

CUDA Graph replay targets a different cost: repeated host submission. As Section 8 explains, replay can be attractive when the GPU executes small regions faster than Python and the CUDA runtime can submit them. The static PyTorch capture/replay path considered here records a fixed launch sequence with stable tensor addresses, then replays that sequence with much less host work. CUDA also provides graph-update APIs that can change node parameters, including memory addresses; those mechanisms are outside this static prototype (NVIDIA 2026g; PyTorch Contributors 2026b).

Host submission can dominate a short operation. In a standalone single-GPU grouped-linear microbenchmark, the captured device sequence for a short MXFP8 expert operation runs more than twice as fast as its BF16 baseline while under ordinary eager submission the same operation is slower than BF16 (Table 48 in Appendix B.6).

Stable graph addresses conflict with live pipeline activations. A pipeline stage can execute several microbatch forward passes before their corresponding backward passes. If every replay of one layer uses the same activation slot, a later forward overwrites data that an earlier backward will read. A pipeline-compatible design can assign a graph runner and static activation slot to every simultaneously live \((\text{layer},\text{microbatch})\) interval, or capture a larger fixed schedule containing the forward–backward reuse order. Separate slots increase graph pools and retained activation memory. Copying live tensors into graph-owned staging storage stabilizes their addresses but adds device-to-device copies. Staging, graph-pool memory, activation retention, and capture cost all count against replay savings. Figure 65 illustrates the separate-slot design with three in-flight microbatches.

Figure 65Interactive · conceptual, no measured quantities

Replay one layer through the schedule: when does a shared slot break?

Microbatches in flight

1 (a)23 (b)4

GPU order

4 / 8

B0 · shared slot overwritten · per-microbatch slot safe

Step 4 of 8 · B0, backward of microbatch 0

One shared slotOverwrittenB0 finds microbatch 2’s context: microbatch 0’s was overwritten by F1.1 graph · 1 wrong backward read so far

One slot per in-flight microbatchSafeB0 reads slot A: microbatch 0’s context, intact. The slot can be recycled for microbatch 3.3 graphs and static buffer sets · 0 wrong reads

With 3 microbatches in flight, a shared slot is safe only if nothing else is live; per-microbatch slots cost 3× the graphs and static buffers.

• Forward
• Backward
• Saved context held in a slot
• Overwritten before its backward

After the report’s two panels: 1 in flight is panel (a), 3 is panel (b); 2 and 4 apply the same 1F1B rule. Slot counts are illustrative; the report measures no memory or timing for its own static prototype. For Megatron-Core’s separate-slot implementation the report cites “about a ten percent end-to-end speedup at roughly 7 GB of extra memory on a DeepSeek-V3 configuration” (Yan et al. 2026); see Section 15.6.

Figure 65. Under a pipeline schedule, static per-layer replay needs one saved-context slot per in-flight microbatch. (a) Without a pipeline the backward of each microbatch consumes the layer’s saved context before the next forward overwrites it, so one graph per layer serves every microbatch. (b) With three microbatches in flight, a shared slot is overwritten by F1 and F2 before B0 reads it; three slots, each with its own graph and static buffers, are recycled as each microbatch’s backward for the layer completes. Conceptual illustration after the design described for Megatron-Core (Yan et al. 2026); no measured quantities.

The prototype could capture a narrow rowwise forward core in non-pipeline and smoke-test configurations, but it did not establish a beneficial complete PP training path. Marking more tensors as static did not solve the scheduling problem, because the number and reuse order of safe slots depends on the pipeline schedule. We therefore retired the graph-specific backend rather than maintain a second rowwise implementation with an incomplete memory trade-off.

Megatron-Core implements the separate-slot design: one graph per layer and microbatch under pipeline parallelism, graphs captured in execution order so that they share one memory pool, and static buffers reused by the pipeline order, for which it reports about a ten percent end-to-end speedup at roughly 7 GB of extra memory on a DeepSeek-V3 configuration (Yan et al. 2026). This result includes the schedule-dependent storage that our isolated forward capture lacked.

15.7 Local Savings Must Survive the Training Schedule

These decisions favor the complete training path over savings in an isolated region. The offload budget exceeds host-link bandwidth; wave overlap must pay for smaller GEMMs and interference; and TBO’s small gain did not justify its additional schedule coordination. For Wgrad, global conversion followed by a tuned GEMM beat direct consumption of the forward layout. The graph and megakernel prototypes lacked the storage or distributed execution needed to carry their local operations through training.