EUROPA联盟在JUPITER超算上训练1T参数MoE模型的技术复盘

TL;DR

Scaling pretraining for a ~1-trillion-parameter Mixture-of-Experts (MoE) model on shared, public infrastructure was previously considered unfeasible due to unique architectural and interconnect bottlenecks. However, by leveraging the fine-grained flexibility of PyTorch and TorchTitan, the EUROPA consortium successfully unlocked frontier-scale training on JUPITER, boosting throughput from 90 to ~1,300 tokens/sec per GPU across 384 NVIDIA GH200 GPUs per DP replica (96 nodes) as part of the EuroHPC Frontier AI Grand Challenge.

The EuroHPC Frontier AI Grand Challenge was designed to prove that sovereign, open-source AI could compete at frontier levels using shared public resources. As part of that process, we – as the EUROPA consortium, formed by Domyn and Fraunhofer IAIS and IIS – were selected to an open European Mixture-of-Experts (MoE) model of at least 400B parameters, support for all 24 official EU languages, and strict compliance with the EU AI Act.

In this post, we describe how we scaled pretraining to 1T MoE on JUPITER, the EuroHPC exascale supercomputer hosted at Forschungszentrum Jülich, proving that public infrastructure can indeed sustain frontier scale.

The table lists the configuration used for the experiments in the post.

Models of this size are usually trained on clusters that their developers operate or rent for exclusive use. JUPITER is instead publicly funded research infrastructure shared by hundreds of projects. It was also designed to support scientific research at large, which led to architectural choices that are somewhat atypical in the AI space, as we’ll describe in detail later in the post. As such, reaching the 1T parameter scale while maintaining viable throughput requires careful optimizations that go beyond standard distributed techniques.

In order to have access to the necessary level of flexibility to perform fine-grained optimizations to the pre-training setup, we chose to go with PyTorch and TorchTitan. Under these conditions, we raised throughput from 90 to ~1,300 tokens/s per GPU on 384 GH200 GPUs (96 nodes). We want to share the main milestones that get us there, each of which builds on TorchTitan’s design and led to contributions back to TorchTitan and PyTorch:

• FlexShard placement for the Muon optimizer, which removes most of the optimizer’s communication.
• Activation-checkpoint offloading to the Grace CPU.
• Running FSDP’s all-gather setup on the compute stream (unshard_async_op), which frees ~14 GiB of HBM per GPU.

The full write-up covers the remaining changes in detail, with interactive figures and profiler traces.

The machine and its constraints

Each JUPITER Booster node has four NVIDIA GH200 Grace Hopper Superchips. Each superchip pairs a Hopper GPU (96 GB HBM3) with a 72-core Grace CPU (120 GB LPDDR5X), connected by a 900 GB/s NVLink-C2C link. GPUs in a node are connected by NVLink 4 at 150 GB/s per direction, and each node has four InfiniBand NDR200 adapters, about 25 GB/s per GPU.

Figure 2. Bandwidths within a JUPITER node, drawn to scale. All inter-node traffic leaves through the thin InfiniBand links.

In addition to the accelerators, two key parts of the system are also relevant for the design:

• Inter-node bandwidth. InfiniBand is 36 times slower than the link between a GPU and its own CPU, and distributing a 1T parameter model over hundreds of accelerators requires crossing that boundary continuously during training. Traffic over InfiniBand must therefore aggressively overlap with computation.
• Host memory capacity. GH200 is available with up to 480 GB of LPDDR5X per Grace, but JUPITER’s compute nodes actually use the 120 GB configuration (as opposed to login nodes which sport 480GB), which favors bandwidth over capacity. After the kernel’s share, about 108 GiB per GPU is usable, and we found that in practice some nodes start jobs with much less. This limits how much state can be moved out of the HBM.

These constraints dictate the parallelism layout: 384-way FSDP with 4-way expert parallelism inside each node and no pipeline parallelism. The MoE token all-to-all is hard to overlap, so we kept it on NVLink. Without pipeline parallelism steps tend to be shorter, which has the advantage of providing fast feedback at scale, and fixed costs are amortized via a large batch of 384 sequences of 65,536 tokens (~25M tokens per step). FSDP already provides us with the overlap of the gather of layer n+1 with the compute of layer n, so most of the remaining work goes into making the overlap effective. We achieve this across three dimensions: moving to longer sequences to compute for longer, reducing the bytes transmitted over InfiniBand, and leveraging Grace memory to provide enough total capacity to hold the model, the gradients and the optimizer state.

FlexShard: Muon without the communication

We use Muon for every weight matrix except the router and embeddings (about 99.5% of the parameters) and AdamW for the rest. Muon requires one fp32 state per parameter, while Adam requires two. Removing the second state is instrumental in making a 65,536-token sequence fit in memory.

Muon orthogonalizes each update with a few Newton-Schulz iterations, which operate on the whole matrix. Because FSDP splits every matrix into 384 shards, a simple implementation gathers each matrix on one rank, runs Newton-Schulz there and scatters the result back. On our interconnect, this communication is too expensive.

FlexShard, now part of TorchTitan, lets us control which rank owns which parameters. We shard the experts by whole expert, so every rank owns complete expert matrices and runs their Newton-Schulz iterations with no communication at all. The dense matrices, such as the attention projections, are still collected on an owning rank, and each rank processes its own experts while the dense matrices are in transit, so that communication is overlapped.

With FSDP’s CPU offload, the optimizer runs on Grace, where the parameters are stored, and Newton-Schulz on 72 Arm cores takes ~35 s per step. We contributed support to TorchTitan for running Newton-Schulz on the GPU while parameters stay offloaded. The optimizer step went from 41.7 s to 2.2 s, and throughput rose from 441 to 597 tokens/s per GPU (+35%).

Figure 3. Rank 0’s profiled step from 80 s onward. Top: Newton-Schulz on Grace, with the GPU idle for ~34 s after the backward pass while the CPU computes. Bottom: Newton-Schulz on Hopper, where the optimizer takes a small fraction of the step.

Activation offloading to Grace

The volume of weights gathered per layer does not depend on the number of tokens in our setup, so a longer sequence increases compute time while not impacting communication. We first stored activations only at layer boundaries and recomputed everything inside a layer during the backward pass. In TorchTitan this is named FullAC. Since GPU was idle for most of the baseline step, the extra compute did not lengthen it. As a result, throughput went from 90 to 370 tokens/s per GPU at 32,768 tokens.

At 65,536 tokens, the boundary activations alone take 74 layers × 65,536 tokens × 6,144 × 2 bytes ≈ 55 GiB per GPU, which does not fit in HBM next to the weights. GH200 has a 900 GB/s link to 120 GB of CPU memory, so we decide to offload them. Each boundary activation is copied to Grace after the forward pass produces it, and copied back to Hopper one layer before the backward pass recomputes it. An offloading parameter sets the share of layers that offload, which trades HBM for Grace memory.