How PyTorch Scaled a Near-1T MoE on Europe’s Public Supercomputer

Introduction

Training a mixture-of-experts model at close to one trillion parameters is normally associated with a dedicated AI cluster. The EUROPA consortium took a different route during the EuroHPC Frontier AI Grand Challenge: it used JUPITER, a publicly funded supercomputer shared by research projects, together with PyTorch and TorchTitan to scale pretraining to the frontier-model range.

The model used in the experiments has 951.9 billion total parameters, while approximately 32.8 billion are active for each token. Its configuration includes 74 layers, 384 routed experts with top-8 routing, a hidden size of 6,144, a 256,000-token vocabulary, and a 65,536-token pretraining sequence. The run used 384 NVIDIA GH200 GPUs across 96 nodes.

The real problem was the system balance

Each JUPITER Booster node contains four GH200 superchips. A GPU provides 96 GB of HBM3 and is linked to a 72-core Grace CPU with 120 GB of LPDDR5X memory through a high-bandwidth NVLink-C2C connection. Inter-node communication, however, goes through InfiniBand and is much slower than the local GPU-to-CPU link.

That imbalance made a conventional distributed setup insufficient. The team used 384-way FSDP, placed four-way expert parallelism within each node, and avoided pipeline parallelism. MoE token all-to-all traffic was kept on NVLink where possible. Long sequences also played a systems role: they increased the amount of computation available to hide communication latency.

Three optimizations that changed the economics

• FlexShard reduced Muon communication. Muon was used for almost all weight matrices. Its Newton-Schulz update operates on full matrices, which creates expensive gather-and-scatter operations when every matrix is sharded by FSDP. FlexShard assigned whole experts to ranks, allowing local orthogonalization without communication. Dense matrices were handled while expert computation continued. Running Newton-Schulz on the GPU instead of the Grace CPU reduced the optimizer step from about 41.7 seconds to 2.2 seconds and raised throughput from 441 to 597 tokens/s per GPU.
• Activation checkpoints moved to Grace. With a 65,536-token sequence, boundary activations alone require roughly 55 GiB per GPU. TorchTitan’s checkpointing support was extended to copy selected activations to Grace memory after the forward pass and bring them back before backward recomputation. At 32,768 tokens, full activation recomputation increased throughput from 90 to 370 tokens/s per GPU; offloading later raised it to 508 tokens/s per GPU.
• FSDP communication was overlapped with compute. FSDP can gather the next layer while the current layer is being processed. The project used unshard_async_op on the compute stream, freeing about 14 GiB of HBM per GPU and improving the handoff between communication and computation. Across 384 GH200 GPUs, the optimized setup reached roughly 1,300 tokens/s per GPU.

Why it matters

The significance of the project goes beyond a throughput figure. It shows that a public research supercomputer, even one not designed exclusively for AI training, can host frontier-scale pretraining when memory tiers, parallelism, optimizer behavior, and communication scheduling are designed together. PyTorch’s flexibility and TorchTitan’s distributed-training components made it possible to tune the stack for JUPITER and contribute several capabilities back to the open-source ecosystem.

The result should not be treated as a universal recipe for every one-trillion-parameter MoE. It depends on the GH200 topology, the available Grace memory, the sequence length, and the chosen parallelism layout. Its broader lesson is more useful: when inter-node bandwidth is the scarce resource, model architecture and systems software must be optimized as one integrated system.

Source: PyTorch Blog