EUROPA 用 PyTorch 与 TorchTitan 在 JUPITER 公有机群训练 1T 参数 MoE

PyTorch：Blog· The EUROPA consortium (Domyn, Fraunhofer IAIS and IIS)·2026-11-12 20:00· 1 小时前AI 评分40

Training a 1T-Parameter MoE on EU public infrastructure with PyTorch and TorchTitan

AI 导读

EUROPA 联盟借助 PyTorch 与 TorchTitan，在 EuroHPC 的 JUPITER 超算上把约 1T 参数 MoE 的预训练吞吐从每 GPU 90 tokens/s 提升到约 1,300 tokens/s，使用 384 块 NVIDIA GH200（96 节点）。

正文 · 原文

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:

1. FlexShard placement for the Muon optimizer, which removes most of the optimizer’s communication.
2. Activation-checkpoint offloading to the Grace CPU.
3. 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.

TorchTitan’s activation-checkpointing wrapper made this change convenient. We contributed the hooks for activation-checkpoint offloading to TorchTitan and built our offloading strategy on top of them. With offload, throughput rose from 370 to 508 tokens/s per GPU (+37%), and the sequence length grew to 49,152 tokens and then, with Muon, to 65,536.

Figure 4. Offload schedule, drawn with 16 layers in place of 74. During the forward pass, each boundary activation moves to Grace. During the backward pass, it returns to HBM one layer before its recomputation.

De-fragging FSDPs Streams

After an upgrade to our software stack, the CUDA caching allocator started retrying allocations on every rank at ~70% memory utilization. At the time of the retries, 25-30 GiB of HBM was free, but the allocator could not use it.

After investigating the root cause, we identified it in how FSDP overlaps communication with computation. FSDP uses one compute stream and four side streams, and the allocator keeps a separate pool per stream. Memory freed by a stream can only be reused by that stream. FSDP allocates each buffer on the stream that writes it, and some of our other communication optimizations allocated temporary buffers inside FSDP’s calls, so these buffers also landed on side streams. Replaying the allocator’s event log at peak memory showed the side streams holding 27.3 GiB, of which only 6.4 GiB was in use. The remaining 20.9 GiB was out of reach of the compute stream, which then could not find a 4 GiB block for a MoE buffer.

Figure 5. HBM held by each allocator stream’s cache at peak live memory, reconstructed from the allocator’s event trace on rank 0. Solid bars are live tensors, and dashed bars are cached memory that only the owning stream can reuse.

Luckily, FSDP2 has a private flag, unshard_async_op, that runs the all-gather setup on the compute stream and lets NCCL run on its internal stream. The gather buffers are then allocated from, and returned to, the compute stream’s pool. We adapted our custom gather to run asynchronously, so the compute stream waits for the network only when it needs the weights. This brought the number of ranks with allocator retries from 384 to zero. Peak reserved memory went from 88.1 to 74.6 GiB, and throughput rose from 996 to 1,210 tokens/s per GPU.

We proposed a TorchTitan option to enable the flag. Following the review, unshard_async_op is being moved into a public FSDP API (PyTorch #202343), so other FSDP users can rely on it.

Upstreaming contributions and closing remarks

The configuration we just described allows to scale pre-training runs to 1T parameters at a workable throughput on public European infrastructure. In order to make it possible for others to reproduce these results (more details are available in a companion post), we upstreamed all the fixes to PyTorch and TorchTitan.

In addition to the three contributions we already mentioned, we contributed a shared RoPE cache to TorchTitan (2.56 GiB per GPU at our sequence length), skipping the embedding re-gather in the backward pass, and configurable seed-offset mesh axes. In PyTorch, we fixed CPU initialization resetting the RNG state and trunc_normal_ all-gathering DTensors. We also reported unattributed host memory at job start to the Jülich Supercomputing Centre, since it affects every user of the machine.

Thanks to the level of control provided by PyTorch and TorchTitan, training now runs at a sustained throughput of ~1,300 tokens/s per GPU on 384 GH200s per DDP replica, for a total of 1536 GPUs for 4 DDP replicas, and fits in memory by fully leveraging the HBM and the Grace architecture. The companion post covers the rest of the path in Figure 1: data and expert load balancing, reducing the bytes sent over InfiniBand, node-local resharding, dynamic tile scheduling in SonicMoE and capping NCCL’s SM usage.

We are grateful to the European Commission for Frontier AI Grand Challenge award. We hope that this work will motivate further investments in public compute infrastructure and the continued growth of the European AI ecosystem.

Learn more:

• Training a 1T-Parameter MoE on Jupiter (full post)
• PyTorch
• TorchTitan
• FlexShard
• Frontier AI Grand Challenge announcement

The post Training a 1T-Parameter MoE on EU public infrastructure with PyTorch and TorchTitan appeared first on PyTorch.

来源：PyTorch：Blog · pytorch.org