TorchTitan在GH200上训练1T MoE的6条工程优化经验

Hi everyone!

Over the last few months we've been in the trenches with the TorchTitan and PyTorch maintainers, trying to find out how far GH200 (the GPUs inside the JUPITER supercomputer) can be pushed. Short answer: far enough to train a ~1T-parameter MoE at ~1,300 tokens/s per GPU on 384 GPUs. That's 14X where we started, and our starting point was, let's say, humbling.

Along the way we collected a handful of lessons. If you've ever trained a big model, or just stared at a CUDA out of memory, some of this will feel familiar.

Here's the tour. If you want the full thing with all the plots, it's on the Domyn blog, and there's a condensed version on the PyTorch blog if you're short on coffee.

Hope you like it!

The setup: FSDP + EP, and why I'm not a PP fan

A 1T MoE doesn't fit on one GPU (shocking, I know), so the first question is how to shard it. This choice ends up dictating basically every optimization that follows, so it pays to think about it. The usual suspects are:

• Pipeline Parallelism (each GPU gets a chunk of layers),
• FSDP (each GPU holds a slice of every weight and gathers the rest just in time),
• and Expert Parallelism (each GPU owns a subset of experts and tokens travel to them).

We skipped PP. To hide its bubbles you need lots of microbatches, and lots of microbatches means long steps. Long steps mean that when something breaks at scale, and something always breaks at scale, you wait forever to find out. I'd rather iterate fast.

So: 384-way FSDP plus 4-way EP, with EP kept strictly inside a node. The token all-to-all for the experts is comunication that is difficult to hide so it needs to happen fast. On the other hand, comms. on FSDP are already hidden: while it computes layer i, it is already gathering i+1.

Why does that matter so much? On JUPITER, InfiniBand gives you ~25 GB/s per GPU. That's 6X slower than NVLink and 36X slower than the Grace-Hopper link. Our baseline step took 84 seconds, and 68 of them were spent waiting on the network. The GPUs were mostly getting paid to sit around. The rest of this post is basically us trying to put them back to work.

Lesson 1: If your GPU is idle, make it work for its memory

FSDP moves weights, and the weights gathered per layer don't care how many tokens you feed through them. More tokens per sequence means more compute per byte on the wire, which is exactly what we wanted. The problem is that longer sequences need more memory, so we went shopping for it:

• Full activation checkpointing. Recomputing every layer in the backward sounds wasteful, until you remember the GPU was idle anyway. Free compute is free. This alone gave us 4X.
• Offloading checkpoint activations to Grace over the 900 GB/s C2C link, and prefetching them back one layer before recompute. Grace was right there, so we used it.
• Muon instead of Adam. One optimizer state per parameter instead of two, and that's what finally got us to 64k-token sequences.

Lession 2: If you don’t shard it, you don’t have to gather it

Quick refresher: Muon orthogonalizes each update with a few Newton-Schulz iterations, and Newton-Schulz needs the whole matrix. FSDP, meanwhile, has carefully cut every matrix into 384 slices and scattered them across the job. The naive fix is to gather each matrix on one rank, run Newton-Schulz, and scatter the result back, which means more communication, and more communication is the one thing we absolutely did not need.

The trick is to shard smarter. With FlexShard (already in TorchTitan), you choose which rank owns which parameters, so we sharded the experts by whole expert. Every GPU owns complete expert matrices and can orthogonalize them locally with zero communication. The dense matrices (attention projections and friends) still need gathering, but they’re in transit while each rank is busy with its own experts, so that traffic hides behind compute.

There was one catch. We offload the optimizer to Grace, which meant Newton-Schulz was running on the CPU too. Grace is a lovely CPU, but asking it to grind through matmuls is like asking a librarian to move furniture: technically possible, deeply inefficient. The optimizer step took 42 seconds. So we contributed a fix to TorchTitan that runs Newton-Schulz on the GPU even when the parameters live on the host. The optimizer step dropped to ~2 s, and throughput went up 35%.

Lesson 3: You're only as fast as your slowest rank

384 GPUs in lockstep is a bit like a group hike: everyone waits for whoever's at the back. We had two stragglers:

• Attention cost. A rank that gets a single 60k-token document does far more attention work than one that gets a hundred short ones, even though both see the same number of tokens. We pre-pack rows, bucket them by attention cost, and build each step from rows of similar cost. The worst rank went from 7.9X the average to 1.15X, for +13% throughput.
• The router. Group-limited routing (two experts picked from each node's group) makes expert load perfectly even, for +7%. The catch is that you're telling the router where it's allowed to send tokens, and early runs suggest that can hurt training. Jury's still out on this one boys.

Lesson 4: Accumulate in fp32, ship in bf16

Gradients come out of the backward in bf16. Summing hundreds of them in bf16, though, is a great way to slowly lose precision, so the standard FSDP mixed-precision setup does the reduction in fp32. That’s the right call. The problem is how it gets there.

NCCL’s reduce-scatter runs as a ring. Each rank receives a partial sum from its neighbor, adds its own contribution, and passes the result along. The thing travelling around the ring is a running sum, so for the accumulation to happen in fp32, the data on the wire has to be fp32 too. FSDP therefore upcasts the bf16 gradients to fp32 before the collective, and every byte of gradient crosses InfiniBand at twice the size it needs to be. On a cluster where InfiniBand is already the bottleneck, that hurts.

But only the summation needs fp32, not the transfer. So we swapped the reduce-scatter for an all-to-all on the expert gradients. Each rank sends every other rank the raw bf16 slice that rank owns, no partial sums involved. The owner receives all the slices, upcasts them, and sums them locally in fp32. Same math, same precision where it matters, half the bytes on the wire.

The result matches the fp32 reduce-scatter up to summation order.

Lesson 5: When in doubt, send fewer bytes

Eventually longer sequences stopped helping (more on why in a second), so we went after communication volume directly:

• FP8 expert all-gather with blockwise scales. Master weights stay in full precision, only the transfer gets squeezed: +18%.
• Node-local resharding of dense weights after forward, so the backward re-gather runs over NVLink instead of fighting gradients for InfiniBand.

And none of it moved the loss in our comparison runs, which was a relief.

Lesson 6: Your OOM might be lying to you

This is my favorite one. After a software upgrade, every single rank started hitting allocator retries at ~70% HBM utilization, with 25–30 GiB sitting free and cached. That's not a memory shortage, that's a memory hostage situation.

The culprit: the PyTorch caching allocator keeps a separate pool per CUDA stream, and FSDP allocates buffers on its side streams. Our shiny FP8 and bf16 collectives allocated their temporaries there too, so ~21 GiB ended up parked on side streams where the compute stream couldn't touch it.

The fix was to move those allocations onto the compute stream (via FSDP2's private unshard_async_op flag, which is now on its way to becoming a public API) and release gradients a bit earlier. Retries went from 384 ranks to zero, peak reserved memory dropped by 13.5 GiB, and throughput jumped 21%, while sending exactly the same bytes.

Wrapping up