More parameters, longer training steps

You want to train a language model across several graphics processors, or GPUs. A training step processes a batch of text, measures prediction errors and updates the model’s parameters, the coefficients it learns. You want more parameters without increasing computation per step by the same proportion. That is the opportunity behind a mixture of experts, or MoE.

Suppose you already have a model without expert selection, called a dense model. You replace selected blocks with experts while keeping comparable computation per token. The new model contains more parameters but activates only a subset at a time. Yet a training step now takes longer. What are the GPUs doing during that extra time?

Understanding that result requires watching the machines. Some multiply matrices while others wait for inputs. Experts receive different token counts, transfers cross different links and temporary memory buffers consume space before contributing useful computation. Saving operations does not guarantee saving elapsed time.

Hugging Face’s Ultimate MoE Training Handbook, by Nouamane Tazi, Leandro von Werra and Thomas Wolf, provides the starting point. This article develops its diagnostic approach using the primary research and implementations linked below. LeCompute has not run multi-GPU benchmarks for this article. The numerical scenario is a constructed example.

An expert transforms numbers inside the model

The model splits text into tokens, pieces that can correspond to a word or part of a word. It associates each token with a list of numbers, a vector. Network layers transform that vector to prepare a prediction. Attention exchanges information between text positions; a feed-forward network, or FFN, then transforms the vector at each position.

In the MoE Mixture-of-Experts. An architecture where the network is split into many experts, of which a router activates only a small subset per token. Per-token compute follows the number of active parameters; memory follows the total count, since every expert must stay resident in VRAM, ready to be called. architectures discussed here, that FFN block offers multiple transformations: the experts. Each subnetwork has its own learned coefficients, or weights. Two experts can have the same structure and receive the same vector while producing different outputs because their weights differ. An expert returns another vector to the network; it does not write a complete answer to the user.

The router is a small network that reads the input vector and scores the experts. It retains the highest-scoring ones. Selecting two is called top-2; top-k is the general term for selecting k experts. Other experts in that layer do not compute outputs for this token. The Mixtral paper describes this selection and the weighted sum that follows.

For an illustrative mixture, suppose the router selects A and B with weights of 0.75 and 0.25. The block returns 0.75 × output_A + 0.25 × output_B, coordinate by coordinate. Those coefficients specify each output’s contribution here, rather than the probability that an expert is correct. The next layer receives the combined vector, and another MoE layer can make a different selection.

Selecting a limited subset lets the architecture add experts without executing all of them for each token. Total weights increase while only some participate in this layer’s computation. The remaining experts’ weights still need storage for tokens that select them.

How the router learns without an expert-selection answer key

When training a model to predict text continuations, the dataset supplies the token actually present at the next position. The model produces a probability distribution over possible tokens. An error function, called the loss, penalizes low probability assigned to the expected token. Router and experts share the goal of reducing this final error.

The program then works backwards through the operations producing that prediction. This backpropagation calculates gradients, which describe how small parameter changes would affect the loss. An optimizer uses those signals to adjust weights. Repeating this process across many text batches builds prediction quality.

Return to A and B’s contributions. If slightly increasing A’s mixing weight would reduce final error, the gradient carries that information to the router. If a small change in A’s parameters would also help, a signal reaches that expert. These signals are calculated through the chain of operations without rerunning a complete prediction for every possible change.

Here, the signal comes through the weighted mixture of selected outputs. The choice retaining two experts and excluding the others does not itself provide a derivative comparing every possible selection. Transforming scores into mixing weights therefore belongs to the learning mechanism, rather than merely displaying percentages. Routing and learning mechanism in Switch Transformers.

The router does not run every expert to compare their answers against a reference. It learns through the paths the model executes. As its weights change, expert scores and selections can change. Specialized behaviour can emerge without an engineer assigning a “code” or “French” expert in advance.

When experts live on different GPUs

So far, we have described mathematical operations. Executing them on a cluster, a set of connected machines, requires deciding where weights live. expert parallelism A way of splitting a MoE where each GPU holds a subset of the experts, instead of receiving a slice of every matrix as in tensor parallelism. Each token therefore has to travel to whichever GPUs hold the experts picked for it, so the collective becomes an all-to-all rather than an all-reduce. Usually shortened to EP. places different experts on different GPUs. Selecting an expert on another device means delivering the vector to that device before computation can begin.

These intermediate vectors are activations. Dispatch groups and delivers them to experts; combine returns outputs and merges their weighted contributions. A dispatcher is the software component organizing those exchanges. With all-to-all A collective where every rank sends distinct data to every other rank. An expert-parallel MoE layer runs two of them: dispatch ships each token's hidden state to the GPUs holding its experts, combine brings the outputs back. Its volume scales with the token count, so unlike weight reads it never amortizes over a larger batch. , participating processes can each send different data to the others. In distributed training, a process typically drives one GPU and has an identifier called its rank. DeepEP communication operations.

Router selects experts for each token
Dispatch delivers activations
Experts compute on their GPUs
Combine returns and weights outputs
Next block uses the combined result
Forward-pass schematic. Backpropagation subsequently follows reverse dependencies and is not shown here.

The schematic shows the forward pass, which calculates results from inputs. Backpropagation subsequently exchanges information needed for gradients. Network cost therefore extends beyond sending input vectors once. The router selects transformations, and expert placement turns those choices into transfers and device workloads.

Four GPUs, with one holding up the group

Illustrative assumption. A layer has eight experts with the same structure and computational cost, each with its own weights. Two experts live on each of four GPUs. Together, the four processes handle 4,096 tokens and select two experts per token. Each token therefore triggers two expert computations, giving 8,192 token-expert assignments. Ignore other forms of parallelism here.

An even distribution would give each expert 1,024 assignments and each GPU 2,048. Instead, suppose routing produces the following loads:

GPUAssignments to its two expertsDevice total
GPU 03,072 + 1,0244,096
GPU 11,024 + 1,0242,048
GPU 2512 + 5121,024
GPU 3512 + 5121,024
Calculated example by LeCompute, without hardware measurements. There are still 8,192 assignments; their distribution has changed.

GPU 0 receives twice the average device load. Operations depending on its outputs must wait even when other devices finish earlier. Step time need not double: matrix operations have different efficiencies at different sizes, and some work overlaps. Nevertheless, a cluster-wide average already hides the device worth investigating.

Observe both levels. Per-expert counts identify underused and overloaded experts. Summing the counts for each GPU exposes the physical distribution of work. A different expert placement could change that device distribution without changing how many assignments each expert receives.

To make the memory implications concrete, assume each vector delivered to an expert contains 2,048 values at two bytes each. The assignments represent 32 MiB of input vectors: 8,192 × 2,048 × 2, assuming a copy for each selected expert. This is not network-interface traffic: some computation stays local, transfers can be aggregated, and metadata, outputs and gradients add costs. It already shows why selecting another expert also increases data to process.

Why add a load-balancing objective?

We have seen how prediction error trains the router. That objective does not directly ask it to distribute work across GPUs. In our scenario, it can keep favouring the expert receiving 3,072 assignments even when that expert’s device holds up the others.

Rarely selected experts also receive less task-learning signal. An auxiliary load-balancing loss adds another objective encouraging work distribution. Its coefficient controls the importance of that constraint relative to prediction error. Switch Transformers describes this mechanism and selective higher precision for routing stability. Small score differences can change selection; retaining more precision protects decisions without running the whole model in FP32, the 32-bit floating-point format.

Auxiliary-Loss-Free Load Balancing instead adjusts a selection bonus or penalty, called a bias, according to each expert’s recent load. Favouring an underused expert increases its chance of selection. This correction affects selection separately from the scores weighting outputs. “Auxiliary-loss-free” means the load correction does not come from a second objective added to prediction error. The model still learns its task.

Assess both goals: learning useful transformations and avoiding a concentration of work on a subset of devices. An isolated balancing-loss curve cannot tell you what each GPU actually receives.

Which tokens are supposed to balance?

A sequence, a microbatch and a global batch contain different mixtures of data. A microbatch is the small group processed at a time by a process; multiple microbatches and processes can contribute to one global weight update. Demanding balance within each small group is a stronger constraint than demanding balance across their union.

Suppose one microbatch contains mostly code and another mostly French text. Enforcing balance inside each pushes both data types towards every expert. Enforcing balance only across their union allows different experts to receive more code or French. However, devices hosting the code-favoured experts may experience a load spike during the first microbatch.

The Qwen team’s Demons in the Detail examines that distinction. In its experiments, global balancing allows more domain specialization while permitting local imbalance. Measure loads in the microbatches the GPUs actually execute even when the balancing objective uses the global batch.

MaxVio measures how far the busiest expert exceeds average load: (maximum − mean) / mean. Our scenario gives (3,072 − 1,024) / 1,024 = 2, meaning the busiest expert receives three times the mean. It is neither prediction error nor the percentage of idle GPUs. Definition in the auxiliary-loss-free study.

Load spikes on different experts can offset one another over a longer period while still slowing each microbatch. Keep short-interval counters for each layer before calculating global summaries. That also prevents a satisfactory whole-model average from hiding one problematic layer.

Establish what the allocation can do

Before tuning the model, establish what the machines can do. An allocation is the set of resources reserved for training. A node is one cluster machine, potentially containing several GPUs. Record nodes, their connections or topology, and software versions. Allocations with equal GPU counts can include a degraded machine or use different network paths. Changing the machines and router together would obscure what caused an improvement.

On supported NVIDIA hardware, DCGM, the GPU diagnostic tool, provides readiness checks. Level 1 is intended for readiness; deeper diagnostics need an appropriate maintenance window. Passing a check does not remove the need to monitor errors and changes during training.

Next, test the communication pattern you use. All-reduce combines process contributions, for example to sum gradients; all-to-all distributes messages among destinations. Good performance on one does not establish good performance on the other. NCCL Tests checks performance and correctness for NCCL, NVIDIA’s communication library. Keep process count, message sizes, nodes, versions and the throughput column with each result. Synthetic busbw bandwidth is not directly the number of tokens your model can process per second.

The baseline distinguishes an unhealthy allocation from a different application workload. Compare network settings on the same machines. Establish another baseline when changing cluster size instead of applying one bandwidth threshold to every allocation.

Make the dense baseline answer a defined question

A dense model, which applies the same FFN block to every token in a layer, provides a useful control. Check that data arrives quickly enough, attention works, loss decreases and saved weights can be restored. This avoids attributing an existing training-loop problem to the MoE architecture.

State what stays equal. Matching active parameters brings part of the computation closer, but an MoE also stores weights of unselected experts. The optimizer may keep additional tracking information to update those weights, consuming more memory. Matching total parameters instead asks how much computation selection avoids in a model of that size. Specify the comparison before interpreting its result.

OLMoE varies expert size and count, experts that run for every token, and upcycling, the conversion of a dense model into an MoE. Isolating these changes helps connect a result to a particular choice. When comparing complete training budgets, the pretraining cost of a dense starting checkpoint belongs in the accounting.

For an engineering comparison, we recommend fixing data, tokenization, sequence length, tokens per update and evaluation protocol. Optimization settings can then be tuned transparently for each architecture. These include the learning rate, which controls weight-update magnitude. The same value may suit one model poorly while working for the other.

Read the timeline before blaming the network

An execution trace represents computation and communication as bars on each rank’s timeline. It shows who computes while someone else waits. If processes reach an exchange at different times, its observed duration can include waiting. The width of an all-to-all bar therefore includes more than time spent moving bytes.

In our scenario, compare GPU 0’s compute completion, token counts and subsequent exchanges. If the heavily loaded rank arrives late while others wait, balancing or placement deserves investigation. If loads are comparable but transfers remain long, prioritize topology and the dispatcher.

Artificially balanced routing can serve as a diagnostic control. It changes model decisions and therefore does not validate model quality. A shorter step reveals headroom associated with that workload distribution, including its effects on matrix sizes and communication.

Local work also matters. Tokens must be grouped by expert and restored to their original order after expert matrix operations. Grouped GEMM means executing groups of matrix multiplications. Grouping helps occupy a GPU when individual experts receive too few tokens. The Megatron Core MoE documentation describes these optimizations and their constraints. Small operations can leave compute units idle even after all their data has arrived.

Save memory without overlooking the new cost

Training memory holds weights, optimizer tracking information and activations required for gradients. It also holds buffers, temporary storage for transfers and computation. In our scenario, GPU 0’s inputs can require more space than its neighbours’. Record each rank’s peak: a comfortable average does not prevent the busiest device from running out of memory. Distinguish storage used by number arrays, called tensors, from memory the allocator reserves for current or future allocations.

Recomputation avoids keeping selected activations and recreates them when backpropagation needs them. Offloading temporarily moves them into system RAM, then brings them back to the GPU. The former adds computation; the latter adds host-device transfers. The Megatron Core technical report examines these trade-offs alongside communication and computation optimizations. Measure memory savings together with the effect on step time.

Overlap runs independent computation while communication progresses to avoid waiting. It can require retaining multiple batches and their buffers simultaneously. Megatron’s configuration reference notes that some expert-parallel overlap increases peak memory. Speeding up a step this way can therefore exceed available memory.

Capacity handling can also change learning. If an expert accepts only a fixed number of assignments, overflow assignments may be dropped from its computation. With a capacity of 1,024, our first expert would have 2,048 overflow assignments while others still had room. That removes this expert’s contribution for those assignments without necessarily removing the token from the entire model. MegaBlocks groups work into blocks suited to varying loads to avoid dropping assignments and padding every expert to a fixed size. A dropless recipe retains assignments, whose work still needs memory and execution time.

Choose parallelism by following the data

Spreading an expert group over more GPUs can reduce expert weights per device while sending activations to more destinations. Record where the group crosses from fast intra-node links onto the inter-node network.

Other forms of parallelism divide different work. Tensor parallelism assigns parts of a matrix operation to several GPUs. Pipeline parallelism places groups of layers on different devices. Data parallelism processes different batches on copies of the model and synchronizes their gradients. These dimensions must fit the GPU count and available memory. MoE Parallel Folding configures attention and expert parallelism separately because the blocks have different requirements.

Changing expert size and count, placement and communication groups simultaneously makes results harder to explain. Start from a supported configuration, test one coherent change at a time and retain traces and evaluations. A placement working on one machine does not predict its performance on eight networked machines.

Optimize time to a defined quality target

High throughput shows that the system processes tokens quickly. You still need to establish how many tokens it takes to learn what matters to you. In a hypothetical example, a recipe processing twice as many tokens per second but needing three times as many to reach the same result takes 1.5 times as long in training computation. Evaluations, checkpoints and interruptions add to total elapsed time.

Define success before selecting a winner. Validation loss measures error on texts that are not used to update weights, while additional evaluations check the tasks you care about. Set a resource limit and report tokens processed, elapsed time, GPUs used and restarts. GPU-hours multiply the GPU count by the duration of use; they do not make different GPU models equivalent.

MOSAIC formalizes this relationship by connecting a predictive quality model to a calibrated systems performance model. It does not replace measurements on your cluster. It explains why the best architecture under an operation budget can change under elapsed-time and memory constraints.

Our practical recommendation from reading the Hugging Face guide is to keep hardware and learning baselines before stacking optimizations. Investigate expert load when ranks wait, transfers when comparable loads still communicate slowly, and local kernels or memory when links no longer explain lost time. Once training finishes, serving an MoE presents another batch-size and networking problem, with different workloads and latency constraints.

Sources and method

Documentary analysis completed on November 8, 2026. We examined the Hugging Face handbook at revision ef70726e6f8b99d74298a16a4894340f288a58c2. Its environment pins Megatron revision d6316eb45bf1bfbb6f1f848c0b56e55166564d80. Documentation labelled latest was consulted on this date and may change.

Architecture and learning mechanisms are checked against Mixtral, Switch Transformers, OLMoE, Auxiliary-Loss-Free Load Balancing and Demons in the Detail. Systems mechanisms are documented by DeepEP, NCCL Tests, DCGM, MegaBlocks, Megatron Core and Parallel Folding. MOSAIC provides the quality-resource-time comparison framework.

The 0.75/0.25 mixture, four-GPU scenario, 32 MiB calculation and throughput-versus-data example are calculated illustrations. The distributed scenario assumes eight equally expensive experts, top-2 routing and two-byte values. They are neither network measurements nor throughput predictions. The diagnostic sequence and comparison protocol are our engineering recommendations, based on the cited mechanisms.