# Training a 1T-Params MoE on EU public infrastructure

> Source: <https://www.domyn.com/blog/training-1t-moe-eu-public-infrastructure-pytorch>
> Published: 2026-11-12 13:44:30+00:00

At Domyn, we have recently started training a frontier-scale model as part of the European Frontier AI Grand Challenge, together with Fraunhofer IAIS and Fraunhofer IIS. The first challenge we faced was understanding if and how a model at the trillion-parameter scale could be trained on public computing European supercomputers, which are typically optimized for scientific workloads.

This post describes the technical details of how we achieved this goal. We believe this, in itself, is a breakthrough for Europe’s position in the AI race. While dedicated AI factories are being built, we demonstrate that public European supercomputers can already unlock frontier scale today.

In June 2026, the EUROPA consortium, led by Domyn with Fraunhofer IAIS and Fraunhofer IIS, was awarded the Frontier AI Grand Challenge. The goal of the project is the creation of an open frontier model, trained on European public infrastructure and released with its weights and training recipes. The consortium was given access to 2.5% of the EuroHPC compute capacity on public [European supercomputers](https://www.eurohpc-ju.europa.eu/supercomputers/our-supercomputers_en). 

In order to align with the current scale at the frontier, we set out to train a Mixture-of-Experts (MoE) language model with ~1T total parameters. Out of the European supercomputers in EuroHPC, we decided to focus on Jupiter, the exascale supercomputer of the EuroHPC Joint Undertaking, hosted at Forschungszentrum Jülich, for the pretraining run.

Models of this size have so far been trained mostly on dedicated Blackwell-generation clusters. That hardware provides higher bandwidth between GPUs, more memory per GPU and more freedom in the choice of parallelism layout. Jupiter's Booster is a Hopper-generation system shared by many users. We have user access only, so we cannot change the kernel, the interconnect, the node memory configuration or the system software. The nodes were configured for scientific simulation, and their memory and network balance differs from that of a cluster built for training. Hundreds of other projects run on the same machine.

The training configuration we use under these conditions, and the changes that raised throughput from 90 to ~1,300 tokens/s per GPU on 384 GH200 GPUs (96 nodes), are described below. We give enough detail for other users of Jupiter and similar EuroHPC systems to reuse it. Where we found problems in shared software, we fixed them in PyTorch and TorchTitan, and we reported system issues to the Jülich Supercomputing Centre (JSC). Refer to [Upstream contributions](#References) for the full list.

Figure 1 shows the improvements in throughput for each of the changes described in the post, starting from the initial configuration. Hovering over a row shows step time, memory and allocator retries.

Given the configuration of the hardware, described below, most of the changes focused on reducing the time the GPUs spend waiting on the network. This is shown in the next figure, where the profiled step of each traced run are split into compute, communication and idle time.

## System and model

### Hardware

[Jupiter](https://apps.fz-juelich.de/jsc/hps/jupiter/configuration.html) is the first European supercomputer in the exascale class. Its Booster module has 5,884 compute nodes, each with four NVIDIA GH200 Grace Hopper Superchips. Each superchip has the following components:

- **Grace:** an Arm CPU with 72 cores and 120 GB of LPDDR5X at 512 GB/s.
- **Hopper:** an H100-class GPU with 96 GB of HBM3 at 4 TB/s.
- **NVLink-C2C:** a 900 GB/s link between Grace and Hopper.

Within a node, GPU pairs are connected by NVLink 4 at 300 GB/s (150 GB/s per direction), and the four Grace CPUs are connected at 100 GB/s per direction. Each node has four InfiniBand NDR200 adapters for inter-node traffic, each with a throughput of 200 Gbit/s. Overall, they provide 800 Gbit/s (100 GB/s) per node, or about 25 GB/s per GPU.

The following bandwidths are referenced throughout the post:

### Software

For the EUROPA model we chose to adopt [PyTorch](https://pytorch.org/) and [TorchTitan](https://github.com/pytorch/torchtitan) ([Liang et al., 2025](https://arxiv.org/abs/2410.06511)), PyTorch’s native library for large-scale LLM pre-training. We picked TorchTitan for the following reasons:

- It is the training framework our team has the most experience with.
- Its code base is small and straightforward to modify for our setup.
- The project is actively maintained, and pull requests are reviewed and merged quickly.

### Model

The configuration we currently use for experiments is listed below. Some values may change before the training run starts.

1. Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., & Idreos, S. (2025). *TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training* . [https://arxiv.org/abs/2410.06511](https://arxiv.org/abs/2410.06511)

## Constraints and parallelism layout

### Constraints

As mentioned previously, raw GPU FLOPS were not the limiting resource. When it comes to pretraining a 1T MoE, the challenges resided in two other parts of the system.

#### Inter-node bandwidth

InfiniBand traffic between nodes is 6 times slower than between GPUs on the same node, and 36 times slower than between GPU and CPU. This link is crossed many times during a training step distributed over 384 GPUs per data parallel replica, like in our setup. We therefore need traffic over InfiniBand to be either small or overlapped with computation happening on the GPUs.

#### Host memory capacity

On Jupiter’s login nodes, GH200 has up to 480 GB of LPDDR5X at 384 GB/s per Grace. Compute nodes, however, have 120 GB at 512 GB/s, which favors bandwidth over capacity. While this configuration suits simulation workloads that are limited by memory bandwidth, for training, this limits how much state can be processed and stored on each node across HBM and Grace memory.

The usable amount is in fact less than 120 GB. On GH200, Linux exposes each Hopper’s HBM as a NUMA node without CPUs. Each node reports 858 GiB of memory, of which 456 GiB is Grace memory. Out of the total, the kernel reserves around 24 GiB, which leaves ~108 GiB of usable host memory per GPU. Additional losses are described in [Host memory losses](#Host-memory-losses).

### Parallelism layout

For a 1T MoE, choosing the correct parallelism layout for the hardware architecture is critical to achieve production throughputs. We initially considered pipeline parallelism, but given the principles below, we ended up focusing on FSDP and expert parallelism and spent our efforts optimizing them, which ultimately paid off. Since each layout dictates all further downstream optimizations, we cannot say whether a layout based on pipeline parallelism would have been faster. The decision was based on three principles:

1. **Expert parallelism within a node.** The token all-to-all of an MoE layer is on the critical path between the router and the experts, and it is difficult to overlap with computation. Its cost is acceptable over NVLink and too high over InfiniBand. Expert parallelism is therefore limited to 4, with one group per node.
2. **Short steps.** Deep pipeline parallelism needs many microbatches per step to keep its bubble small, which makes steps longer. This makes discovering opportunities for optimization and problems arising at scale harder, leading to slower feedback loops.
3. **Large batch.** A larger batch reduces gradient noise. It also gives each GPU more work per step, which amortizes fixed costs such as kernel launches and the optimizer step.

As per the above considerations, we landed on 384-way fully sharded data parallelism (FSDP), with 4-way expert parallelism inside each node. All runs in this post use a single data-parallel replica (DP = 1), so one FSDP group spans the whole job. In practice, we can run up to 4 data-parallel replicas across 1536 GPUs with no appreciable effect on throughput.

## Sequence length, optimizer and MoE kernels

### Baseline

FSDP overlapped only part of the communication. Computation on a layer often finished before the weights of the next layer had arrived, which led the GPU to stay idle. In the baseline (8,192 tokens per sequence), a profiled step on rank 0 lasted a total of 83.6 s, of which computation took 14.3 s, while 68.3 s were spent waiting for communication to take place.

Since communication in FSDP is limited to weights, and the volume of weights gathered per layer does not depend on the number of tokens, computing over longer sequences increases the computation per byte communicated. Therefore, we increased the sequence length as far as memory allowed. However, even with the baseline’s activation checkpointing, a sequence of 12,288 tokens already exceeded HBM capacity.

### Activation checkpointing and offload to Grace

Longer sequences need more activation memory. First, we stored activations only at layer boundaries and recomputed everything inside a layer during the backward pass. Full recomputation adds computation, but the GPU was idle for most of the step, so the added compute did not lengthen the step. With this change, the sequence length reached 32,768 tokens and throughput 370 tokens/s per GPU, 4× the baseline.

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 together with the weights and the remaining activations. We therefore decided to offload them over the 900 GB/s NVLink-C2C link. Each layer’s boundary activation is copied to Grace after the forward pass, and it is copied back to the GPU before recomputation during the backward pass. We set the `offload_fraction` parameter, which sets the share of layers that do the offloading and trades HBM for Grace memory, to 0.86. With offloading, the sequence length reached 49,152 tokens and throughput 508 tokens/s per GPU.

We [contributed](https://github.com/pytorch/torchtitan/pull/5090) the hooks for activation-checkpoint offloading to TorchTitan together with a minimal CPU offloader as a documented example. We then built the offloading strategy described here on top of the contributed hooks.

### Muon with expert-sharded Newton-Schulz

The optimizer state requires a significant share of memory. We use Muon ([Jordan et al., 2024](https://kellerjordan.github.io/posts/muon/)) for every weight matrix except the router, about 99.5% of the parameters, and AdamW for the embeddings, the LM head, the router and the norms. Muon stores one fp32 state per parameter (momentum), while Adam stores two. Removing the second state made a sequence length of 65,536 tokens fit in memory. Muon has been reported to train large models more efficiently than AdamW (Liu et al., 2025), and it has been used for trillion-parameter MoE models (Kimi Team, 2025; [Kimi Team et al., 2026](https://arxiv.org/abs/2607.24653)).

With Muon, each update is orthogonalized with a few Newton-Schulz iterations, which operate on the whole matrix. However, with our FSDP setup matrices are split into 384 shards. A straightforward implementation would gather each matrix on one rank, run Newton-Schulz there and scatter the result back, which significantly adds to communication.

[FlexShard](https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/flex_shard/README.md), now part of TorchTitan, enables us to choose which rank owns which parameters. Empowered by this, we sharded the experts by whole experts, so every rank owns complete expert matrices and runs their Newton-Schulz iterations without communication. The dense matrices, such as the attention projections, are still sharded and collected per rank. This setup allows us to further leverage computation-communication overlap, since each rank can process its own experts while the dense matrices are in transit, leading to a full communication overlap.

As a key step, we decided to offload the optimizer to run on the Grace processor. However, contrary to other optimizer-related operations, Newton-Schulz is computationally intensive and significantly benefits from GPU acceleration. The following figure compares Newton-Schulz on our expert shapes on Grace and on Hopper.

During training, the optimizer step took 41.7 s in total with Newton-Schulz on Grace. We [contributed](https://github.com/pytorch/torchtitan/pull/5352) support to TorchTitan for running Newton-Schulz on the GPU when parameters are offloaded. With this change, the optimizer step took 2.2 s, which led throughput to rise from 441 to 597 tokens/s per GPU (+35%).

On a CPU device like Grace, PyTorch streams and events have no effect. This causes the overlap that FlexShard provides on CUDA to be silently lost. We patched the CPU-side code to run the all-to-alls of the dense matrices asynchronously instead of blocking the main thread. The patch reduced the optimizer step to 1.6 s.

### Dynamic tile scheduling in SonicMoE

For the experts, we use [SonicMoE](https://github.com/Dao-AILab/sonic-moe) ([Guo et al., 2026](https://arxiv.org/abs/2512.14080)), a fused MoE kernel built on the CuTe DSL kernels of [QuACK](https://github.com/Dao-AILab/quack). It runs the expert GEMMs directly on the token-dispatch layout, without a permuted copy of the tokens. On a single GH200 at a test shape, it reduced the layer’s peak memory from 12.5 to 5.4 GiB, and forward plus backward time from 213 to 153 ms.

Communication kernels run on the same GPU SMs as compute kernels. Therefore, an FSDP all-gather in flight occupies some of the SMs with NCCL thread blocks. SonicMoE’s GEMMs use a static persistent scheduler, which assigns each of the 132 persistent thread blocks a fixed list of tiles at launch. A block whose SM is occupied by the all-gather starts its share of tiles only after the all-gather ends, and the GEMM finishes when its last block finishes. In one run, the same up-projection took 55-60 ms when executed alone and ~100 ms during an all-gather.

QuACK also provides a dynamic scheduler, in which blocks take tiles from a shared atomic counter, so running blocks also process the tiles of blocks that have not started. SonicMoE hard-codes the static scheduler, so we wrap its GEMM calls to enable the dynamic one. We benchmarked one MoE layer at our training shape on one GH200, with and without a dummy kernel occupying 40 SMs.

With 40 SMs occupied, the dynamic scheduler leads GEMM to be 17-45% faster, with the layer time decreasing from 304 to 217 ms. On an idle GPU, it is 3-32% faster per GEMM and 11% faster on the layer. Outputs are bitwise identical, except for fp32 rounding in the routing-score gradient. In training, we enabled it on the four expert GEMMs whose token dimension varies per expert. The step time decreased from 108.7 s to 103.8 s, which resulted in a further 5% throughput gain.

We have [asked the SonicMoE authors](https://github.com/Dao-AILab/sonic-moe/issues/75) why the static scheduler is the default. [Limiting NCCL SM usage](#Limiting-NCCL-SM-usage) addresses the same contention by reducing the number of SMs NCCL occupies.

## Load balancing across ranks

The 384 GPUs run in lockstep, so the slowest rank dominates step time. Two sources of imbalance made some ranks slower than others.

### Attention-cost balancing of training data

Attention cost grows with the square of document length. At 65,536 tokens per row, one rank could receive a single 60k-token document while another a hundred short ones. Both process the same number of tokens, but the first does far more attention work, which leads computation on that rank to take longer.

To alleviate this effect, we balance the data offline. We pre-tokenize the corpus and pack documents, in random order, into rows of 65,536 tokens. A document that does not fit continues in the next row, and for each row we compute its attention cost as the sum of squared segment lengths, where a segment is the part of a document inside the row. We sort rows into narrow cost buckets, shuffle the rows within each bucket, and fill each step with consecutive rows from this order, one per rank. The rows in a step therefore have nearly equal cost. Finally, we shuffle the order of the steps.

Over the 1,057 steps, the most expensive rank in a step did a median of 7.9× the mean attention work before balancing, and 1.15× after.

In training, data balancing reduced the step time from 103.8 s to 91.7 s, leading to a 13% throughput gain.

Balancing changes the data distribution within a step. Long documents are grouped into the same steps instead of arriving in their natural mix. If one rank receives a row with a single 65k-token document, every other rank in that step receives a similar row, which can affect training dynamics. We expect the effect to be limited to long documents, and we plan to compensate for it with a dedicated long-context phase.

### Group-limited expert routing

The second source of imbalance is the MoE router. When it sends more tokens to the experts on one GPU, that GPU becomes the slowest in its node. In earlier runs without data balancing, expert imbalance cost ~16-17 s per step, about 20% of the step.

We experimented with making the expert load exactly even. The 384 experts are split into four groups of 96, one per expert-parallel rank. The router selects the same number of experts from each group (two per group with top-8 routing). Every GPU then receives the same number of token-expert pairs in every layer. On top of data balancing, this reduced the step time from 91.7 s to 85.9 s, with a 7% throughput gain. Of note, peak HBM decreased from 87.2 to 81.3 GiB, because no GPU needs memory for a surplus of tokens.

The constraint however reduces routing freedom, which may impact training efficacy. As an example, a token whose three highest-scoring experts are in the same group is forced to be routed to only two of them. Selecting more experts per group increases the chance that each token reaches its highest-scoring experts while keeping the load even. Our early runs indicate that this constraint can indeed hurt training, and we are still evaluating whether to keep it or relax it with an intermediate strategy.

## Reducing communication volume

Beyond a certain sequence length, longer sequences no longer improved throughput. When activations put pressure on memory, the CUDA allocator starts retrying and fragmenting. Fragmentation ultimately costs more time than the time saved by using longer sequences. From that point, the remaining option was to send less data. Before the changes in this section, a profiled step spent 39.5 of its 98.7 s waiting on the network ([step breakdown](#Step-time-by-activity)), and all-gathers and reduce-scatters occupied the GPU for more than 50 s.

### FP8 expert weight all-gather

From profiling, the expert weight all-gather is the largest communication cost in a step, at about 7 GB per rank in bf16, or 28 GB per node. Gathering the weights in FP8 would halve this volume.

To this end, the weights are split into blocks of 128 values, and each block is scaled to the FP8 (E4M3) range and cast. Each rank sends FP8 values plus one fp32 scale per block; the receivers then dequantize them to bf16 before the expert computation. The forward pass therefore uses the weights at lower precision, with an error of up to ~6% of the block maximum per value. Master weights, gradients and optimizer state remain in full precision. This approach is close to the one used for DeepSeek-V3 (DeepSeek-AI, 2024).

With this strategy, all-gather time on the GPU decreased from 53.6 to 35.1 s per step; throughput rose from 763 to 898 tokens/s per GPU (+18%). Early results show no appreciable effect on training.

Because the weights already arrive in FP8, the expert GEMMs could also run in FP8 with a small additional loss of precision, as in DeepSeek-V3’s FP8 MoE training. We may incorporate this as a further optimization.

### BF16 gradient all-to-all with FP32 accumulation

With FP8 gathers, the forward-pass communication was fully overlapped, but the backward pass remained communication-bound because of gradient reduction.

In the standard FSDP mixed-precision setup, gradients are computed in bf16 and reduced in fp32. They are cast to fp32 before the reduce-scatter, so the ring accumulates in fp32 and gradients are communicated in fp32 as well.

However, only the summation requires fp32, while the transfer can happen in bf16. Following this observation, we decided to send gradients in bf16 during the all-to-all, in which each rank sends every other rank the slice of its gradient owned by the other rank. The owner upcasts the slices and sums them locally in fp32. The result matches the fp32 reduce-scatter up to summation order, and it was bitwise identical in our 4-rank tests. The all-to-all sends half the bytes and halves the gradient buffer held between layers.

GPU time in gradient collectives decreased from 50.7 s to 33.2 s per step (10.4 s of reduce-scatter and 22.8 s of all-to-all). Time spent waiting on communication decreased from 27.3 s to 13.9 s, throughput rose from 898 to 1,022 tokens/s per GPU (+14%). With all other changes in this section applied, disabling the bf16 all-to-all reduces throughput from 1,210 to 1,004 tokens/s per GPU.

Neither FP8 gathers nor bf16 gradients change the loss measurably. In a set of comparison runs, using the same data order and initialization, after ten steps losses are within 0.01 of each other.

### Node-local resharding of dense weights

With FSDP the layer’s weights are freed after each layer’s forward pass, and then gathered again in the backward pass for recomputation. Each backward re-gather competes for InfiniBand bandwidth with the gradient reduce-scatter.

The experts make up most of the model and cannot stay in memory during that time. The dense parameters (attention projections, router and norms), however, can stay on the GPU if sharded. After the forward pass, we reshard each layer’s dense weights to 1/4 within the node instead of 1/384 across the job, so the backward re-gather runs over NVLink instead of InfiniBand. This uses ~55 MiB of HBM per layer, or ~4 GiB in total. Conversely, keeping the dense weights fully unsharded (~16 GiB) caused allocator retries at our sequence length, with a consequent loss in throughput.

The change required little new code. TorchTitan registers parameters on two meshes, a dense mesh spanning all ranks and a sparse mesh for experts, and the expert-parallel dimension of the sparse mesh corresponds to the node. We set the post-forward resharding of the dense parameters (`post_forward_mesh_info`) to the node-local group, and relaxed a few checks in FSDP that were stricter than necessary.

On its own, this change improved throughput by 2% (1,022 → 1,038 tokens/s per GPU). The additional 4 GiB caused allocator retries on all 384 ranks.

### Allocator pools on FSDP side streams

After an upgrade of our software stack, the CUDA allocator started retrying on GPUs with ~70% memory utilization. A retry frees the allocator’s whole cache and requires a full-device synchronization, which is slow. At the time of the retries, 25-30 GiB of HBM was free, mapped and in the allocator’s cache, but the allocator could not use it.

The cause is the way FSDP overlaps communication with computation. A GPU executes work in streams, which are queues that run concurrently. FSDP uses one compute stream and four side streams:

The PyTorch caching allocator keeps one memory pool per stream. Freed memory returns to the pool of the stream that allocated it, and only that stream can reuse it, because another stream would need a synchronization to know that the previous user has finished. FSDP allocates each buffer on the stream that writes it. Our FP8 all-gather and bf16 all-to-all allocate their temporary buffers inside FSDP’s calls, so these buffers were also allocated on the side streams. Each buffer is used briefly once per layer, but each stream’s cache ends up retaining the largest blocks it has allocated.

We replayed the allocator’s event log from a memory snapshot of rank 0. At peak live memory, the side streams held 27.3 GiB, of which 6.4 GiB was in use. The remaining 20.9 GiB was unavailable to the compute stream, which had 4.5 GiB of free cache. At each retry, the compute stream could not find a free block for a 4 GiB SonicMoE buffer. Reducing the number of layers did not help, because it reduced live memory and left the unused cache unchanged.

We made three changes so that this memory is allocated from pools the compute stream can reuse:

1. **Earlier gradient release.** Our gradient-norm hook (see[Smaller memory and latency changes](#Smaller-memory-and-latency-changes) ) runs on every layer. When a hook is registered, FSDP keeps each layer’s reduced gradient alive until the end of the backward pass, so that the reduce-scatter stream does not reuse it while the hook reads it. This is safe for one hook, but for 74 hooks it costs 9 GiB. We release each gradient two layers later, after making the reduce-scatter stream wait for the hook.
2. **All-gather on the compute stream.** FSDP2 had a private, undocumented flag, unshard_async_op, that runs the all-gather setup on the compute stream and lets NCCL run on its internal stream. The buffers are then allocated from and returned to the compute stream’s pool. This required an asynchronous mode for our FP8 gather, which now returns a handle that dequantizes when FSDP waits on it, so the compute stream waits for the network only when it needs the weights. We [proposed](https://github.com/pytorch/torchtitan/pull/5627) a public TorchTitan option to enable the flag. On a smaller test model, it reduced peak reserved memory by ~11% at a ~0.7% throughput cost.****
3. **Gradient receive buffer on the compute stream.** The receive buffer of the bf16 all-to-all, which is as large as the gradient, is allocated from the compute stream’s pool.

On the full 384-GPU job, with the changes in [Smaller memory and latency changes](##Smaller-memory-and-latency-changes) already applied, the number of ranks with allocator retries decreased from 384 to zero. Peak reserved memory decreased from 88.1 to 74.6 GiB, and throughput rose from 996 to 1,210 tokens/s per GPU (+21%). The volume of data sent is unchanged; only the pools that memory is allocated from differ.

Following the review of our [TorchTitan PR](https://github.com/pytorch/torchtitan/pull/5627), the `unshard_async_op` flag is being moved into a public FSDP API ([PyTorch #202343](https://github.com/pytorch/pytorch/pull/202343)).

1. DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. *arXiv Preprint arXiv:2412.19437* .

## Smaller memory and latency changes

The changes in this section have a small effect on step time individually. We applied them while every rank was retrying allocations (see [Allocator pools on FSDP side streams](#allocator-pools-on-fsdp-side-streams)), which led to throughput to go from 1,038 to 996 tokens/s per GPU. We have not measured their individual effect on throughput, however they reduced memory use.

### Gradient norm during the backward pass

Gradient clipping requires the global gradient norm. With CPU offload, TorchTitan computes it after the backward pass over gradients already on Grace, which takes ~0.43 s of `_foreach_norm` on the CPU and ~0.19 s of gloo all-reduces.

Each gradient shard is on the GPU when its reduce-scatter completes. At that point, we add its squared norm to a running scalar. After the backward pass, one NCCL all-reduce of the scalar gives the global norm, and clipping proceeds as before.

### No embedding all-gather in the backward pass

By default, FSDP gathers the embedding table again in the backward pass. The embedding gradient depends only on the input token IDs, which every rank already has, so the gather is unnecessary. On our communicator, it also waited behind the reduce-scatter of layer 1, which exposed up to ~126 ms per step. The fix is one line:

We have [submitted this contribution upstream](https://github.com/pytorch/torchtitan/pull/5579).

### Shared RoPE cache

TorchTitan builds a RoPE cache for each layer, although layers with the same configuration get identical caches. At 65,536 tokens, our 37 RoPE layers held 2.75 GiB of copies of two distinct caches, which is enough to cause an out-of-memory failure. Sharing the caches reduces this to 0.19 GiB, a saving of 2.56 GiB per GPU.

Sharing a single buffer was [proposed upstream](https://github.com/pytorch/torchtitan/pull/4111), but the proposal stalled over concerns about fragility and pipeline parallelism. We have [submitted](https://github.com/pytorch/torchtitan/pull/5583) a smaller version that ties the cache where TorchTitan ties embeddings, so it is compatible with pipeline parallelism.

## Limiting NCCL SM usage

Overlapping communication with computation means running both on the same GPU at the same time. This section covers the SMs that communication kernels take from computation kernels.

NCCL runs each collective as a GPU kernel with one cooperative thread array (CTA, NVIDIA’s term for a thread block) per channel. Each CTA occupies an SM for the duration of the collective and spends most of that time waiting for the network. Our GEMM and attention kernels need most of an SM’s shared memory, so they cannot run on an SM that hosts an NCCL CTA. NCCL selects the number of channels needed to drive the links at full speed. On our job, it used 40 CTAs for the expert all-gather and 64 for the bf16 gradient all-to-all. Hopper has 132 SMs. When these two collectives overlapped, NCCL held up to 88 SMs, which left 44 for the GEMMs overlapping with them. Over the profiled step, NCCL held 50 SMs on average while any collective was running, and 64 or more for 20 s.

The [dynamic scheduler](#dynamic-tile-scheduling-in-sonicmoe) reduces the effect of lost SMs on the expert GEMMs. The change described here reduces the number of SMs NCCL takes. These collectives cross InfiniBand at ~25 GB/s per GPU, which limits their speed. Beyond a certain number of CTAs, additional CTAs do not increase throughput. `NCCL_MAX_CTAS` caps the number of CTAs per collective, and we set it to 16.

The two large collectives affected by the cap were slightly slower. The expert all-gather took 155 ms with both 16 and 40 CTAs, and the all-to-all was 3% slower. Other collectives changed more, in both directions, as described below. Both runs process the same data in the same order, so each compute kernel launch in one run does the same work as its counterpart in the other. The figure below shows the time of the main compute kernels over the step in each run, as well as the number of SMs NCCL left free while they ran.

All of these kernels became faster, approximately in proportion to the number of SMs they regained. The FA3 backward kernels ran mostly while no collective was in flight. They had 126 SMs by default and 130 with the cap, and took 3% less time. The QuACK GEMMs with a fused activation went from 81 to 118 SMs and took 22% less time. The cuBLAS GEMMs with the TN layout went from 82 to 112 SMs and took 20% less time.

In total, these kernels went from 52.5 to 44.8 s, and the expert GEMMs alone from 23.5 to 19.2 s. A few in-node communication kernels (an NVLink all-gather and the HybridEP sync) spent longer waiting, so the step’s total GPU time decreased by 6 s, from 60.6 to 54.7 s.

Exposed communication increased by 1.6 s, which is less than the computation time saved. The dense collectives kept their 8 CTAs and finished sooner (119 to 40 ms for the all-gather and 131 to 78 ms for the reduce-scatter). A likely explanation is that a collective also waits for the slowest rank, and every rank now reaches the collective earlier. With this environment variable alone, throughput rose from 1,210 to 1,293 tokens/s per GPU (+7%).

We tested only the value 16, and other values may perform better. The cap also applies to collectives that need more SMs. The dense re-gather over NVLink ran on 16 CTAs instead of 24, and its time increased from 1.3 to 17 ms. The difference is small for this job, but a job with more intra-node traffic would need a higher cap.

## Host memory losses

The most time-consuming problems in this work were differences between the host memory we budgeted for and the memory available on the nodes. Grace has ~108 GiB of usable memory per GPU, and the training configuration is set to use nearly all of it on every rank. Memory lost on any rank lowers the limit for the whole job. This section describes four causes of lost memory, approximately in the order we found them.

### Host allocations placed on HBM

Because HBM is exposed as a NUMA node, Linux places ordinary host allocations on it when Grace memory is full. The allocation succeeds without an error, however, CUDA later runs out of memory with gigabytes in use that are not accounted for by any process. `numactl --membind` prevents the kernel from doing this, but it does not cover allocations initiated by the GPU. A pageable host buffer whose first write is a GPU-to-CPU copy has its pages placed on the GPU’s HBM, even under a strict bind. We found that our activation-offload arena was behaving that way, leading to 13-17 GiB of HBM usage on every GPU. We fixed this by writing every page of the arena from the CPU once at startup.

### Pinned memory rounding

FSDP’s CPU offloading pins all sharded parameters and gradients. The PyTorch pinned-memory allocator rounds each allocation up to the next power of two, which leads a 70 MiB tensor to occupy 128 MiB. Across the model, this padding amounted to 14-28 GiB per rank, and it was not visible in /dev/shm, ipcs or smaps.

Setting `pinned_max_round_threshold_mb:1` in `PYTORCH_CUDA_ALLOC_CONF` disables the rounding for allocations above 1 MiB. For four test allocations, the setting reduces the memory retained from 1,728 MiB to the 1,077 MiB requested.

### Replicated QKV weight in checkpointing

A `state_dict` hook that splits the fused QKV weight for checkpointing gathered the full weight on every rank, which led the checkpointer to cache the result for the whole run. This used 28 GiB per rank, independent of the number of nodes. We therefore resorted to storing the fused layout for our checkpoints. TorchTitan has since moved the QKV conversion out of training checkpoints and into its Hugging Face adapters ([#4815](https://github.com/pytorch/torchtitan/pull/4815)), which removes the cause upstream.

### Memory unavailable at job start

After these fixes, some ranks still ran out of Grace memory. Interestingly, out-of-memory failures were reported as segmentation faults, probably due to a lower maturity of the underlying stack.

We found that nodes were not starting our jobs with their memory free. We logged every Grace socket at the start of 62 of our jobs, for 20,436 socket observations in total. A socket has 119.7 GiB. The median socket had 100.7 GiB free at job start, and the lowest had 72.7 GiB. Much of the difference is memory that the kernel attributes to no process and no meminfo category. This unattributed memory had a median of 10.9 GiB per socket, 17.6 GiB at the 90th percentile and a maximum of 39.9 GiB. Most sockets release it under memory pressure, however this didn’t happen for a few sockets, which had global impact.

We also found one leak in our own environment. The `libc10` we build against uses mimalloc, which backs its arenas with 512 MiB transparent huge pages. Setting `MIMALLOC_ALLOW_THP=0` removed the leak.

We reported the unattributed memory to JSC support. They confirmed the issue and recommended staying under ~75 GB per rank until the job epilogues are fixed. The issue affected all users of the machine, and it has since been fixed.

## Parameter initialization with CPU offload

CPU offload exercises code paths that are generally not utilized. This section describes two initialization problems we found on these paths.

### Identical random values across ranks

With CPU offload, TorchTitan initializes each parameter shard directly on the CPU, which saves HBM. On the GPU, PyTorch’s distributed RNG offsets each rank’s random state so that shards are drawn independently. On the CPU, the initializer does not detect that the parameter is on the CPU and takes the CUDA code path. That path draws from the CPU generator correctly, but it resets the CPU RNG state when it finishes.

As a result, every initializer call starts from the same random state, which leads every tensor on every rank to use the same random values. For TorchTitan’s truncated normal, the initialization does not terminate. Values outside the truncation range are resampled, but because the RNG state never advances, the same out-of-range values are drawn on every iteration.

We fixed the problem in two places:

- **PyTorch:** an [issue](https://github.com/pytorch/pytorch/issues/196562) and a [fix](https://github.com/pytorch/pytorch/pull/196072) that makes the RNG state advance. The fix guarantees termination, but it does not ensure that ranks draw different values.
- **TorchTitan:** a[PR](https://github.com/pytorch/torchtitan/pull/4600) that lets users choose the mesh axes along which the seed is offset.

With the PyTorch patch applied and `dp_shard` as the seeding axis, the model initializes correctly with offload enabled.

### Slow truncated-normal initialization

Initialization took ~10 minutes, and the cause was a second problem in the truncated normal. Its rejection loop uses `.any()` to check whether any value still needs resampling. On a sharded DTensor, this full reduction redistributes the mask to every rank, so each iteration runs an all-gather of a mask as large as the tensor to compute a single boolean. The cost is negligible for small tensors or fast links. For a large fused expert tensor over InfiniBand, however, each check can take seconds, which adds up to minutes across all layers.

The fix reduces each shard locally and all-reduces a single flag. Every rank gets the same result, so the RNG states of all ranks stay synchronized. Refer to the [PR](https://github.com/pytorch/pytorch/pull/201840) for details.

## Upstream contributions

The configuration in this post depends on open-source software and on publicly funded compute. We contribute all the fixes upstream, and we report system issues to the operators of the supercomputer, so that other users of PyTorch, TorchTitan and Jupiter can benefit from them. The table lists the contributions described in this post.

## Conclusion and next steps

The final configuration for the EUROPA model reaches ~1,300 tokens/s per GPU on 384 GH200 GPUs per DP replica, 14× the baseline. At this rate the job processes about 43B tokens per replica per day. With four data-parallel replicas on 1,536 GPUs, at the 89% scaling efficiency we measured, it would process about 150B tokens per day. With this work, we demonstrated that a publicly funded, shared, Hopper-generation system with user-level access can host training at the 1T scale, which is a significant achievement for the European AI ecosystem.

The configuration we converged to is built around the two main limits we found: host memory per GPU and bandwidth between nodes.

Each of the two limits required a dedicated set of changes. For inter-node bandwidth, increasing sequence length and switching to FP8 weight gathers and BF16 gradient all-to-alls increased the ratio between the computation and communication. Node-local resharding moved the backward re-gather of the dense weights onto NVLink, and `NCCL_MAX_CTAS` freed up SMs for the compute kernels. For host memory, we moved activations and optimizer state to Grace and kept Newton-Schulz on the GPU, and we tracked down causes of lost host memory. We want to highlight that several of these problems appear only with CPU offload at this model size, which is a less beaten path for PyTorch and TorchTitan in general. Nevertheless, the fixes are now already in PyTorch and TorchTitan, as a testimony to the activity of these projects and the dedication to their community.

Several settings remain unexplored in this post. We have not documented how severely group-limited routing costs in model quality in the long run, whether the expert GEMMs can run in FP8 without a measurable loss, whether values of `NCCL_MAX_CTAS` other than 16 perform better, or whether a layout based on pipeline parallelism would have been faster. We plan to report on them as the training run progresses.

The model and all its training recipes will be released openly by the time the EUROPA model is published.

## Acknowledgments

We thank the Jülich Supercomputing Centre for their support, in particular with the host memory investigation, and the PyTorch and TorchTitan maintainers for reviewing and merging our pull requests. EUROPA is funded by the EuroHPC Joint Undertaking under the AI-BOOST Frontier AI Grand Challenge.

## References

1. Guo, W., Mishra, M., Cheng, X., Stoica, I., & Dao, T. (2026). *SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations* .[https://arxiv.org/abs/2512.14080](https://arxiv.org/abs/2512.14080)
2. Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., & Bernstein, J. (2024). *Muon: An optimizer for hidden layers in neural networks* .[https://kellerjordan.github.io/posts/muon/](https://kellerjordan.github.io/posts/muon/)
3. Kimi Team. (2025). Kimi K2: Open Agentic Intelligence. *arXiv Preprint arXiv:2507.20534* .
4. Kimi Team, Bai, T., Bai, Y., Bao, Y., C., M., Cai, J., Cai, X., Cao, P., Cao, Y., Chai, Z., Charles, Y., Che, H. S., Chen, G., Chen, G., Chen, G., Chen, H., Chen, J., Chen, J., Chen, J., … Zu, X. (2026). *Kimi K3: Open Frontier Intelligence* .[https://arxiv.org/abs/2607.24653](https://arxiv.org/abs/2607.24653)
5. Liu, J., Su, J., Yao, X., Jiang, Z., Lai, G., Du, Y., Qin, Y., Xu, W., Lu, E., Yan, J., & others. (2025). Muon is Scalable for LLM Training. *arXiv Preprint arXiv:2502.16982* .
