The Ultimate MoE Training Handbook

1. Introduction

The frontier is sparse.Reading time: 1-2 days. For the best reading experience, we recommend not using a mobile phone. On the Artificial Analysis index, 18 of the 20 highest-scoring open-weights models are mixtures of experts (the two exceptions, Nanbeige4.3-33B and Qwen3.8 27B at twelfth and thirteenth place, are dense)¹. They include the top model of each of the five leading labs: DeepSeek V4.1 Flash 1017, GLM-5.3, MiMo-V2.6-Pro, Ling 3.1 Flash 1019 and MiniMax-M3.1-Flash¹.

What an MoE offers lies in how it spends its inference compute. Each token runs through only a few of its experts, so the model can hold many more parameters than it uses on any one token: the small MoE you will train in chapter 2 has 6.86B parameters, of which 1.22B are active per token. Turning that inference saving into faster training is the hard part.

Our previous deep dive on training, the Smol Training Playbook, covered training a dense model but only touched on MoE design. Because the star of that book, SmolLM3, was a dense transformer, its training needed no technique specific to MoE, like expert parallelism. Still, we glanced at MoE in the playbook's "Going Sparse: MoE" section to cover its high-level design choices: sparsity, granularity, shared experts and load balancing. That section named the main reasons distributed MoE training is complex: expert placement, load balancing and all-to-all communication, but without diving into the particulars². This handbook starts from that list.

Going from a dense model to an MoE gives you three new challenges to manage: a router that decides which experts each token uses, a load that has to stay balanced across the experts, and, once the experts are split across GPUs, an exchange of tokens between those GPUs in every MoE layer. Each aspect brings settings a dense run never had and new ways for a run to fail. The chapters that follow take those settings one at a time and examine them in detail, with code you can run and, in most cases, direct results and data from our own cluster experiments.

A run that goes south will give you a symptom but not the reason behind it. The map below is a high-level view to help you go back to the source: start from how a run fails, find the first signal of bad behavior (the canary in the coal mine), and go read the part of the post that explains the possible cause. You can keep this map at hand for the day a run goes wrong.

Now, before we start, a quick note on the intended audience. Who is this handbook for? You will benefit from it if you have trained a dense model on a GPU cluster and want to train an MoE next. Chapter 2 sets up what you need and has you train that small MoE on a single node. If you are more curious about how MoEs work and what frontier models choose than about running one, skip ahead to chapter 3, Anatomy of MoE Training.

Once the experts sit on different GPUs, every MoE layer exchanges tokens between them, a pattern our dense run never used. Can you trust your cluster's links with it before you change a single line of the model?

2. Setup and your first run

2.0 What you need

Every command in this chapter was written for one kind of machine, so that everyone runs the same environment: a Slurm cluster of eight-GPU nodes. This section describes that setup, how much of the post you can follow with one GPU or none, and how to rent an eight-GPU node on Hugging Face Jobs. If your machines differ, you can still follow the explanations and adapt the commands.

The cluster the commands assume

The commands expect a Slurm cluster whose GPU nodes hold eight GPUs each. Containers start through Enroot and Pyxis, which let srun launch a task inside an image.^3,4 The per-node readiness check in 2.1 uses NVIDIA's DCGM diagnostics. Our runs used AWS P5 nodes connected by Elastic Fabric Adapter (EFA), each with eight H100 GPUs of 80 GB.

Most of the chapter fits on one node. The dense baseline and the first MoE each request a single node with eight GPUs of about 80 GB. Only the network checks need more: the all-to-all tests in 2.1 run on two, four and eight nodes, and the EFA image in 2.2 is verified on two.

The software comes from two pinned sources. The first is the NGC PyTorch image 26.06-py3, which 2.2 imports and extends with the cluster's network libraries. The second is a Megatron checkout at commit d6316eb45, which you mount into the container rather than build into it. You also need a filesystem that the CPU workers and the GPU nodes can both read, because 2.3 tokenizes the data on a CPU partition and the training jobs read the result from there.

Reading with one GPU, or none

You can get most of the post without a cluster. Chapter 3 takes the MoE layer apart: what each part does, what frontier models choose for it and the evidence behind each choice. You can read it end to end without running anything, and save its flags for when you have the hardware. Chapter 4 carries the same choices past pretraining, into serving and reinforcement learning.

In this chapter, the tokenization in 2.3 runs on CPUs. Its small sample runs on a single machine without Slurm, so you can produce and inspect the same files the training runs read. The training recipes in 2.4 and 2.5 assume eight GPUs of about 80 GB, and the MoE spreads its experts across all eight. With fewer GPUs, read them for the method: what stays fixed between the two runs, how the expert width is chosen and what to compare. The all-to-all tests in 2.1 and the EFA setup in 2.2 need several nodes to run as written.

Renting a node on Hugging Face Jobs

If you want to run the single-node parts without owning the hardware, Hugging Face Jobs runs a Docker image and a command on Hugging Face's machines, billed by the minute. It is open to any user or organization with a positive credit balance, and the monthly credits of a PRO, Team or Enterprise subscription can pay for jobs.⁵ The largest hardware options that hf jobs hardware lists have eight GPUs, as many as 2.4 and 2.5 use; h200x8, for example, holds eight H200s of 141 GB each.⁶

A job is defined by a Docker image from Docker Hub or a Hugging Face Space, a command and a hardware option, so running this chapter there means packaging the image from 2.2 as a Docker image and replacing the Slurm launcher with a Jobs command. A job stops after 30 minutes unless you set a different timeout.⁷ A first command only checks that the job sees all eight GPUs:

hf jobs run --flavor h200x8 --timeout 10m \
    pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel \
    python -c "import torch; print(torch.cuda.device_count())"

Whether you run the recipes on your own cluster or on a rented node, the machines you were given may not perform the way their specification says. How would you notice a GPU or a link that quietly underperforms?

2.1 Test your hardware

Before you start tuning a model, establish what the machines can do. With Megatron's alltoall dispatcher, each MoE layer sends tokens to the GPUs that hold their experts and returns the results through another all-to-all. The links between GPUs are therefore part of the work you need to test: a healthy all-reduce result alone does not tell you how quickly those token exchanges will run.

Run the checks below inside your job allocation, before launching training. With Pyxis, srun --container-image starts a container; the next section explains how to prepare the training image.³

Check each node

Start with a readiness check on every allocated node. NVIDIA's DCGM provides several diagnostic levels, with level 1 intended for readiness checks and the higher levels adding tests that help investigate failures. The command below runs one task per node so you can inspect each result before starting training.⁸

NNODES=$SLURM_JOB_NUM_NODES
srun --nodes="$NNODES" --ntasks="$NNODES" --ntasks-per-node=1 \
    dcgmi diag -r 1 --fail-early

If every node passes, proceed to the communication test. If a node fails, keep its diagnostic output and investigate it before proceeding. For example, a level 2 diagnostic reported thermal violations alongside failed PCIe and memory tests on three GPUs in one of our nodes.

Longer diagnostics belong in a separate maintenance window. NVIDIA describes levels 3 and 4 as administrator post-mortem tests, whereas level 1 is a readiness check.⁸

Prepare the diagnostic image

If your cluster already provides an nccl-tests image, set NCCL_TESTS_IMAGE to its path. Otherwise, on AWS with EFA you can import the public image maintained with the AWS NCCL test recipes. Choose a destination on a filesystem accessible to the allocated nodes:⁹

NCCL_TESTS_IMAGE=/shared/images/nccl-tests.sqsh
mkdir -p "$(dirname "$NCCL_TESTS_IMAGE")"
enroot import -o "$NCCL_TESTS_IMAGE" \
    'docker://public.ecr.aws#hpc-cloud/nccl-tests:latest'

Keep the imported SquashFS file with your baseline, since importing latest again can give you a different image. For a repeatable import, select a versioned tag from the AWS recipe in place of latest.^10,9

The AWS recipe also provides a Dockerfile if you need to build the image. For another fabric, use its supported CUDA, NCCL, MPI and network-plugin stack; NVIDIA's nccl-tests build instructions explain how to select those library paths. Multi-node tests require a build with MPI=1. The benchmark below assumes that MPI can launch through Slurm's PMIx integration, so match the image to the cluster's MPI setup.¹¹

Measure the collective you will use

Once the nodes pass, measure all-to-all bandwidth across the allocation. The command below launches one rank per GPU and exchanges an 8 GiB buffer per rank, matching the large-message test used for our baseline. Use the NCCL_TESTS_IMAGE path prepared above, adjusting the executable path if your image installs nccl-tests elsewhere.

srun --nodes="$NNODES" --ntasks="$((8 * NNODES))" \
    --ntasks-per-node=8 --gres=gpu:8 --mpi=pmix --cpu-bind=none \
    --container-image="$NCCL_TESTS_IMAGE" --no-container-mount-home \
    /opt/nccl-tests/build/alltoall_perf \
    -b 8G -e 8G -g 1 -c 1 -n 20

Check the rank listing before reading the result: an allocation of two eight-GPU nodes should list 16 ranks. Then inspect the busbw and #wrong columns for the tested size. Here is the out-of-place result from our two-node baseline, with the other columns omitted:

size (B)      busbw (GB/s)    #wrong
8589934592    77.14           0

The zero in #wrong means the correctness check passed for that result. The bandwidth gives you a reference for subsequent allocations with the same node count, GPUs per node and test settings. Keep the node list with it so that a later slow result can be traced back to the machines involved. In a separate two-node allocation, run with NCCL tuning variables set, the same message size reached 39.84 GB/s, so record which nodes and settings produced each baseline.

All-reduce remains useful as a comparison, but the two collectives can behave very differently on the same machines. That slower two-node test delivered 398.52 GB/s for all-reduce alongside its 39.84 GB/s all-to-all result. For an MoE that exchanges tokens across those nodes, the second measurement is the one that tests the relevant communication pattern.

Choose a baseline for each allocation size

As you add nodes, more of an all-to-all exchange crosses the inter-node network. Our multi-node sweep shows why the baseline must depend on node count: it reached 77.14 GB/s on two nodes, 39.25 GB/s on four, and 30.91 GB/s on eight. Compare each allocation with a baseline at its own node count, using the same message size and the same result column.

Before choosing a reference, repeat the check on a known-good configuration with the hardware, software and test settings held fixed. Use the resulting range to decide whether a new allocation is ready for training. A practical starting policy is to require 90% of the baseline bandwidth and a passing correctness check. With the two-node result above, the bandwidth threshold would be about 69 GB/s. Adjust the threshold to your cluster's observed variability. A fixed 33 GB/s threshold, for example, would accept the slow two-node result while rejecting the eight-node sweep result.

Keep this baseline unchanged while you investigate a slow allocation. If you also want to try NCCL tuning variables, compare the settings on the same nodes with the same test parameters. Our default and tuned runs used different node pairs, so their bandwidth difference cannot identify the effect of those variables.

Check for changes during training

For periodic hardware checks during training, use Megatron's --gpu-sniff-test-interval <n>. It runs before training and every n iterations (training.py#L4850-L4852, training.py#L4484-L4489), including an all-to-all over the expert-parallel group (gpu_sniff_test.py#L5-L21). Its outlier warnings identify ranks to investigate when a slowdown appears (gpu_sniff_test.py#L104-L118).

The diagnostic image has let you test the GPUs and their links. Your training image needs access to the same hardware, so let's prepare it and check its network libraries.

2.2 Build your image

A training environment has two parts to get right: the libraries that run GPU computations, and the network libraries that connect those GPUs across nodes. We start from an NVIDIA PyTorch container for the first part, then add the network support required by the cluster. Keeping Megatron in a separate checkout lets you change the training code while retaining the same image.

The examples here use Enroot and Pyxis on Slurm. The base image is NGC PyTorch 26.06-py3; the network-specific instructions apply to our AWS P5 nodes with EFA. On another cluster, use its supported network stack for that part of the image.^3,4

Import the base image

Enroot imports a container image into a SquashFS file that Pyxis can launch through srun. Choose an output path accessible from your compute nodes, then import the base image:⁴

IMAGE=/path/to/images/pytorch-26.06.sqsh
enroot import -o "$IMAGE" docker://nvcr.io#nvidia/pytorch:26.06-py3

Keep the Megatron checkout outside this image and mount it into the container when you launch training. The flags in this post target commit d6316eb45; checking it out explicitly makes the code you read alongside the examples match the code you run. The source notes list the image and commit used for each measurement.

Put this checkout, and later the data and tokenizer, on a filesystem that both the CPU workers and GPU nodes can read. In the commands below, SHARED_ROOT is that directory and PROJECT_DIR is your project directory, with the companion scripts cloned into its scripts/ folder.

export SHARED_ROOT=/path/to/shared/work
export PROJECT_DIR="$SHARED_ROOT/moe-blogpost"
git clone https://gist.github.com/nmntz/5de952d686f66d7bce5bce8cba9d7139.git "$PROJECT_DIR/scripts"
git clone https://github.com/NVIDIA/Megatron-LM.git "$SHARED_ROOT/Megatron-LM"
cd "$SHARED_ROOT/Megatron-LM"
git checkout d6316eb45bf1bfbb6f1f848c0b56e55166564d80

Verify GPU access

Before adding the cluster's network libraries, check that the base image can import PyTorch and Transformer Engine and see the GPUs. Run this inside an allocated eight-GPU node, using the IMAGE path from above:

srun --nodes=1 --ntasks=1 --gres=gpu:8 \
    --container-image="$IMAGE" --no-container-mount-home \
    python -c '
import torch, transformer_engine
print("torch", torch.__version__)
print("transformer_engine", transformer_engine.__version__)
print("device_count", torch.cuda.device_count())
x = torch.randn(8).cuda().sum()
print("cuda sum ok", x.item())
'

The stock image reported the following versions and device count. The final line of the command also prints a sum computed on the GPU, so inspect that result along with the imports before proceeding.

torch 2.13.0a0+8145d630e8.nv26.06
transformer_engine 2.16.0+4220403e
device_count 8

The CUDA result checks execution on the GPU within this node. For EFA communication across nodes, add the network plugin and its dependencies described below.

Add EFA support

On our AWS nodes, the network layer consists of the EFA userspace libraries and the aws-ofi-nccl plugin built against the NCCL inside the image. The build used EFA installer 1.49.0 and aws-ofi-nccl v1.19.2. Record both versions alongside the base image in your build recipe.

The EFA image build script creates a writable Enroot container from the base image, installs the EFA userspace libraries with --skip-kmod, and compiles the plugin against the image's NCCL. It then exports a new SquashFS image and removes the temporary container. Run it on an allocated compute node with Enroot and GPU access, choosing a new output path:

EFA_IMAGE=/path/to/images/pytorch-26.06-efa.sqsh
bash "$PROJECT_DIR/scripts/build_efa_image.sh" "$IMAGE" "$EFA_IMAGE"

Use EFA_IMAGE as the --container-image path for subsequent training launches. The script prepares the network layer; keep the following environment settings with that image so the job can find and select its libraries.

Here, the library paths refer to the copies installed inside the container, and NCCL_NET_PLUGIN=ofi selects the plugin by name. Passing a filesystem path as the plugin value caused NCCL to construct an invalid library name in our setup.

export LD_LIBRARY_PATH=/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:${LD_LIBRARY_PATH:-}
export NCCL_NET_PLUGIN=ofi
export FI_PROVIDER=efa
export FI_EFA_USE_DEVICE_RDMA=1
export NCCL_SOCKET_IFNAME='^lo,docker0,veth_def_agent'
export NCCL_DEBUG=INFO

The socket interface exclusion addresses the initial connection NCCL uses to set up communication. On our nodes, NCCL selected veth_def_agent, tried to connect to its link-local address, and failed with Connection refused. Excluding that interface allowed the launch to proceed. Adapt the exclusion list to your cluster's interfaces before reusing it.

Verify the network path

On your first multi-node training launch, inspect the NCCL initialization log before interpreting training performance. This checks the network path from the training image itself. For the EFA image described here, the line to look for is:

NCCL INFO Using network Libfabric

If you see Using network Socket instead, read the plugin-loading messages immediately above it. An invalid plugin name or a missing library version can leave NCCL using sockets. For example, the EFA_1.4 error identifies a missing library version to resolve before trying again.

The log confirms which transport the job selected; measuring its bandwidth remains a separate check. Record the training image and network settings alongside the diagnostic result from 2.1, keeping the two environments explicit when comparing their performance.

With the hardware checked and the image prepared, we can turn to the first training run. It starts with the data: turning raw text into token sequences and storing them in the format Megatron expects.

2.3 Tokenize

Megatron reads training text from pairs of indexed files. The .bin file stores token IDs, while the .idx file tells the loader where each sequence begins and ends. Preparing these files once lets the dense and MoE runs read the same tokens, so changing the model does not also change the data.

We will use FineWeb-Edu and the Llama 3.2 tokenizer. Datatrove's MegatronDocumentTokenizer writes the format directly; the remaining work is to choose an output directory, check a small sample, and submit the larger tokenization job.

Prepare the tokenizer

Keep the data and tokenizer under SHARED_ROOT, the shared directory from 2.2. The tokenization script runs locally for the sample and submits a Slurm array for the larger job.

export TOKENIZER_DIR="$SHARED_ROOT/tokenizer"
python3.11 -m venv "$SHARED_ROOT/tokenize-venv"
source "$SHARED_ROOT/tokenize-venv/bin/activate"
python -m pip install \
    'datatrove @ git+https://github.com/huggingface/datatrove.git@2c895af7a95010a96fa53e2dc77cc0ce92799e3f' \
    pyarrow tokenizers huggingface_hub

The original tokenizer repository, meta-llama/Llama-3.2-1B, requires access approval. If your account has access, run hf auth login and use that repository in the download command. Otherwise, unsloth/Llama-3.2-1B provides an ungated copy: our comparison found matching vocabulary and token IDs. Both tokenization and training will read the same downloaded tokenizer files.¹²

TOKENIZER_REPO=unsloth/Llama-3.2-1B
hf download "$TOKENIZER_REPO" \
    --include 'tokenizer*' --include 'special_tokens_map.json' --include 'config.json' \
    --local-dir "$TOKENIZER_DIR"

Datatrove reads tokenizer.json from this directory and appends <|end_of_text|> after each document. Megatron later receives the directory through --tokenizer-model, which also gives it the tokenizer's configuration.

Check a small shard

Tokenize a small sample on a CPU compute node before submitting the full array. This command reads up to 1,000 documents and keeps them all, writing the result to a separate directory so the sample cannot be mistaken for the training corpus.

python "$PROJECT_DIR/scripts/tokenize_first_run.py" \
    --tokenizer "$TOKENIZER_DIR/tokenizer.json" \
    --output "$SHARED_ROOT/data-smoke" --limit 1000

The output prefix is fineweb-edu_00000_tokens. Check it with Megatron's own reader inside the training image from 2.2, using an allocation where Pyxis can launch the container. Deactivate the tokenization environment first so its Python does not override the image's interpreter. Set TRAIN_IMAGE to EFA_IMAGE on the EFA setup, or to the prepared training image for your cluster.

deactivate
export TRAIN_IMAGE="$EFA_IMAGE"
export MEGATRON_PATH="$SHARED_ROOT/Megatron-LM"
srun --nodes=1 --ntasks=1 \
    --container-image="$TRAIN_IMAGE" \
    --container-mounts="$SHARED_ROOT:$SHARED_ROOT" \
    --container-workdir="$MEGATRON_PATH" --no-container-mount-home \
    env PYTHONPATH="$MEGATRON_PATH" python -c '
import sys
from megatron.core.datasets.indexed_dataset import IndexedDataset
data = IndexedDataset(sys.argv[1])
assert len(data) > 0
print("sequences", len(data))
for i in range(min(3, len(data))):
    tokens = data[i]
    assert len(tokens) and tokens[-1] == 128001
    print("length", len(tokens), "last token", int(tokens[-1]))
' "$SHARED_ROOT/data-smoke/fineweb-edu_00000_tokens"

You should see a nonzero sequence count and last token 128001 for each inspected sequence. That final token is the document boundary we asked the tokenizer to append. Reading the shard through Megatron also checks that Megatron can parse the index and read the first sequences from the token file before GPU training starts.

Submit the tokenization job

The sample checks the file format; the larger job prepares a corpus to reuse as your training runs grow. Submit it from the login node with the same virtual environment activated. Set CPU_PARTITION to your cluster's CPU partition; add --qos if the scheduler requires one. The script requests 12 CPUs and 24 GB of memory per task, and the virtual environment must be accessible at the same path on those workers.

source "$SHARED_ROOT/tokenize-venv/bin/activate"
export OUT_DIR="$SHARED_ROOT/data/fineweb-edu"
CPU_PARTITION=your_cpu_partition
python "$PROJECT_DIR/scripts/tokenize_first_run.py" \
    --tokenizer "$TOKENIZER_DIR/tokenizer.json" --output "$OUT_DIR" \
    --tasks 100 --limit -1 --sample-rate 0.074 --partition "$CPU_PARTITION"

Datatrove prints the submitted job ID. Use that ID with squeue -j <job-id> to follow the array, and inspect failed tasks under $OUT_DIR/logs before relaunching the same command.

Each task reads a share of the input files and keeps a random 7.4% of the documents before tokenizing them. Because documents have different lengths, this fraction does not fix the number of tokens in the output. Keep the task count, sampling rate and output directory unchanged when restarting an interrupted job. Datatrove uses completion markers under the logging directory to skip finished tasks.¹³

Pass the shards to Megatron

Wait for the array to finish before building the training manifest. The following check requires all 100 task-completion markers and expected shard pairs and writes the shard prefixes to data-args.txt, one per line. Each prefix names both files, so neither .bin nor .idx belongs in the manifest.

export DATA_ARGS="$OUT_DIR/data-args.txt"
python - <<'PY'
import os
from pathlib import Path
root = Path(os.environ["OUT_DIR"])
prefixes = [root / f"fineweb-edu_{rank:05d}_tokens" for rank in range(100)]
for rank, prefix in enumerate(prefixes):
    assert (root / "logs/completions" / f"{rank:05d}").is_file(), rank
    for suffix in (".bin", ".idx"):
        path = Path(str(prefix) + suffix)
        assert path.is_file() and path.stat().st_size > 0, path
Path(os.environ["DATA_ARGS"]).write_text("\n".join(map(str, prefixes)) + "\n")
print(f"wrote {len(prefixes)} prefixes")
PY
deactivate

Training will read this file through --data-args-path. That option supplies the same list as --data-path, but keeps a long set of shards out of the launch command. The launcher also sets --data-cache-path to a persistent directory for Megatron's sample indices, allowing later runs with the same dataset configuration to reuse them.

For the first runs we use --split 100,0,0, assigning the data to training. Reserve a held-out split before using the run to compare model quality.

With a readable dataset and a manifest pointing to its shards, you can launch the dense baseline. Its loss, gradient norms and step times will establish a reference before we introduce experts.

2.4 Dense baseline

A dense run gives you something concrete to compare with when the MoE behaves unexpectedly. Both models will use the same data, batch size and optimizer, so a loss spike or a slow step can first be checked against the simpler model. This baseline also exercises the training image and data loader together before routing enters the picture.

The model has 16 transformer layers, a hidden size of 2,048 and an 8,192-wide feed-forward network. With the 128,256-entry vocabulary, it contains about 1.22 billion parameters. The recipe targets one node with eight GPUs with roughly 80 GB of memory each. We will reuse the launch script for the MoE.

Launch on one node

We will use Trackio to record the training metrics. Trackio is a lightweight, open-source experiment tracker from Hugging Face. Its API mirrors wandb, so import trackio as wandb is usually the only change existing code needs. Runs are logged to a local database you own, and you can browse them in a local dashboard or sync them to a Hugging Face Space to share live curves with your team, with no account or paid plan required.¹⁴

First, create a Python environment on shared storage from inside the training image. It inherits the image's GPU libraries and adds the pinned tracking package:

srun --nodes=1 --ntasks=1 \
    --container-image="$TRAIN_IMAGE" \
    --container-mounts="$SHARED_ROOT:$SHARED_ROOT" --no-container-mount-home \
    bash -c 'python -m venv --system-site-packages "$1/train-venv" &&
        "$1/train-venv/bin/python" -m pip install "trackio==0.42.0"' _ "$SHARED_ROOT"
export TRACKIO_PROJECT=moe-first-run

The launcher uses the Trackio entry point to send Megatron's scalar metrics to one run per job. Both runs share the project above, with their data stored under $SHARED_ROOT/trackio. The adapter uses Megatron's existing scalar-logging hooks, so flags such as --log-memory-to-tensorboard control which metrics reach Trackio.

The first-run script groups the training flags into Bash arrays. It uses the training image, pinned Megatron checkout, tokenizer directory and data manifest prepared in the previous sections. Keep the mounted paths under SHARED_ROOT and set TRAIN_IMAGE to the prepared SquashFS image.

From the shared blog checkout, submit the dense job to your GPU partition. Keep the variables from 2.3 exported when you submit the job. The script records the image path, Megatron commit and data manifest with its output.

cd "$PROJECT_DIR"
GPU_PARTITION=your_gpu_partition
sbatch --partition="$GPU_PARTITION" scripts/first_run.sbatch dense

Slurm reserves one eight-GPU node. Inside it, a single srun task starts torchrun, which launches one training process per GPU. Look for the run output under $SHARED_ROOT/runs/dense-<job-id> and the training log at first-run-<job-id>.out in the submission directory.

The initial run stops after 120 iterations. Checkpoint saving and validation are disabled for this setup check. Keep the log and Trackio data when you move on to the MoE; they are the reference for the next section.

To inspect the curves, run the dashboard on a machine that can read the saved directory, after installing trackio==0.42.0 there. Select the dense and MoE runs in the project to compare them.

TRACKIO_DIR="$SHARED_ROOT/trackio" trackio show --project "$TRACKIO_PROJECT"

Understand the configuration

With tensor and pipeline parallelism both set to one, each GPU holds a copy of the dense model. Data parallelism distributes the batch across the eight GPUs, while the distributed optimizer and communication overlaps are enabled in both runs.

--tensor-model-parallel-size 1 --pipeline-model-parallel-size 1
--expert-model-parallel-size 1 --context-parallel-size 1
--use-distributed-optimizer --overlap-grad-reduce --overlap-param-gather
--micro-batch-size 8 --global-batch-size 512

A micro-batch contains eight sequences per GPU. Across eight GPUs, each micro-batch processes 64 sequences; eight gradient accumulation steps produce a global batch of 512. At a sequence length of 4,096, each optimizer step consumes 2,097,152 tokens. Keeping these quantities fixed will make dense and MoE step times directly comparable.

The learning rate follows a warmup, stable, decay schedule, selected by --lr-decay-style WSD. It warms up to 5e-4, holds that learning rate, then decays near the end of the configured token budget. The 120-iteration run ends during the stable phase; its stop condition does not shorten the learning-rate schedule.

Full activation recomputation keeps the starting configuration within the memory budget used for our baseline. With micro-batch size eight and the native cross-entropy implementation, the corresponding run without recomputation exhausted memory in the loss calculation. We keep recomputation in both models so that the first comparison does not mix a model change with a memory-saving change.

--recompute-granularity full --recompute-method uniform --recompute-num-layers 1
--cross-entropy-loss-fusion --cross-entropy-fusion-impl native

Keep the dense run's loss, gradient norm and step time in Trackio as your reference. When you replace the feed-forward layers with experts, these measurements will help you see how the change affects training.

2.5 Your first MoE

The dense model applies the same feed-forward network to every token. In the MoE, a router chooses eight of 64 smaller experts for each token, allowing the model to hold more parameters than it uses on any one token. We choose the expert width so that the eight selected experts contain the same number of feed-forward parameters as the dense layer they replace.

This gives us a 6.86-billion-parameter model with about 1.22 billion active parameters per token. The extra capacity comes with a new cost: the tokens have to reach their selected experts, and the expert outputs have to return to the processes that need them.

Replace the feed-forward layers

The launcher selects the MoE configuration when you pass moe. Submit it using the same exported paths and GPU partition as the dense run, keeping the image, checkout, tokenizer and data manifest unchanged:

sbatch --partition="$GPU_PARTITION" scripts/first_run.sbatch moe

Expert parallelism increases from one to eight, distributing the 64 experts across the node's eight GPUs. Each GPU holds eight experts, while the router selects eight experts per token from the full set.

--expert-model-parallel-size 8 --expert-tensor-parallel-size 1
--num-experts 64 --moe-ffn-hidden-size 1024 --moe-router-topk 8
--moe-token-dispatcher-type alltoall --moe-grouped-gemm
--moe-router-load-balancing-type aux_loss --moe-aux-loss-coeff 1e-3
--moe-router-dtype fp32 --moe-per-layer-logging

alltoall moves the routed tokens between GPUs, and grouped GEMM executes the local experts' matrix multiplications. The auxiliary loss encourages the router to distribute work across experts, with a coefficient of 1e-3. We also compute router logits in fp32 and enable per-layer logging, giving us the initial settings and measurements that chapter 3 will examine.

Everything else stays in the common configuration, including full activation recomputation. In particular, the dense --ffn-hidden-size 8192 setting remains in the script, while --moe-ffn-hidden-size 1024 supplies the width of each expert.

Count active and total parameters

A bias-free SwiGLU expert has three weight matrices. With hidden size 2,048 and expert width 1,024, their combined size is 3 × 2048 × 1024 = 6,291,456 parameters. Selecting eight experts therefore uses 50,331,648 feed-forward parameters per token, matching the dense layer's 3 × 2048 × 8192 parameters.

All 64 experts together contain eight times as many feed-forward parameters. The whole model grows by a smaller factor because attention, embeddings and the other non-expert parameters stay the same size, and each layer gains a small router matrix. Summing those components across 16 layers gives 6.86 billion total parameters and 1.221 billion active parameters, compared with 1.219 billion in the dense model.

The result is about 5.6 times as many total parameters at nearly the same active parameter count. Matching active parameter counts does not guarantee equal step time: distributing tokens among experts changes how the GPU work is organized and introduces communication.

Keep a dense reference

An MoE with one expert and top-1 routing computes the same function as the dense model it came from, so it should train like it. In our run, a 1-expert MoE upcycled from a dense checkpoint and continued beside the dense run stayed within 0.00087 of it over 15 aligned steps. Start the MoE from its own initialization instead and the comparison says little, since the two runs sit 2.72 apart at iteration 4. The MoE side also logs a load-balancing loss of exactly 1 and about 32,770 zeros in its gradient against at most 17 for the dense run: the router weights, which get no gradient with a single expert. The parity run upcycles the dense checkpoint, a path 3.6 returns to.

At the pin, that upcycling path ran for us only with two local fixes to training.py. The first passes the dense model to the checkpoint loader under the keyword model (training.py#L3074-L3078), while load_checkpoint names that parameter ddp_model. The second concerns the dense model built for the load: the upcycling code switches the dense shape on in args (training.py#L3056-L3062), but the model builder reads the model configuration stored at startup (training.py#L2917-L2926), so that model keeps the MoE layer layout and the load stops on a missing pre_mlp_layernorm key.

Deterministic mode makes a run repeatable, so that a gap between two runs means something. At the pin, --deterministic-mode requires cross-entropy loss fusion, TP communication overlap and the fused aux-loss kernel to be off (determinism.py#L28-L37, determinism.py#L195-L210), and accepts NVTE_ALLOW_NONDETERMINISTIC_ALGO only at 0 (determinism.py#L85-L90); it sets NCCL_ALGO=Ring, NCCL_EP_HT_EM_AG_SCAN_MODE=1, NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 and CUBLAS_WORKSPACE_CONFIG=:4096:8 itself when they are unset (determinism.py#L40-L50). The 2.4 launcher passes --cross-entropy-loss-fusion and exports NVTE_ALLOW_NONDETERMINISTIC_ALGO=1, so drop both before adding the flag. With these settings, two runs of this MoE shrunk to micro-batch 2, global batch 64 and 12 iterations printed the same lm loss, load-balancing loss, grad norm and gradient zeros at every iteration; without them the lm loss differed from the first iteration on, by up to 0.00063.

Switch on deterministic mode before you compare two runs or chase a NaN.

Keep the dense run and its configuration as a reference as you develop the MoE. Whenever a change makes training slower or the loss stops improving, return to the dense run in Trackio and compare step time, loss and gradient norm at the same batch size and sequence length. A problem that appears in both runs points you toward their shared setup; one that appears only in the MoE gives you a reason to investigate routing and expert execution.

The dense reference helps you decide where to begin an investigation. Chapter 3 opens the MoE layer piece by piece, introducing the measurements you need to understand its behavior.

3. Anatomy of MoE Training

The MoE from 2.5 added parts the dense run never had: a router that picks eight experts for each token, experts that can end up with very different amounts of work, and an all-to-all that moves tokens between GPUs. Each of them can slow training or hurt the loss, and a gap against the dense reference does not, on its own, tell you which part caused it.

The figure below draws one MoE layer at a smaller scale than ours, with 16 experts and top-4 instead of 64 and top-8, plus a shared expert that ours does not have. It shows the tokens of a batch, the router, the experts and the sum that combines their outputs. Each section then opens on its own view of the same drawing, zoomed onto the part it covers, so you can skip to the part you need. A few views also carry a result from their section, such as the rogue experts that Step 3.5 Flash's training loss did not reveal (3.6)¹⁵.

After the monitoring setup in 3.0, we start inside the router, with how its scores become the weights of the chosen experts and whether tokens choose experts or experts choose tokens. Then come uneven load across experts, how many experts to use and how large, which layers get experts at all, and how to start a run and keep it stable. We end with where a training step spends its time on the GPUs.

3.0 Monitoring: turn it on once

Before changing any setting, turn on what tells you whether it worked. You do not need to know what every number means yet: each section introduces a metric when it first needs it (z-loss in 3.1, MaxVio in 3.2, timings and memory in 3.7). Appendix A.1 gives the capture commands, calculations and limits of each measurement.

What to turn on

The 2.5 run already logs the router losses per layer (--moe-per-layer-logging), so you can see which layer misbehaves (3.1, 3.3, A.1). Keep the Trackio launcher from 2.4 and add these flags to the 2.5 flags:

--save <dir> --save-interval <m>
--save-tokens-per-expert-interval <n>
--log-throughput --log-params-norm
--log-num-zeros-in-grad --log-max-attention-logit

Megatron writes the token counts and the activation dump below under --save (training.py#L3398-L3401) and asks for --save-interval beside it (arguments.py#L1321-L1323), so the run also saves checkpoints there. The next flag records how many tokens each expert received (3.3, A.1). The other four track the run's health: throughput, the parameter norm, the zeros in the gradient and the largest attention logit (3.6, A.1). The scalar metrics appear in the same Trackio project as your baseline, while routing captures stay in their own files. The NaN check is already on (arguments.py#L2825-L2827).

What to capture when you need it

Two more measurements are bounded captures: run them when a question calls for them.

The routing trace records which experts each token picked (3.1, 3.3, A.1). Its hooks on the routers do not fire under CUDA graph replay (router_trace.py#L23), and by default it traces every iteration (training_config.py#L363-L367), so turn graphs off and bound it:

--moe-routing-trace-path <dir> --moe-routing-trace-max-training-iters <n>
--cuda-graph-impl none

The activation dump saves the experts' activations under the same --save directory, so you can compute per-expert activation norms offline (3.4, A.1):

--save-activations-interval <n>

With monitoring on, let's start where every token starts: the router. It picks the experts, but how much does each one count in the token's output?

3.1 Router scoring: how logits become weights

The router turns each token into one logit per expert, and in plain top-k routing the top-k logits pick the experts the token goes to. Between those logits and the weights that mix the experts' outputs sit four choices. The gating function, the output scale (a constant multiplied onto the weighted sum of the chosen experts' outputs) and the sharpness (how peaked the router's scores are, so how much the top expert dominates) change how much each chosen expert counts in the token's output, and how large that output is. In plain top-k routing, for a given token's logits, none of them changes which experts win: each keeps the experts in the same order, or applies after the top-k.

Precision is the exception. A bf16 router can round two close logits to the same value and so change the winners, which makes an fp32 router the one setting to always use here.

The table below shows what ten current open MoEs chose for each of the four, one row per model.

how current open MoEs set up their routers, one column per part of this section. Values come from each model's config.json and modeling code, training precision from its tech report; √k is computed, and shown wherever the weights sum to 1.

Let's start with the gating function. If it never changes the winners, what does it change?

Gating function

Softmax, sigmoid and sqrt(softplus) rank the experts the same way, so what differs is the weight each chosen expert gets. The figure below raises one token's logits from an unsure router to a confident one: softmax makes the top expert count more and more, while sigmoid saturates at large logits, so confident experts end up with equal weights. Gating functions keep the ranking when selection uses their scores alone. The expert bias of 3.3 is added to the gated scores before selection (moe_utils.py#L936-L938), so how far it moves an expert depends on the size of the scores, and a different gating function can change the top-k. Group-limited routing (3.7) can too, because it picks groups by summing their best scores (moe_utils.py#L712-L715).

In the figure, sqrt(softplus) lands between softmax and sigmoid: for positive logits, the ratio of two weights tends to the square root of the ratio of their logits as the scale grows. DeepSeek-V4 switched to sqrt(softplus) from sigmoid²⁷. In Megatron, the expert bias of 3.3 needs sigmoid or sqrt(softplus): it refuses the bias with softmax at startup (transformer_config.py#L3139-L3147), and its softmax path never adds it. Wang et al., who proposed the bias²⁸, explain why it fits softmax badly: softmax's normalization makes the gap between two experts depend on every other expert's score, so the bias is harder to tune. They still balanced a softmax router with it²⁸. The auxiliary loss works with all three (moe_utils.py#L998-L1014).

913if score_function == "softmax":914if use_pre_softmax:915scores = torch.softmax(logits, dim=-1, dtype=torch.float32)916if precomputed_indices is not None:917top_indices = precomputed_indices918probs = torch.gather(scores, dim=1, index=top_indices)919else:920probs, top_indices = compute_topk(scores, topk, num_groups, group_topk)921else:922if precomputed_indices is not None:923top_indices = precomputed_indices924scores = torch.gather(logits, dim=1, index=top_indices)925else:926scores, top_indices = compute_topk(logits, topk, num_groups, group_topk)927probs = torch.softmax(scores, dim=-1, dtype=torch.float32)

928elif score_function in ("sigmoid", "sqrtsoftplus"):929if score_function == "sigmoid":930scores = torch.sigmoid(logits.float())931else:932scores = torch.nn.functional.softplus(logits.float()).sqrt()933if precomputed_indices is not None:934top_indices = precomputed_indices935scores = torch.gather(scores, dim=1, index=top_indices)936elif expert_bias is not None:937scores_for_routing = scores + expert_bias.float()938_, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)939scores = torch.gather(scores, dim=1, index=top_indices)940else:941scores, top_indices = compute_topk(scores, topk, num_groups, group_topk)

3139if self.moe_router_enable_expert_bias and self.moe_router_score_function not in (3140"sigmoid",3141"sqrtsoftplus",3142):3143raise ValueError(3144"Expert bias for aux-loss-free routing only supports 'sigmoid' and 'sqrtsoftplus' "3145"score functions. Please set --moe-router-score-function to 'sigmoid' or "3146"'sqrtsoftplus', or unset --moe-router-enable-expert-bias."3147)

Sigmoid gives the chosen experts equal weights once all their logits are large and positive, which is where the simulation's scaled-up token ends: a confident router with the flat gate that Sharpness returns to. To see whether a run's router gets there, record the routing trace with its logits on (A.1) and follow two numbers per layer: the top expert's weight, and the routing entropy, which is low when one expert takes most of the weight and high when the weights are spread.

Use sigmoid (or sqrt(softplus)) when you balance with the expert bias (3.3), since Megatron refuses the bias with softmax. Otherwise use softmax, Megatron's default.

The router's weights do two jobs. They pick the experts, and they are also the numbers that mix the experts' outputs: the layer returns each chosen expert's output times its weight, summed. The gating function decides how much each chosen expert counts against the others. It does not decide how much they count together: with softmax, the top-k weights sum to less than 1, by a different amount for each token, so the size of the layer's output moves with them. So how big should that combined output be?

Output scale

Two settings decide how big the experts' combined output is. The first is renormalization: rescale the k chosen weights so they sum to 1, and the output's size no longer moves with their sum. In Krajewski et al. renormalizing lowered validation loss from 2.219 to 2.183 at granularity 8Granularity 8: experts 8 times smaller, 8 times as many, and 8 times as many chosen per token, so parameters and compute stay the same; 3.4 compares granularities.²⁹. Megatron renormalizes by default (moe_utils.py#L913-L942), and --moe-router-pre-softmax turns renormalization off. Softmax at top-1 requires that flag: with one expert chosen, softmax over its single logit is always 1, so the router gets no gradient (transformer_config.py#L3213-L3221).

The simulation sweeps toy tokens from an unsure router to a confident one, with uncorrelated experts of equal RMS: without renormalization the routed output's size follows the top-k probability mass, and renormalizing keeps it between 1/√k of one expert's RMS and one expert's RMS.

The second setting is the routed scaling factor, the constant multiplied onto the result. It is needed because, even renormalized, the size of the routed output depends on how much each expert counts. With all the weight on one expert, the output has that expert's root mean square (RMS). With equal weights on k uncorrelated experts of equal RMS, it drops to 1/√k of that, and a routed scaling factor near √k restores it. DeepSeek-V3 uses 2.5 at k = 8³⁰, and Kimi K2 2.827 ≈ √8. Shipped models disagree on its value: in the table above only GLM-5.3 and Step-5-Preview stay near √k, and Kimi went from √8 in K2 to 1.0 at k = 16 in K3²⁰. Scaling the routed output also changes its weight against the shared expert's output (3.4).

The figure's last row scales the expert's input instead of its output: Llama 4 sends each token to one expert and multiplies the token by that expert's sigmoid score before the expert runs³¹. Its SwiGLU experts have no biases, so the output shrinks roughly as the score squared, and the score also softens the SiLU gate.

942probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores943else:944raise ValueError(f"Invalid score_function: {score_function}")945946if scaling_factor:947probs = probs * scaling_factor

3213if (3214self.moe_router_topk == 13215and self.moe_router_score_function == "softmax"3216and not self.moe_router_pre_softmax3217and self.moe_router_load_balancing_type != "sinkhorn"3218):3219# Requires applying softmax before selecting the top-k when k is 1,3220# since softmax on a [num_tokens, 1] would yield a zero gradient.3221raise ValueError("Please use --moe-router-pre-softmax when topk is 1.")

Without renormalization, the size of the routed output follows the top-k probability mass from token to token, the setting that gave the higher loss above, so the number to watch is how large the routed output is for each token, the root of the sum of its squared weights. Compute it offline from the routing trace (A.1).

Keep renormalization on at top-k > 1, and use --moe-router-pre-softmax for softmax at top-1, or the router gets no gradient. The routed scaling factor has no settled value: √k restores the routed output's RMS only for equal weights on uncorrelated experts, and shipped models disagree on its value.

Both the weights and the output's size depend on how peaked the logits are. A router can get too sharp, with one expert taking all the weight, or too flat, averaging its top-k. What sets how peaked it is?

Sharpness

With softmax, a linear router gets sharper as the token's norm or the router's weights grow: the logits spread apart, and the top expert takes more of the weight. The figure below compares three ways to control that against the linear router. Logit soft-capping compresses mostly the logits near its cap. Normalizing the router input removes the part that comes from the token's norm. Gating logit normalization standardizes the logits across experts³², which sets their spread to a chosen scale.

The plain linear router in the figure sharpens as the token's norm or the router matrix grows, and soft-capping, which Grok 2 applies at 30 without renormalizing³³, pulls the large-norm token's top weight back only once its logits near the cap. Normalizing the input, as Gemma 4 does¹⁷, routes both tokens alike but still sharpens as the matrix grows, and gating logit normalization is unaffected by both. Skywork-MoE used gating logit normalization against a gate so flat that the layer averages its top-k experts³⁴. With it, training loss came out moderately lower³⁵. Normalizing the router input was ablated only together with the expert embeddings³⁶.

1150# Apply input jitter1151input = self.apply_input_jitter(input)1152logits = self.gating(input)

824if self.config.moe_input_jitter_eps is not None:825eps = self.config.moe_input_jitter_eps826if self.input_jitter is None:827self.input_jitter = torch.distributions.uniform.Uniform(828torch.tensor(1.0 - eps, dtype=input.dtype, device=input.device),829torch.tensor(1.0 + eps, dtype=input.dtype, device=input.device),830).rsample831return input * self.input_jitter(input.shape)

A flat gate, one that averages its top-k experts and makes the model learn worse, shows up as Max1/Max2, the top expert's weight over the second's, falling toward 1, as it did in Skywork-MoE without normalization³⁵. Follow it per layer together with the top expert's weight, both from the routing trace (A.1).

If the gate goes flat, gating logit normalization fixed it in Skywork-MoE. On the router input, Megatron offers Switch-style multiplicative input jitter, --moe-input-jitter-eps.

Sharpness lives in the gaps between logits, and softmax ignores a shift shared by all of them. bf16 does not: if every logit drifts upward together, the values bf16 can represent get farther apart, and close experts can tie. Do those ties change which experts win?

Precision

Top-k orders tied experts in a way PyTorch does not specify³⁷. The figure below simulates routing the same tokens with 64 experts and top-8 through a bf16 router, with and without z-loss, a penalty on the logits' logsumexp. Its simulated logits drift up by one shared offset, and it counts the tokens whose experts differ from what an fp32 router picks.

As the simulated offset rises, the bf16 router's logits land on a coarser grid and more tokens get different experts than under fp32. Z-loss lowers the logsumexp by pushing the logits negative. At top-8 of 64 the 8th logit sits about 3.5 below the logsumexp, so in the simulation it ends near −2.2 with z-loss and +3.6 without, two sizes bf16 rounds in steps of the same length, and about as many tokens change experts; only the fp32 router has no flips, since fp32 rounds 65,536 times more finely than bf16³⁸. DeepSeek-V3.2, GLM-4.5 and Kimi K2.5 compute the router in fp32 in their inference code³⁹. The case for fp32 rests on the training evidence, which is about stability. In Switch Transformers, a bf16 router diverged while an fp32 router trained at bf16 speed⁴⁰. In ST-MoE, the router z-loss kept every run stable⁴¹. Labs rarely say what precision they train the router in: in the table above, 2 of the 10 reports mention it, and 5 of the 10 run the router's matrix multiply in bf16 in their inference code. How far below bf16 the rest of the layer can go, while the router keeps its precision, is the question 3.7 Precision returns to.

124# Convert to specified datatype for routing computation if enabled125router_dtype = input.dtype126if self.config.moe_router_dtype == 'fp32':127router_dtype = torch.float32128elif self.config.moe_router_dtype == 'fp64':129router_dtype = torch.float64130logits = router_gating_linear(input, self.weight, self.bias, router_dtype)

1424if te_general_gemm is not None and router_dtype != torch.float64:1425# cuBLASLt's non-FP8 bias epilogue expects bias and output to have the same1426# dtype. Router parameters may be BF16 while router logits are FP32, so cast the1427# small bias vector before passing it to TE.1428gemm_bias = bias.to(router_dtype) if bias is not None else None1429output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=gemm_bias)[0]1430elif bias is None:1431output = torch.mm(inp.to(router_dtype), weight.to(router_dtype).t())1432else:1433output = torch.addmm(1434bias.to(router_dtype), inp.to(router_dtype), weight.to(router_dtype).t()1435)

171logsum = torch.logsumexp(logits, dim=-1)172z_loss_values = torch.square(logsum)173174if padding_mask is not None:175# Invert padding_mask: True (padding) -> 0, False (valid) -> 1176valid_mask = ~padding_mask177# Only compute z_loss for valid (non-padding) tokens178z_loss_values = z_loss_values * valid_mask179# Compute mean over valid tokens only180num_valid_tokens = valid_mask.sum()181z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff182else:183z_loss = torch.mean(z_loss_values) * z_loss_coeff184return z_loss

761if self.config.moe_z_loss_coeff is not None and self.training and torch.is_grad_enabled():762# Skip Z loss calculations when using torch.no_grad() or checkpointing.763moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()764z_loss = z_loss_func(logits, moe_z_loss_coeff, padding_mask=padding_mask)

A bf16 router can make training diverge, as it did in Switch Transformers, so two signals are worth following: the z-loss, which tracks how far the logits drift up, and the flip rate, the share of tokens whose experts change when the router is recomputed in fp32. Megatron logs the z-loss when --moe-z-loss-coeff is set, already divided by the coefficient (router.py#L805-L807), so its square root is the RMS logsumexp: our run logged 20.9, an RMS logsumexp of about 4.6. The flip rate needs a matched capture of the same inputs and current router state; A.1 provides a bounded live hook, since the trace's initial weight snapshot does not track later training updates.

Always set --moe-router-dtype fp32. Add --moe-z-loss-coeff 1e-3, ST-MoE's value, to a softmax router whose logits drift up. DeepSeek-V3's report never mentions z-loss for its sigmoid router⁴².

What to watch

So far we have shaped how a token's logits become weights. The gating function sets how much each chosen expert counts, renormalization and the scaling factor set how big their combined output is, and sharpness and precision set how peaked and how exact the scores are. All of it assumed one setup: a learned router, and the token picking its top k. Both can change: experts can pick their tokens instead, and a fixed table can replace the learned choice. Who should make the choice, the token or the expert, and does the choice need to be learned at all?

3.2 Routing strategy: who selects, and how

Token-choice top-k is the default: each token scores every expert and keeps its k best. This section looks at two ways to change that. The expert can pick its tokens instead of the token picking its experts, or a fixed hash can pick, with no learned router at all.

How current open MoEs pick their experts, and which ones look up a table by token id, with Zhou et al.'s research model as an expert-choice comparison. Current-model settings come from their config.json and modeling code; the expert-choice row comes from the paper⁴³.

Before asking whether the choice needs a learned router, we have to settle who makes it, because with token choice many tokens can want the same experts and loads skew. Should each token keep picking its experts, or should each expert pick its tokens instead?

Who selects

With token choice, many tokens can want the same experts, so loads skew. Expert choice flips the roles: each expert takes its best tokens, so every expert gets the same load⁴³. The catch is future token leakage⁴⁹: whether a token gets an expert depends on the tokens after it. At inference those tokens do not exist yet, so routing differs from training⁵⁰.

Expert threshold, the third rule in the figure below, keeps loads near even without looking at future tokens⁵⁰. In exchange, each token gets a variable number of experts. In a 575M-parameter model it reached a lower loss than token choice, with or without balancing (2.844 against about 2.89)⁵¹.

In the figure, the sequence leans toward expert 1, and token choice, drawn with no balancer, overloads it. Expert choice balances the loads, but later tokens push out earlier selections, while expert threshold, whose cutoffs come from earlier batches that lean the same way, keeps the loads near the target without revising any selection, at the price of a varying number of experts per token.

891# Sorting top-k turned off during inference892return torch.topk(893scores,894k=topk,895dim=1,896sorted=torch.is_grad_enabled() or is_batch_invariant_mode_enabled(),897)

The skew token choice leaves behind makes the step slow. It shows in the token counts, tokens per expert per layer, which --save-tokens-per-expert-interval logs (A.1) and which read best as MaxVio, how far the busiest expert's load sits above the meanMaxVio, short for maximal violation: (busiest expert's load minus the mean load) divided by the mean load. 0 is perfectly even; 1 means the busiest expert gets twice the average. Wang et al., who proposed the expert bias of 3.3, define it.⁵².

Use token choice, the routing Megatron trains, with --moe-router-topk setting k.

Whether the token or the expert does the picking, the choice so far always comes from a learned router, so the routing itself has to be trained. Does routing really need to be learned, or could a simpler rule pick the experts?

Hash layers

In DeepSeek-V4's first 3 MoE layers, a fixed table maps each token id to its experts⁵³, and the router is kept only for the weights⁵⁴. The table is a fixed hash of the token id²⁷, and in our probe it spreads frequent tokens no better than a random table. There is nothing to learn for the selection, and hash layers matched Switch and BASE layersLewis et al.'s balanced assignment of experts, which routes by solving a linear assignment that gives every expert an equal number of tokens. in models up to 4.5B⁵⁵. The figure below puts learned and hash routing side by side, above a third lookup by token id: Engram, which hashes each token's last 2 to 4 tokens into tables of embeddings and adds the rows it reads to the hidden state⁵⁶. DeepSeek-V4.1-Flash keeps Engram at two layers, but picks every layer's experts with its learned router^57,58.

The figure's two routers treat the word "the" differently: the learned router selects each one by its context, while the hash layer sends every "the" to the six experts its table row names, even past an expert with a higher score. Both rows weight their six experts to sum to V4-Flash's routed scaling factor of 1.5⁵⁹, and the figure measures how alike the four "the"s' expert sets are by their mean Jaccard similarity. What the hash layer gives up is context: the same token id always selects the same experts, whatever surrounds it. It also scales worse than learned routing. Clark et al. varied how many of a model's layers are MoE layers. With a single MoE layer, hash routing reached the lowest loss of the routing methods they compared, but as more layers became MoE layers it stopped improving, while their learned router kept getting better⁶⁰. And no bias can rebalance it⁵⁴: frequent tokens can overload their assigned experts, and nothing moves them. V4.1-Flash dropped hash layers and kept only Engram, which picks no experts⁴⁵.

In Megatron, --moe-num-hash-layers makes the first MoE layers hash layers (router.py#L304-L308). They read each token's experts from a token-to-expert table (tid2eid) with one row per vocabulary entry, and if you do not load one, Megatron fills it with a round-robin placeholder (router.py#L313-L329). Since the expert bias is off in these layers, nothing corrects a bad table during training, so the table matters. Roller et al. compared ways to build it: a random table and one that greedily spreads the most frequent tokens over the emptiest experts reached about the same perplexity, with the second better balanced, while grouping similar tokens by clustering their embeddings did clearly worse⁶¹.

304self.is_hash_layer = (305not self.is_mtp_layer306and hash_moe_layer_threshold > 0307and layer_number <= hash_moe_layer_threshold308)

322token_ids = torch.arange(self.config.hash_moe_vocab_size, device=self.weight.device)323expert_offsets = torch.arange(self.topk, device=token_ids.device)324self.tid2eid = ((token_ids[:, None] + expert_offsets) % self.num_experts).to(325torch.int32326)

898flat_ids = input_ids.T.reshape(-1)899assert flat_ids.numel() == logits.shape[0], (900f"input_ids contains {flat_ids.numel()} tokens, but router logits contain "901f"{logits.shape[0]}."902)903# Token IDs address the real tokenizer vocabulary; reject out-of-range IDs,904# including negative IDs, without clamping them to another token's experts.905default_top_indices = self.tid2eid.index_select(0, flat_ids).long()

914def _compute_hash_topk(scores, topk, num_groups=None, group_topk=None):915del topk, num_groups, group_topk916return scores.gather(1, default_top_indices), default_top_indices

A table that overloads some experts makes the step slow, and since nothing rebalances it during training, check it early in the run by comparing the tokens per expert in the hash layers with those in the learned layers, from the token counts that --save-tokens-per-expert-interval writes (A.1).

Use it only to reproduce DeepSeek-V4's first layers. Load a fixed hash table as V4 does²⁷, or a frequency-balanced one, instead of Megatron's round-robin placeholder.

What to watch

Expert choice and expert threshold keep the load even or close to it, and no bias can move a hash table's load. Megatron trains learned token choice, where nothing so far stops many tokens from crowding onto the same few experts. Crowded experts are the first real failure mode of an MoE, so how do we keep tokens spread evenly over the experts while the router keeps learning?

3.3 Load balancing

When many tokens pick the same experts, two things go wrong. The GPU holding the busiest experts has the most work, so the other GPUs wait for it and it sets the step time, and if the router keeps sending tokens to the same few experts, the others get few tokens and their parameters go mostly unused. The first cost is the easier one to measure: with routing forced even, the same model ran 1.5x faster at init.

Forcing the routing even only shows the speed to aim for, since it throws away what the router learned. To get close to it while the router keeps learning, there are two families of fixes, an auxiliary loss and a loss-free bias, and they balance the load in different ways. The aux loss adds a term that rewards an even load, so the router's gradient pushes its probabilities toward spreading tokens, at the price of a gradient that competes with the LM loss. The bias changes only which experts are selected. Whichever you pick, the scope it balances over (a sequence, a micro-batch or the global batch) decides what "balanced" means.

Balancing recipes reported by model authors. Auxiliary losses and selection biases can be used together. Scope describes the tokens whose loads are balanced; the MiMo row concerns RL rather than pretraining.

Before comparing balancers, we need to know how much speed an uneven load costs us, since every balancer below should be judged against the throughput of routing that is forced to be even. Where does that 1.5x come from, and how do we measure it on our own setup?

Cost of imbalance

To see what imbalance costs, we routed the same model three ways. Forced-even routing, Megatron's --moe-router-force-load-balancing, replaces the router's logits with random ones so that every expert gets about the same load, which makes it the best case, at 149.4 TFLOP/s.

Forced-biased routing, --moe-router-force-biased, adds a shared random bias of growing size to those random logits, so a few experts get more and more of the tokens: a dial for how uneven the load is. At the far end of the dial throughput falls to 100.5 TFLOP/s, and a replay of that routing, whose aux loss matches the logged one, has the busiest expert receive every token, a MaxVio of 7, as uneven as top-8 routing over 64 experts can get. It can get worse: in the replay the busiest experts are spread over several GPUs, so the busiest GPU holds only about twice the mean load, and had they all sat on one GPU, that GPU would have done most of the expert work.

The learned router, the one you would train with, starts far from that best case: a freshly initialised router ran its first 40 iterations as slowly as the worst forced case. The figure below compares the three. Training has to close that gap without forcing the router's choices, and the rest of this section compares the balancing methods by how close they bring a learned router to forced-even throughput.

In the figure, throughput falls with each larger forced bias, as a few experts take more of the tokens, and the learned router starts at the level of the most uneven forced run. The busiest experts set the step because the all-to-all hands each GPU every token routed to its local experts, 8 of 64 in our runs, and the next all-to-all waits for the slowest GPU. A GPU's load is the sum over its local experts, so it can be more balanced than the per-expert load. That is why a balance loss can also work per GPU rather than per expert: DeepSeekMoE's device-level balance loss pools each GPU's experts into one term⁶⁹, and the figure that follows compares it with the expert-level loss.

Trained on one toy batch with two experts per GPU, one of them among the four the LM favours, both losses balance the GPUs, and the device-level loss pulls the router less far from what the LM wants, since it sees only each GPU's total, which that placement keeps closer to even. The catch is that the experts within each GPU stay uneven, so one expert can still take most of its GPU's tokens. DeepSeekMoE gives the same reason for weighting the two losses unevenly: strict expert-level balance would cost model quality, so it sets a larger device-level factor and keeps a small expert-level one against routing collapse⁶⁹. DeepSeek-V2 runs them at 0.05 against 0.003⁷⁰. DeepSeek-V3's report no longer mentions a device-level loss: it balances each expert with a bias and a small sequence-wise loss, both covered later in this section⁴².

The same imbalance costs time in serving too, so a frontier MoE in production keeps it small: when DeepSeek serves DeepSeek-V4.1-Flash, its slowest GPU sits at most 1.09 times the mean⁷¹.

1185if self.config.moe_router_force_load_balancing:1186# Apply force load balancing with random logits for benchmark1187logits = apply_random_logits(logits)11881189if self.config.moe_router_force_biased is not None:1190# Apply biased logits with shared random bias across all ranks1191logits = apply_biased_logits(1192logits, self.config.moe_router_force_biased, self.layer_number1193)

1301@staticmethod1302def forward(ctx, logits: torch.Tensor) -> torch.Tensor:1303"""1304Forward pass returns random logits with rank-specific seed.13051306Args:1307logits (torch.Tensor): The logits.13081309Returns:1310torch.Tensor: The random logits.1311"""1312with get_cuda_rng_tracker().fork(get_expert_parallel_rng_tracker_name()):1313random_logits = logits.clone().normal_()1314return random_logits

1353@staticmethod1354def forward(ctx, logits, std, layer_number):1355"""Forward pass: apply random bias to logits."""1356# Check cache if reuse mode (negative std)1357if std < 0 and layer_number in RandomSTEShared._cache:1358return logits + RandomSTEShared._cache[layer_number]13591360# Generate random bias with shared seed across all ranks1361with get_cuda_rng_tracker().fork(get_data_parallel_rng_tracker_name()):1362bias = torch.empty(logits.shape[-1], device=logits.device, dtype=logits.dtype).normal_(1363std=abs(std)1364)13651366# Cache if reuse mode1367if std < 0 and layer_number is not None:1368RandomSTEShared._cache[layer_number] = bias13691370return logits + bias

Imbalance shows up as a slow step, so the number to watch is TFLOP/s against the forced-even control run. To see where the lost time goes, look at the step profiled in 3.7: it ran on an unbalanced router, and 3.7 shows how to tell whether imbalance is what slows its all-to-all.

Use it once per setup, as a control run: --moe-router-force-load-balancing gives the balanced throughput to compare every balancer against. Never train with it.

Forced routing shows how fast an even load runs, but it replaces the router's logits with random ones, so it can only serve as a control run, never as a way to train. How do we steer a learned router toward an even load while it trains, starting with the oldest fix, a loss term?

Aux loss

The oldest fix adds a loss term, the aux loss: the coefficient times the expert count times a sum over experts, each expert's share of the micro-batch's tokens times its mean router probability (moe_utils.py#L148-L152). Its gradient pulls router probabilities away from what the LM wants, toward an even load. Below, four settings route one batch: no balancing, the aux loss at a small and at a large coefficient, and the loss-free bias.

Both balancers spread the load in the figure, in different ways: the aux loss pulls the router's probabilities off what the LM loss wants, further at the larger coefficient, while the loss-free bias only changes which experts are picked and balances best, though it reroutes the most selections. The aux loss's pull makes its coefficient a tradeoff: too large a coefficient costs LM loss, and too small a coefficient leaves the load uneven. Yuan3.0's ablation of DeepSeek-V3's sequence-wise aux loss on a 10B model, plotted in the next figure, shows both sides: test loss rose with the coefficient from 1e-3 to 0.1, while the smallest coefficient, 1e-4, did slightly better than no aux loss at all but still left the busiest expert in a layer with about 13.6 times the tokens of the least used one^72,73.

The smallest coefficient in the figure leaves the load uneven, and the one that balances it, 1e-2, raises the test loss. Megatron sets the coefficient with --moe-aux-loss-coeff next to --moe-router-load-balancing-type aux_loss, one coefficient per loss type, applied in every layer. Motif 3 found its first MoE layers more susceptible to imbalance and gives them stronger auxiliary loss weights⁷⁴. We come back to the first layers later in the chapter.

148aggregated_probs_per_expert = probs.sum(dim=0)149aux_loss = torch.sum(aggregated_probs_per_expert * tokens_per_expert) * (150num_experts * moe_aux_loss_coeff / (topk * total_num_tokens * total_num_tokens)151)152return aux_loss

515aux_loss_coeff = self.get_aux_loss_coeff("aux_loss")516if aux_loss_coeff == 0:517return probs518519global_tokens_per_expert, local_num_tokens, total_num_tokens = (520get_tokens_per_expert_and_token_count(521routing_map=routing_map,522reduce_group=self.tp_cp_group,523topk=self.topk,524with_padding_mask=with_padding_mask,525)526)527528aux_loss = switch_load_balancing_loss_func(529probs=scores_for_aux_loss,530tokens_per_expert=global_tokens_per_expert,531total_num_tokens=total_num_tokens,532topk=self.topk,533num_experts=self.config.num_moe_experts,534moe_aux_loss_coeff=aux_loss_coeff,535fused=self.config.moe_router_aux_loss_fusion,536)

459def get_aux_loss_coeff(self, aux_loss_type: str) -> float:460"""Return the aux loss coeff for the given auxiliary loss type.461If the auxiliary loss type is not found, return 0.0.462"""463if isinstance(self.routing_type, str):464if self.routing_type == aux_loss_type:465return self.config.moe_aux_loss_coeff466if isinstance(self.routing_type, list):467try:468idx = self.routing_type.index(aux_loss_type)469return self.config.moe_aux_loss_coeff[idx]470except ValueError:471return 0.0472return 0.0

The measured row under the four-column figure above comes from Wang et al., who proposed the loss-free bias and compared it with a per-sequence aux loss on a 1B MoE with 64 routed experts and top-6. At their baseline coefficient of 1e-3, MaxVio over the validation set was 0.72 at a perplexity of 9.56, against 0.04 and 9.50 with the bias⁷⁵; with no aux loss MaxVio rose to about 5.1, and at 1e-2 it fell to about 0.1 while perplexity rose to about 9.69⁷⁶.

This aux loss, at a coefficient of 1e-3, is the learned router of Cost of imbalance, and by the end of its 120 iterations it had climbed to 125 TFLOP/s.

A coefficient set too high makes the model learn worse, and one set too low leaves the load uneven, so watch the LM loss and the balance together. Balance is best read from MaxVio rather than from the aux loss itself. The loss multiplies each expert's share of the tokens by its mean router probability, so when the router's probabilities are flat the loss reads the same as for an even load, whatever the token shares are, and a perfectly balanced reading can hide a skewed load. Sigma-MoE-Tiny hit this case: in its lower layers the loss went to its minimum by flattening the probabilities while the load stayed uneven⁷⁷. --moe-per-layer-logging logs this value per layer as load_balancing_loss, divided by the coefficient.

Use it when the simplest balancer will do: aux_loss at a coefficient around 1e-3, as in our first MoE run in 2.5. Raise it if the load stays uneven, lower it if the LM loss suffers⁷².

The aux loss we used so far counts each expert's load within one micro-batch, which leaves a question we have not asked yet: over which tokens should the load be even? You could balance it over each sequence, over each micro-batch, or over the whole global batch, and the next part compares the three.

Scope

Think about what one micro-batch holds. Pretraining pipelines often pack each micro-batch from a single domain, say all code, so a loss that asks for an even load inside that micro-batch spreads the code tokens over every expert, and no expert gets to specialize in code⁷⁸. Balance over the whole global batch instead, which mixes many domains, and the code tokens may crowd onto a few experts as long as tokens from other domains fill the rest. Qiu et al. saw experts specialize by domain only in that second setting, and it also lowered their perplexity⁷⁹. The GPUs pay for it: inside each micro-batch the load is no longer even, and their 43B model's step ran about 6% slower, about 1% of that slowdown from communication overhead and the rest mainly from the uneven load^80,81.

A single sequence is the tightest scope of all. Balancing each sequence on its own spreads every sequence over all the experts, even when a micro-batch mixes domains; DeepSeek-V2-Lite counted its only balance loss, an expert-level one, this way⁸², and DeepSeek-V3 keeps a small one to stop a single sequence from piling onto a few experts while the batch as a whole looks balanced⁴². In the figure below the three scopes run side by side, on GPUs that each hold two sequences from two domains: the tighter the scope, the less the experts can concentrate on one domain, and the looser it is, the more uneven the load each GPU carries.

In that simulation, the sequence-wise loss spreads the code sequence over all the experts, the micro-batch loss lets each GPU's two domains split its experts between them, and the global-batch loss leaves the code tokens mostly on their two favoured experts, at the price of the most uneven load per GPU. Megatron picks the scope through the load-balancing type: aux_loss balances per micro-batch, seq_aux_loss per sequence and global_aux_loss over the global batch. To combine scopes, pass several types and one coefficient for each, in the same order, for example --moe-router-load-balancing-type aux_loss seq_aux_loss --moe-aux-loss-coeff 1e-3 1e-4.

567scores_for_aux_loss = scores_for_aux_loss.reshape(seq_length, -1)568routing_map = routing_map.reshape(seq_length, -1)569570global_tokens_per_expert, local_num_tokens, total_num_tokens = (571get_tokens_per_expert_and_token_count(572routing_map=routing_map,573reduce_group=self.tp_cp_group,574with_padding_mask=with_padding_mask,575topk=self.topk * bsz,576)577)578579aux_loss = (580switch_load_balancing_loss_func(581probs=scores_for_aux_loss,582tokens_per_expert=global_tokens_per_expert,583total_num_tokens=total_num_tokens,584topk=self.topk,585num_experts=self.config.num_moe_experts,586moe_aux_loss_coeff=seq_aux_loss_coeff,587fused=self.config.moe_router_aux_loss_fusion,588)589/ bsz

617global_tokens_per_expert, local_num_tokens, total_num_tokens = (618get_tokens_per_expert_and_token_count(619routing_map=routing_map,620reduce_group=self.tp_dp_cp_group,621with_padding_mask=with_padding_mask,622topk=self.topk,623)624)625626self.global_tokens_per_expert += global_tokens_per_expert627self.ga_steps += 1628averated_tokens_per_expert = self.global_tokens_per_expert / self.ga_steps

233# Initialize global tokens per expert for global aux loss234if self.get_aux_loss_coeff("global_aux_loss") > 0:235self.register_buffer(236'global_tokens_per_expert',237torch.zeros(238self.config.num_moe_experts,239dtype=torch.float32,240device=torch.cuda.current_device(),241),242persistent=False,243)244self.register_buffer(245'ga_steps',246torch.tensor(0, dtype=torch.float32, device=torch.cuda.current_device()),247persistent=False,248)

1121def reset_global_aux_loss_tracker(self):1122"""Reset the global aux loss tracker."""1123if self.global_tokens_per_expert is not None:1124self.global_tokens_per_expert.zero_()1125self.ga_steps.zero_()

Each scope fails in its own way. Balancing per sequence or per micro-batch can keep domain experts from forming, so the model learns worse, while balancing over the global batch leaves the load on each GPU uneven, so the step gets slower. Whichever scope you pick, read balance from MaxVio at that scope, computed per micro-batch from the loads --save-tokens-per-expert-interval saves, per global batch from the sum of those loads, and per sequence offline from the routing trace (A.1).

Use micro-batch (aux_loss) by default. Use global_aux_loss only for domain-specialized experts, at +5.8% step time in Qiu et al.⁸¹.

Whichever scope you pick, a balancing loss still steers the router through its gradient, and that gradient competes with the one from the LM loss. Could you balance the load without touching the gradient at all?

Loss-free bias

The aux loss fixes an overloaded expert by pulling the router's probabilities away from what the LM wants. The loss-free bias leaves the probabilities alone and moves the load with no gradient. Each expert gets a bias, added to its score before the top-k and used only to pick experts, while the weights still come from the scores, and once per step Megatron moves each bias by a fixed amount, the update rate: up for underloaded experts, down for overloaded ones. The next figure runs that loop. In DeepSeek-V3's ablation it matched a batch-wise aux loss in validation loss, at 1B and at 3B⁸³. Like the aux loss, it reroutes some selections away from the LM's own top-k, and in the four-column toy of the Aux loss part it reroutes the most.

The figure runs 16 toy experts with a step far larger than Megatron's default, so a few rounds stand in for many: overloaded experts lower their bias and underloaded ones raise it until the load is close to even, and the back-and-forth that the large step leaves is damped once the toy cuts its step tenfold. In real training the catch is the fixed step, which must suit the expert count. Li et al. swept the expert count of small MoEs, and with an update rate of 1e-3 the busiest expert's load grew with it, to about 50 times the mean at 1,024 experts, where an aux loss of 1e-2 alone stayed under about 7 times⁸⁴. Above 128 experts the bias also hurt LM quality, and the authors suspect a step too coarse for that many experts⁸⁵.

The reference lines on the overload trace in the figure above come from the same 1B runs, measured on each training batch rather than over the validation set: MaxVio settled near 0.72 with the aux loss at 1e-3 and near 0.1 with the bias⁸⁶.

Megatron's rule is the simplest one you could write: every step, each bias moves up or down by the same fixed amount, however far its expert sits from the mean load. That is the sign update, and it is not the only way to run the loop. A linear update moves each bias in proportion to its expert's distance from the mean load, so a badly overloaded expert gets a large correction and a nearly balanced one a small one, while an AdamW update adapts the step for each expert. LongCat-Flash runs a linear update, aimed at the expert loads its compute budget calls for⁶⁵, and ZAYA1-8B feeds each expert's load error to AdamW⁸⁷.

The rules also differ in a subtler way: when an update's steps do not sum to zero, the mean of all the biases drifts. That drift changes no selection, since adding the same number to every expert's score leaves the top-k as it was, but it costs precision where the bias is stored in bf16; Megatron keeps it in fp32. Moonlight and Trinity Mini avoid the drift by subtracting the mean of the steps^88,89. The released weights differ the same way: DeepSeek-V3's biases sit at a per-layer mean of 2.2 to 8.0, while Moonlight's and ZAYA1-8B's average 0 and LongCat-Flash's sit just below it⁹⁰. Below, the sign update, a centered sign update, the linear update and the AdamW update run on the same batch.

The figure's toy batch has a few hot experts and many just under the mean load, so most sign steps point up: the sign and AdamW updates lift the mean of the biases while the centered sign and linear updates keep it at zero, yet the sign update balances as well as the centered one at every step, since top-k ignores a shift shared by all experts. Whatever rule moves it, the bias can only balance the tokens it counts. A batch-wide bias is fitted on one batch and applied to the next, so it can make the batch as a whole look even without guaranteeing that each sequence is: a single sequence can still pile onto a few experts. The obvious fix, a second bias computed over each whole sequence, would use tokens that come later in that sequence, and so leak the future into the routing of the earlier ones.

Chang's causal bias avoids the leak by building that second bias only from what the sequence has already routed. It resets at the start of every sequence, and after each token it lowers the bias of the experts just picked and raises the others', so each expert's bias tracks how far its picks so far sit above or below its share, k/N per token, and an expert this sequence has overused is pushed down for the tokens that follow. Like the batch bias, it only changes which experts are picked and is never trained by the gradient⁹¹. The next figure routes the same sequence with the global-batch bias alone and with the causal bias from this sequence's past added on top.

The sequence shown switches to code halfway. With only the global-batch bias, its code tokens pile onto one expert, even though the whole batch stays balanced; with the causal bias added, those code tokens spread over more experts and the sequence's load stays much closer to even. On the other hand, the simulation shows the causal bias sending more of the picks to experts outside the router's own top-k, away from what the router itself prefers, and the evidence for it is small-scale, from Chang's experiments on nanochat and TinyStories⁹¹.

Released models handle the single sequence differently. DeepSeek-V3 pairs its bias with the sequence-wise aux loss we saw earlier, at a tiny coefficient of 1e-4⁶²: the bias keeps the batch even, and the small loss guards against extreme imbalance within a single sequence. The figure that follows compares no such loss, a small one and a large one, on an ordinary sequence and on an extreme one.

In that toy, whose two coefficients are illustrative rather than DeepSeek-V3's, a small coefficient brings an ordinary sequence closer to even while moving the router only a little, but leaves an extreme sequence piled onto one expert; spreading that one takes a large coefficient, which pulls the router much further off the routing the LM loss prefers. A sequence is one way a batch-wide count can hide an imbalance, and a modality is another.

When a batch mixes modalities, a bias counted over all the tokens can balance their sum while the tokens of one modality, taken alone, still crowd onto a few experts. DeepSeek-VL2 and Kimi-VL keep one bias for all tokens^92,93, and DeepSeek-V4.1-Flash keeps one per modality⁶³. The figure below compares one shared bias with one bias per modality, each counting only its own modality's tokens.

The toy scripts text to lean toward four of its eight experts and images toward the other four, so a shared bias balances the total while each modality's surplus cancels the other's, and a bias per modality balances each modality on its own. The price of the split is that the bias then pushes every expert to take its share of each modality, away from specializing in images.

To run it in Megatron, turn the bias on with --moe-router-enable-expert-bias and set its fixed step with --moe-router-bias-update-rate, with sigmoid scores (--moe-router-score-function sigmoid), since Megatron accepts the bias only with sigmoid or sqrt(softplus) scores (transformer_config.py#L3139-L3147).

936elif expert_bias is not None:937scores_for_routing = scores + expert_bias.float()938_, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)939scores = torch.gather(scores, dim=1, index=top_indices)

209if self.enable_expert_bias:210self.register_buffer(211'local_tokens_per_expert',212torch.zeros(213self.config.num_moe_experts,214# Keep token counts exact through accumulation and cross-rank reduction.215# Float16Module leaves integer buffers in their original dtype.216dtype=torch.int64,217device=torch.cuda.current_device(),218),219persistent=False,220)221self.register_buffer(222'expert_bias',223torch.zeros(224self.config.num_moe_experts,225dtype=torch.float32,226device=torch.cuda.current_device(),227),228)

1253# All Reduce Across TPxCPxDP group1254torch.distributed.all_reduce(tokens_per_expert, group=tp_dp_cp_group)1255num_experts = tokens_per_expert.shape[-1]1256total_tokens = tokens_per_expert.sum(dim=-1, keepdim=True)1257# Compare each tokens_per_expert value with the row average without converting the integer1258# counts to floating point: tokens_per_expert < total / num_experts iff1259# tokens_per_expert * num_experts < total.1260update_direction = torch.sign(total_tokens - tokens_per_expert * num_experts)1261updated_expert_bias = (1262expert_bias + update_direction.to(dtype=expert_bias.dtype) * expert_bias_update_rate1263)

726if config.moe_router_enable_expert_bias:727if pg_collection is None:728tp_dp_cp_group = parallel_state.get_tensor_and_data_parallel_group(729with_context_parallel=True730)731_update_router_expert_bias(model, config, tp_dp_cp_group=tp_dp_cp_group)

When the update rate is too coarse for your expert count, the load stays uneven, which slows the step as Cost of imbalance showed, and above 128 experts Li et al.'s sweep also saw LM quality suffer. You can catch it in two places. Every checkpoint saves expert_bias, so you can follow how far the biases of each layer spread from one checkpoint to the next, and the token counts that --save-tokens-per-expert-interval writes give you MaxVio per global batch (A.1).

Use it when the balancer must barely pull against the LM loss: the bias adds no gradient at all, and DeepSeek-V3's sequence-wise loss at α = 1e-4⁶² adds very little. Run it with sigmoid scores, and tune the update rate to your expert count, since 1e-3 can be too coarse above 128 experts⁸⁵.

The bias moves by the same small step whatever the load looks like, so it can take many steps to catch up with a router that keeps changing. Instead of stepping toward the right bias, could you compute it directly from the scores of the current batch?

Quantile

Quantile balancing keeps the bias but solves for it. For each expert, it ranks the micro-batch's tokens by their score for that expert and reads off a quantile: the score that only the expert's target share of tokens beats. The bias is set from that quantile instead of moved by a fixed step. Jin et al. compared balancers on a 19B-parameter MoE with 768 experts, and quantile balancing balanced sooner than the sign update at about the same LM loss, halving the average MaxVio, though over the last 1k steps its MaxVio ended slightly higher⁹⁴. That precision has a price: in every micro-batch, each expert has to rank all of the micro-batch's tokens by its score to find its quantile, a top-k per expert on top of the usual top-k per token. Side by side below, the sign update and quantile balancing run on the same batch.

The toy routes the same batch at every step, so each solved bias routes the very tokens it was solved from: after the last step the sign update still leaves the load uneven, while quantile balancing brings it close to even within a few steps. In Megatron, each micro-batch solves for its own quantiles, and the solutions are averaged over the global batch to route the next one, optionally smoothed by an exponential moving average (EMA, --moe-router-quantile-balancing-ema).An EMA keeps most of the previous bias and mixes in only a share of each new solution, so the bias changes slowly from step to step.

233topk_result = (scores - beta).topk(k + 1, dim=1)234indices = topk_result.indices[:, :-1]235236if not update_beta:237return indices, beta238239assert (num_tokens * k) % num_experts == 0, (240"Quantile balancing requires the number of routed assignments "241f"({num_tokens} tokens * top-{k}) to be divisible by "242f"{num_experts} experts."243)244col_target = num_tokens * k // num_experts245alpha = topk_result.values[:, -1:]246beta_local = (scores - alpha).topk(col_target + 1, dim=0).values[-1].contiguous()247return indices, beta_local

441stacked_beta = torch.stack(qb_beta_list, dim=0)442local_avg_list = [443accum / count.clamp(min=1).to(accum.dtype)444for accum, count in zip(qb_beta_accum_list, qb_beta_count_list)445]446stacked_local_avg = torch.stack(local_avg_list, dim=0)447448torch.distributed.all_reduce(449stacked_local_avg, op=torch.distributed.ReduceOp.AVG, group=dp_cp_group450)451452ema = config.moe_router_quantile_balancing_ema453stacked_new_beta = ema * stacked_beta + (1.0 - ema) * stacked_local_avg454stacked_new_beta = stacked_new_beta - stacked_new_beta.mean(dim=-1, keepdim=True)455456for qb_beta, new_beta in zip(qb_beta_list, stacked_new_beta):457qb_beta.copy_(new_beta)

256if self.routing_type == "quantile_balancing":257assert not self.is_aux_loss_enabled(), (258"Quantile balancing handles load balance via the bias update; "259"aux losses must be disabled (set moe_aux_loss_coeff to 0)."

398assert (399not self.config.moe_router_fusion400), "Quantile balancing routing does not support moe_router_fusion."401assert (402self.config.moe_router_num_groups is None and self.config.moe_router_group_topk is None403), "Quantile balancing routing does not support group-limited routing."

Kimi K3 runs quantile balancing over its whole global batch, reading each expert's quantile from a histogram summed across GPUs⁶⁴. It left the fixed-step rule because balancing grew harder once its routed pool reached 896 experts per layer⁶⁴, and Su Jianlin, who proposed quantile balancing⁹⁵, writes that Kimi K2 had used the sign update⁹⁶.

Averaging has a price: a bias shared by the whole global batch can balance that batch while each GPU's micro-batch stays uneven. Neitemeier et al. measured how wide that gap gets in a 7.5B run (one seed). The rank-averaged quantile that Megatron uses balanced the global batch to a MaxVio of 0.92, while in the worst layer the busiest expert of an expert-parallel micro-batch still took about 7 times its fair share, a local MaxVio of 5.96. Solving for the exact global-batch quantile lowered both⁹⁷. They read local MaxVio as a proxy for dispatch cost rather than as measured throughput⁹⁸.

How do the balancers compare in throughput? We ran five router configurations from init on the same model, with one seed, and compared their median throughput over the last 20 of 120 iterations. Quantile balancing reached 139 TFLOP/s, against 125 for the aux loss and 99 for the tested bias-only configuration. The bias-only run used sigmoid scores and an update rate of 1e-3, while the other four used softmax.These runs used an older Megatron commit, which accumulated the bias's token counts in bf16, so a count could round and flip the direction of its bias update; at d6316eb45 the counts are exact. With an EMA of 0.9, quantile balancing reached 121 TFLOP/s. That setting gives each new solution 10% of the weight in the bias update (transformer_config.py#L910-L913).

Jin et al. added an EMA to quantile balancing: stronger smoothing kept the bias steadier between checkpoints but corrected the load more slowly early in training⁹⁹. All five runs are plotted below.

Of the five runs in the figure, the bias-only one stays flat while the other four climb, with quantile balancing fastest over the final 20 iterations. Throughput is still rising at the end, so the ranking describes these first 120 iterations. Quantile balancing also ends with the highest training LM loss of the five, so its throughput lead alone does not establish faster training to the same quality.

To try quantile balancing, set --moe-router-load-balancing-type quantile_balancing.The same flag also accepts sinkhorn: Sinkhorn uses a balanced-assignment procedure to choose experts for each token during training.

The risk is a global batch that looks balanced while each GPU's micro-batch is not, which Neitemeier et al. read as a higher dispatch cost. To see the gap, compute MaxVio twice from the token counts that --save-tokens-per-expert-interval writes, once per micro-batch and once per global batch, and compare the two (A.1).

Use it when the sign-update bias balances too slowly: solving for the bias reached a lower average MaxVio than stepping toward it, at about the same LM loss, with 768 experts⁹⁴, and it had the highest median throughput over the last 20 of 120 iterations in our five-configuration comparison (one seed). Plan around its limits, since Megatron currently refuses it next to any aux loss, --moe-router-fusion or group-limited routing.

What to watch

Every balancer in this section spreads tokens over a set of experts that was fixed before training started, and we have not yet asked what that set should look like. How many experts should a layer have, how wide should each one be, and should some of them see every token?

3.4 Experts: how many, how big, shared or not

Once the router has picked a token's experts, what those experts look like decides both what the layer can learn and what it costs to run. You choose how many experts a layer has and how wide each one is, how to keep the network traffic of many small experts in check, whether one expert should see every token, and which activation runs inside each of them.

Expert choices in leading open MoEs, plus Nemotron for LatentMoE and ReLU² and LongCat-Flash for variable compute. Count and size gives the routed experts each token uses out of the total, then each expert's FFN width; shared experts are listed separately. Latent dimensions show token width before and after projection. "No limit specified" refers to the published configuration.

With the parameter budget fixed, adding more experts means making each one narrower. Which split gives you a better model for the same compute?

Count and size

Picture one MoE layer with a fixed budget of parameters and compute, and cut its experts finer: twice as many experts at half the width, with twice the top-k, keeps both budgets the same. Figure 3.4.1 draws that ladder: coarse, then 16 experts of d/2 with top-4, then 32 of d/4 with top-8. Liew et al. ran the ablation in an 18-layer model with total and active parameters fixed, and loss fell with each finer cut up to 128 experts of d/4 with top-8, from 2.334 to 2.297, while finer cuts did not improve it further¹⁰⁸. In this experiment, then, segmenting stopped paying past 128 experts, while every finer step still ran smaller GEMMs and dispatched more copies of each token.

Released models sit on different rungs of that ladder. gpt-oss-120b and Trinity Large keep each expert as wide as the model, with top-4 of 128 and top-4 of 256 experts, Mistral Small 4 halves that width with top-4 of 128, and Qwen3.5-35B-A3B quarters it with top-8 of 256¹⁰⁹.

In the figure, you can see that splitting the experts changes their number and width while keeping the total parameters and compute per token fixed. The total tile area represents the parameters, and the dark area represents the expert computation for one token. What you are changing is the granularity.Granularity describes how finely you divide the expert capacity: finer granularity means more, narrower experts. To keep the same compute per token, each token uses more of the narrower experts. You can also increase sparsity: add more parameters while keeping the number used by each token fixed. Two ways to spend those extra parameters are compared next: more routed experts, or n-gram embedding tables looked up by token id.

The figure plots LongCat's comparison at 280M active parameters. There, n-gram tables lowered the loss more than the same parameters spent on more experts (by about 0.01 at 17 times the active parameters) once the experts already held about 12 times the active parameters, while tables added when the experts stopped at about 4 times lost from the start. The gain had a limit: on English validation loss, more experts won again once tables held a little over 40% of the parameters, and LongCat's rule, read from that crossing, is to put no more than 50% of the parameters in tables¹¹⁰. In that comparison, expanding the experts paid off first, and once they were large enough, part of the next parameters did more in lookup tables. Where you spend extra parameters therefore depends on how much expert capacity you already have.

Both ways of spending ship in released models: Kimi K2 grows its pool to 384 experts with top-8, and LongCat-Flash-Lite puts 31.4B of its 68.5B parameters in n-gram embedding tables¹¹¹.

Whether you add more experts or divide the existing capacity into narrower ones, the resulting shape also affects training speed. Narrower experts run smaller GEMMs, which use the GPU less efficiently.

For a dense model, compute-optimal scaling laws help you choose how to divide a fixed FLOP budget between model size and training tokens: a larger model leaves room for fewer tokens, while a smaller one can train on more¹¹². An MoE adds another choice. You can increase its total parameters without increasing the parameters used by each token, so you also need to choose its sparsity. For either kind of model, choosing under a time limit also means accounting for how efficiently each configuration uses your hardware¹¹³.

Sarkar et al. combined a fitted loss scaling law with a training-throughput model to compare the sparsity that minimizes loss at fixed model FLOPs with the sparsity that minimizes loss at fixed cluster time. Their predictions for one cluster and training window are plotted below¹¹³.Sparsity here is the fraction of the model's parameters that a token does not use.

In their predictions, choosing by model FLOPs alone points you toward the sparsest model in the grid. Once the authors account for training speed on their cluster, the lowest predicted loss moves to a less sparse configuration, 0.031 nats below the sparsest one. As in dense scaling, you are balancing model capacity against how many tokens you can train on. With an MoE, increasing sparsity can lower hardware utilization without increasing model FLOPs per token, so extra capacity can still cost training time¹¹³.

Megatron sets the shape with three flags: --num-experts for the count, --moe-ffn-hidden-size for each expert's width and --moe-router-topk for how many experts each token picks.

262# Double the output width with gated linear unit, see https://arxiv.org/pdf/2002.05202.pdf263ffn_hidden_size = not_none(self.config.moe_ffn_hidden_size)264if self.config.gated_linear_unit:265ffn_hidden_size *= 2266267self.linear_fc1 = submodules.linear_fc1(268self.num_local_experts,269self.input_size if self.config.moe_latent_size is None else self.config.moe_latent_size,270ffn_hidden_size,271config=self.config,272init_method=not_none(self.config.init_method),273bias=self.config.add_bias_linear,274skip_bias_add=False,275is_expert=True,276tp_comm_buffer_name='fc1',277pg_collection=pg_collection,278name=(name + ".linear_fc1") if name is not None else None,279)280281if self.config.use_te_activation_func and not (submodules.activation_func is None):282self.activation_func = apply_module(submodules.activation_func(config=self.config))283else:284self.activation_func = self.config.activation_func285286self.linear_fc2 = submodules.linear_fc2(287self.num_local_experts,288not_none(self.config.moe_ffn_hidden_size),289(290self.config.hidden_size291if self.config.moe_latent_size is None292else self.config.moe_latent_size293),

77# Initialize the gate weights.78# TODO: Add support for GPU initialization, which requires updating the golden values.79self.weight = torch.nn.Parameter(80torch.empty((self.config.num_moe_experts, self.config.hidden_size), dtype=torch.float32)81)

Cut the experts too fine and each expert's GEMMs get small enough to lower TFLOP/s per GPU, and with the wrong granularity the model learns worse. To see what a shape costs you, read the grouped-GEMM time per MoE layer from the profiler next to TFLOP/s per GPU from --log-throughput. The token counts that --save-tokens-per-expert-interval writes also show dead experts that receive no tokens (A.1), but an expert can also die while it keeps its tokens, which shows in its output norm instead (3.6)¹¹⁴.

Use finer experts at equal total and active parameters, up to a point: in one 18-layer model, the loss fell until 128 experts of d/4 with top-8, and finer splits gained nothing¹⁰⁸. Stop going finer once small GEMMs and token copies start to dominate your step.

Finer experts help the loss only up to a point, and each finer step sends more copies of every token over the network. Could you keep the many small experts and send less data for each copy?

Latent MoE

Each copy of a token carries its full hidden state, d numbers, sent once for every expert it picked. LatentMoE shrinks that copy. It projects each token down to a smaller latent ℓ before dispatch, runs the experts on ℓ, and projects back up after combine. Every copy shrinks from d to ℓ, and the savings pay for more experts: at the same expert parameters, expert compute and traffic, the layer can afford α = d/ℓ times more experts and an α times larger top-k¹¹⁵. Below, a standard layer sits next to its latent counterpart, keeping the parameters and compute inside the experts fixed.

Drawn as tiles, the standard layer selects four of sixteen experts, each receiving a full-width token vector. The latent layer selects eight of thirty-two, but each receives only half as many values. Both send the same amount of data per token. Their total tile area and dark area are also equal, showing unchanged expert parameters and compute; the two projections add work outside those areas¹¹⁵.

Does the trade pay off in quality? In NVIDIA's 73B-A8B ablation, at equal total and active parameters, LatentMoE scored 4.6 points higher on MMLU-Pro than a standard MoE¹¹⁶, and a small-scale replication also lowered perplexity¹¹⁷. Nemotron 3 Super, Nemotron 3 Ultra and Kimi K3 ship it¹¹⁸. In return, every token runs through two d x ℓ projections, and serving throughput ranged from 12% lower at concurrency 1 to slightly higher at some concurrencies¹¹⁹. In Megatron, the latent size is --moe-latent-size.

304self.fc1_latent_proj = linear_cls(305self.config.hidden_size,306self.config.moe_latent_size,307parallel_mode="duplicated",308config=self.config,309init_method=self.config.init_method,310bias=self.config.add_bias_linear,311skip_bias_add=False,312skip_weight_param_allocation=False,313is_expert=False,314name=(name + ".fc1_latent_proj") if name is not None else None,315**linear_gtp_kwargs,316)317if self.config.moe_use_norm_before_up_proj:318self.fc2_norm = TENorm(319config=self.config,320hidden_size=self.config.moe_latent_size,321eps=self.config.layernorm_epsilon,322)323self.fc2_latent_proj = linear_cls(324self.config.moe_latent_size,325self.config.hidden_size,326parallel_mode="duplicated",327config=self.config,328init_method=self.config.output_layer_init_method,329bias=self.config.add_bias_linear,330skip_bias_add=False,331skip_weight_param_allocation=False,332is_expert=False,333name=(name + ".fc2_latent_proj") if name is not None else None,334**linear_gtp_kwargs,335)

554# Project the hidden_states from hidden dimension down to latent dimension.555if self.config.moe_latent_size:556hidden_states, _ = self.fc1_latent_proj(hidden_states)

662output = self.token_dispatcher.combine_postprocess(output)663if self.config.moe_latent_size:664if self.config.moe_use_norm_before_up_proj:665output = apply_module(self.fc2_norm)(output)666output, _ = self.fc2_latent_proj(output)

269self.input_size if self.config.moe_latent_size is None else self.config.moe_latent_size,

A latent MoE can go wrong in two ways: it runs slow, or its activations blow up. When it runs slow, look past the traffic, since the all-to-all now carries ℓ numbers per copy instead of d, to the two extra projections and the smaller expert GEMMs. LatentMoE's NVIDIA authors propose small-GEMM kernels and a separate stream for the shared expert as fixes¹²⁰. When it blows up, as Kimi K3 reports for the activations of its routed branch at its 2.8-trillion-parameter scale¹⁰², one place to look is the output norm of the up-projection, fc2_latent_proj, which you can compute per layer offline from the --save-activations-interval dumps (A.1). K3's fix was an RMSNorm before the up-projection, --moe-use-norm-before-up-proj, which it ships with a routed scaling factor of 1.0¹²¹.

Use a latent when you want more experts and a larger top-k without sending more data per token: shrinking each copy to ℓ = d/α affords α times more of both¹¹⁵, and it scored higher at matched size¹¹⁶, at the cost of two projections.

Every expert so far sits behind the router and sees only the tokens sent to it. What would you gain by putting one expert outside the router, where it sees every token, and is it worth its compute?

Shared experts

A shared expert sits outside the router: it processes every token, next to whichever routed experts the token picked. DeepSeekMoE introduced shared experts to learn common patterns, reducing duplicated knowledge across routed experts and giving them room to specialize¹²². In its ablation at equal compute, adding a shared expert improved results on most benchmarks¹²³. Qwen's models show how open the question is: the Qwen3 MoE models in the Qwen3 report dropped it¹²⁴, Qwen3-Next brought it back behind a sigmoid gate¹²⁵, and Qwen3.5 kept it¹²⁶.

Measured by loss in models trained from scratch at equal compute, the gain is small and not consistent. A shared expert lowered loss by about 0.01 to 0.02 in Pangu Ultra MoE and in Sun et al.'s expert-threshold models^127,128, while in OLMoE and in Li et al.'s sweep of small MoEs it did nothing or slightly worse^129,130.

Sharing also reduces the router's freedom. With 32 experts and 4 active, reserving one slot for a shared expert leaves 4,495 ways to choose the other 3, compared with 35,960 ways to choose 4 from all 32. If you instead add the shared expert without reducing routed top-k, its always-on FFN increases compute per token: Mellum2 added a shared expert without reducing routed top-k and found no quality gain and slower inference¹³¹.

If you do add one, you still have to decide how its output mixes with the routed experts'. The figure below lines up five designs at equal compute per token: routed only, a shared expert at fixed weight (DeepSeek-V3)³⁰, a gated one (Qwen3-Next)¹²⁵, one normalized jointly with the routed experts (Inkling)¹⁹ and a parallel dense FFN (Grok 2)¹³².

Lined up this way, a shared expert at fixed weight holds the same share of every token's weight, Qwen's gate moves that share between 0 and 50%, Inkling's router learns it per token by scoring the shared experts with the routed ones, and a parallel dense FFN without renormalization holds over half of it whatever the router scores.

Megatron's flags are --moe-shared-expert-intermediate-size, --moe-shared-expert-gate (a sigmoid gate on its output) and --use-grouped-gemm-for-shared-expert. Without the gate, Megatron adds the shared output at weight 1 (moe_layer.py#L668-L669), while the routed weights sum to the routed scaling factor when normalized over the selected experts (moe_utils.py#L942-L947, 3.1).Megatron does this with sigmoid routing, or by applying softmax after selecting the top-k experts. For the latter, leave --moe-router-pre-softmax disabled. The shared expert therefore takes 1/(1 + factor) of each token's output weight, 29% at a factor of 2.5, and the larger the factor, the smaller its share; GLM-4.5-Air and Mistral Large 3 set the factor to 1, which gives it half^133,134.

128config.ffn_hidden_size = config.moe_shared_expert_intermediate_size

201def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:202"""Forward function"""203output, _ = super().forward(hidden_states)204if self.use_shared_expert_gate:205logits = torch.nn.functional.linear(hidden_states, self.gate_weight)206gate_score = torch.nn.functional.sigmoid(logits)207output = output * gate_score208return output

668if shared_expert_output is not None:669output = output + shared_expert_output

To track how much the shared expert contributes to each token's output during training, compare the shared branch after its gate with the routed branch after combine and any latent up-projection, using the sparse forward hook in A.1.

A shared expert is a common starting choice in leading MoEs, including GLM-5.3, Kimi K3 and DeepSeek V4.1 Flash^135,136,137. MiMo-V2.6 ships without one¹³⁸. To compare with and without one at equal compute, give it the routed experts' width and replace one routed slot with it.

Shared or not, every expert so far has the same width and every token gets the same top-k, so every token gets the same compute. Could the router give some tokens less compute and others more, while the average stays fixed?

Variable compute

Two designs let the compute per token vary while the mean stays fixed. LongCat-Flash adds zero-computation experts that pass the input through unchanged. Its router selects 12 experts per token from a pool of 512 FFN experts and 256 of these identity paths. A bias adjusts how often the FFN experts are selected, targeting an average of 8 per token while allowing individual tokens to use more or fewer¹³⁹; the standard deviation of that count rose to about 2.5 during pre-training¹⁴⁰. Xiaomi's AI Lab takes a different approach in MoDSE (Mixture of Diverse Size Experts): they mix expert widths, letting the router choose larger or smaller experts for each token¹⁴¹. The figure below simulates the active width per token for homogeneous experts, as in Mixtral¹⁴², and both designs.

In the simulation, where one bias per design holds the mean active width at that of the homogeneous experts, some tokens get more of it and others less, and every token still fills all its top-k slots, since only what lands in them varies. In LongCat-Flash's ablation, zero-computation experts reached lower loss at matched compute¹⁴³. Mixed expert sizes can improve quality, but may cost throughput. Xiaomi's AI Lab and Wang et al. each put one router over mixed-size experts in small models and report gains^141,144. Li et al. found none, but in their setup the experts sit in two pools, each with a fixed top-k, so every token still gets the same compute¹⁴⁵. In GroveMoE, the cost showed at inference: its mixed sizes ran about 30% slower than Qwen3-30B-A3B, which its authors attribute to processing the two expert sizes with separate MoE kernels in each layer¹⁴⁶.

Track the mean and spread of compute per token in each layer: the number of selected FFN experts for zero-computation routing, or the sum of selected expert widths for mixed-size routing. Compare the mean with the target budget, and log token counts per expert alongside compute per expert-parallel rank and step time.

Start with routed experts of one width and a fixed top-k, as MiMo-V2.6-Pro-RL, GLM-5.3 and Kimi K3 do^138,135,136. This also fits Megatron's grouped linear layers, which use a single moe_ffn_hidden_size.

With the count, width and compute of the experts settled, one choice remains inside each expert: the function between its two projections. How much does that choice matter for an MoE?

Activation

As in a dense Transformer, each expert uses an activation function such as SwiGLU, as in DeepSeek-V3, Qwen3 and Kimi K2¹⁴⁷, or ReLU² between its input and output projections. In an MoE, you also need to watch how activations behave across experts: Step 3.5 Flash found that one or two experts could develop rapidly growing output norms while the others stayed stable and the training loss barely changed¹⁵. The activation function, and any clamp applied to it, helps control that growth. The figure below shows how the output of SwiGLU, ReLU² and Expert-Specific PolyNorm grows as the input scale rises.

Scaled up, SwiGLU and ReLU² outputs grow with the square of the input scale, while Expert-Specific PolyNorm normalizes its gate, so its output grows only linearly. A clamp can bound that growth. Two studies found similar quality with and without clipping in their tested settings. In a 7.9B MoE ablation by Ant Group's Ling team, SwiGLU with clipping matched plain SwiGLU's average over 17 benchmarks¹⁴⁸. Step 3.5 Flash's activation clipping held its outlier experts' norms down at indistinguishable loss, although that branch started mid-run from the unclipped run¹⁵. Clipping did not eliminate loss spikes in the Ling team's FP8 run, the clipped SwiGLU delayed the loss spikes, yet one still came around step 77,000¹⁴⁹. The next figure sweeps the input scale through plain SwiGLU, a hard clamp and SiTU-GLU's smooth cap at the same bound, so you can see how each one bends.

Through that sweep, plain SwiGLU's largest hidden unit keeps growing with the square of the gate input up to 6,862, the peak K-EXAONE 2.0 measured at its last layer with its clamp disabled at inference¹⁵⁰, the hard clamp follows it up to the bound and stops there, and SiTU-GLU bends away well before the bound. The sweep sets both to stop a hidden unit at about 100, the clamp at 10 and SiTU-GLU with a gate cap of 4 and an up cap of 25, and from about the same input both pass almost no gradient to those units' gate weights. Replacing the activation altogether rests on thinner evidence. Zhuo et al.'s PolyNorm lowered average validation loss by 0.02 against SwiGLU in OLMoE-1B-7B¹⁵¹, while Motif 3 compares its Expert-Specific PolyNorm with SwiGLU only on effective rank¹⁵². Kimi K3 states that SiTU-GLU's smooth cap gives better training behavior than hard clamping¹⁵³.

In Megatron, --activation-func-clamp-value sets the hard-clipping limit for SwiGLU in MoE layers; a value of 10 matches DeepSeek V4.1 Flash's limit¹⁵⁴, which GigaChat 3.5 and Hy4-preview also use¹⁵⁵. Kimi K3's smooth cap, SiTU-GLU, is --activation-func-tanh-clamp-scale 4 --activation-func-tanh-clamp-scale-linear 25, available from Megatron Core v0.20 on, and Nemotron 3's ReLU² is --squared-relu¹¹⁸.

1100if tanh_clamp_scale is not None:1101intermediate_parallel = situ_glu(1102intermediate_parallel,1103tanh_clamp_scale,1104self.config.activation_func_tanh_clamp_scale_linear,1105self.config.glu_linear_offset,1106)1107else:11081109def glu(x):1110x_glu, x_linear = torch.chunk(x, 2, dim=-1)1111if (val := self.config.activation_func_clamp_value) is not None:1112x_glu = x_glu.clamp(min=None, max=val)1113x_linear = x_linear.clamp(min=-val, max=val)1114return self.config.activation_func(x_glu) * (1115x_linear + self.config.glu_linear_offset1116)11171118intermediate_parallel = glu(intermediate_parallel)

17def tanh_soft_clamp(x: torch.Tensor, scale: float) -> torch.Tensor:18"""Tanh Soft Clamp to precondition activation inputs."""19return (scale * torch.tanh(x.float() / scale)).to(x.dtype)202122@jit_fuser23def situ(x: torch.Tensor, scale: float) -> torch.Tensor:24"""Sigmoid-Tanh Unit from Kimi K3: ``s * tanh(x / s) * sigmoid(x)``"""25return tanh_soft_clamp(x, scale) * torch.sigmoid(x)262728@jit_fuser29def situ_glu(30x: torch.Tensor,31gate_scale: float,32linear_scale: Optional[float] = None,33linear_offset: float = 0.0,34) -> torch.Tensor:35"""SiTU-GLU: ``situ(x_gate, gate_scale) * (x_linear + linear_offset)``."""36x_gate, x_linear = torch.chunk(x, 2, dim=-1)37if linear_scale is not None:38x_linear = tanh_soft_clamp(x_linear, linear_scale)39return situ(x_gate, gate_scale) * (x_linear + linear_offset)

2993if (2994self.activation_func_clamp_value is not None2995and self.activation_func == F.silu2996and self.gated_linear_unit2997):2998if (2999not math.isfinite(self.activation_func_clamp_value)3000or self.activation_func_clamp_value <= 03001):3002raise ValueError(3003"activation_func_clamp_value for SwiGLU must be finite and greater than zero."3004)3005if self.num_moe_experts is None:3006raise ValueError(3007"activation_func_clamp_value for SwiGLU is only supported with MoE."3008)3009if self.glu_linear_offset != 0.0:3010raise ValueError(3011"glu_linear_offset must be zero when activation_func_clamp_value "3012"is set for SwiGLU."3013)

Outlier activations in the experts can blow up a run. To catch them, compute Step 3.5 Flash's signal, the max-to-median ratio of expert output norms per layer¹⁵⁶, offline from the --save-activations-interval dumps (A.1): one expert pulling away from the others shows as a rising ratio. Each dump keeps only the last micro-batch and copies it to the CPU, so sample sparsely.

Start from your reference model's activation recipe. Leading MoEs use different ways to bound activations: Step 5 Preview clips SwiGLU at 7, DeepSeek V4.1 Flash at 10, and Kimi K3 uses SiTU-GLU's smooth caps^107,154,153. For a new architecture, SwiGLU is a reasonable baseline. If expert norms diverge, test clipping and compare training loss as well as the norms before choosing a limit.

What to watch

Everything in this section happened inside one MoE layer, while a model stacks many of them, which raises two questions we have not asked yet. Should every layer be an MoE layer, and how should the block around each one be wired?

3.5 Block layout

Picture the model as a stack of blocks, each one attention followed by either a dense MLP or an MoE layer. Every MoE layer sends its tokens out to their experts and back with an all-to-all, so where the MoE layers sit decides how many all-to-alls a step makes, and also which layers end up balancing worst. How each block is wired then decides how much of every all-to-all can hide behind attention.

Where models put their MoE layers and how they connect them to attention. Dense-layer counts refer to the backbone, excluding prediction modules.

The first layout decision is which layers get experts at all, since you can put an MoE in every layer, keep a few layers dense, or alternate the two. Which of these should you pick?

Dense layers

A dense layer has no router and no experts, so it makes no all-to-all. Figure 3.5.1 counts the all-to-alls of three layouts side by side: MoE in every layer, as in Qwen3, gpt-oss, MiniMax M2 and Llama 4 Scout¹⁶², leading dense layers, and MoE every other layer. Leading dense layers replace only the earliest MoE layers, as in DeepSeek-V3, which keeps its first three dense¹⁶³. Putting an MoE in every other layer halves the all-to-alls instead, as in Llama 4 Maverick, which alternates dense and MoE layers¹⁶⁰. Either way, you give up some MoE capacity in exchange for fewer all-to-alls and fewer parameters.

The figure's toy scripts its routing to follow the token id less with depth, so its first layer balances slowest: leading dense layers remove its two slowest layers but cut only 4 of 16 all-to-alls, while MoE every other layer halves the all-to-alls but keeps the second slowest layer. In Ling's 60-layer models, leading dense layers cost little: making the first 1 to 3 layers dense barely changed efficiency at matched FLOPs, an efficiency leverage close to 1¹⁶⁴.

The other argument for them is balance, and it rests on less evidence. DeepSeekMoE observed that the first layer's balance converges more slowly¹⁶⁵, and Ling says leading dense layers ease early-layer imbalance¹⁶⁶, though its ablation of those layers reports only loss¹⁶⁴. Later reports pull in different directions. Some keep the first layers away from a learned router: Nucleus-Image, an image-generation MoE, diverged without its three leading dense layers¹⁶⁷, and DeepSeek-V4 hashes its first 3 MoE layers⁵³.

Others keep routing there and treat the imbalance instead. Solar Open uses MoE in every layer and fixes its early-layer imbalance with a sequence-wise auxiliary loss¹⁶⁸. Sigma-MoE-Tiny routes each token to only 1 of 96 experts, and with an aux loss alone its lower layers collapsed: the loss took a shortcut by making the router's probabilities uniform, while some experts still got almost no tokens. Its fix was to activate more experts per token in the first 8 layers early in training, then drop back to one⁷⁷. Dense layers do not always end the problem either: Motif 3, whose first two layers are dense, saw its earliest MoE layers balance worst early in training⁷⁴, while ERNIE 5.0 found no severe imbalance in its first layer on text¹⁶⁹.

Megatron takes the layout as a list in --moe-layer-freq, one entry per layer, 0 for dense and 1 for MoE: "([0]+[1]*15)" gives one leading dense layer, and "([0,1]*8)" a dense-first interleave like Maverick's. The flag also takes an integer, and that form puts an MoE at layer 0, so --moe-layer-freq 2 starts with an MoE rather than with a dense layer as Maverick does. The dense layers take their width from --ffn-hidden-size, separately from the expert size.

635if isinstance(config.moe_layer_freq, int):636moe_layer_pattern = [6371 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.num_layers)638]639elif isinstance(config.moe_layer_freq, list):640moe_layer_pattern = config.moe_layer_freq641assert len(moe_layer_pattern) == config.num_layers, (642f"Invalid length of moe_layer_pattern: {len(moe_layer_pattern)}, "643f"expected {config.num_layers}, "644f"current moe layer pattern: {config.moe_layer_freq}"645)

652layer_specs = []653for layer_number in range(config.num_layers):654if moe_layer_pattern[layer_number] == 1:655layer_specs.append(moe_layer_spec)656elif moe_layer_pattern[layer_number] == 0:657layer_specs.append(dense_layer_spec)658else:659raise ValueError(f"Invalid layer pattern: {moe_layer_pattern}")

The reports above saw imbalance in the early layers, and with dense layers in those places, the first MoE layers after them can still balance worst, as they did in Motif 3. Those first MoE layers are the ones to watch. Follow their aux loss with --moe-per-layer-logging, whose per-layer values the launcher from 2.4 sends to Trackio, and their MaxVio from the routing trace (A.1). The dense layers have no router to watch. The z-loss measures how large the router logits grow, not how evenly tokens spread.

Use leading dense layers when you want fewer all-to-alls at little cost: one to three of them barely changed efficiency at matched FLOPs in Ling's 60-layer models¹⁶⁴. Do not count on them to fix early-layer imbalance: neither DeepSeekMoE nor Ling shows data for it^165,166, and Motif 3 still saw it in its first MoE layers after two dense ones⁷⁴.

Every MoE layer you keep still sends its tokens out and back, and in an MoE block, where the MoE runs on attention's output, the GPUs spend much of that time waiting. Could you rewire the block so that some of that traffic moves while attention computes?

Block wiring

In an MoE block, each all-to-all is tied to compute on both sides. The router reads attention's output, so dispatch cannot start until attention is done, and the next layer reads the MoE's output, so nothing after the MoE can start until combine is done. The Ultra-Scale Playbook covers how to overlap communication with compute in general.

Rewiring the block removes one of these dependencies or both, and the communication it frees can then overlap with compute. A parallel block feeds attention and the MoE the same input, as Cohere's North Mini Code and Moondream 3 do¹⁷⁰, so dispatch runs while attention computes. Combine is harder to hide, because it can only overlap with work left in the same layer, which in Megatron's shortcut block is attention's output projection and the shared expert; with neither, it stays exposed. A shortcut-connected MoE hides combine too, behind the dense compute that runs beside it, as in LongCat-Flash¹⁷¹ and in Pos-2, the preferred variant of ScMoE, Cai et al.'s shortcut-connected MoE¹⁷². Figure 3.5.2 compares the three wirings inside one micro-batch, and in its illustrative durations the exposed all-to-all share of the layer falls from 33% in the MoE block to 7% with the shortcut.

102compute_layer = layers[physical_index]103paired_type = layer_type_list[physical_index]104if paired_type not in SUPPORTED_SHORTCUT_PREDECESSORS:105raise ValueError(106"Shortcut MoE must be preceded by a Mamba, GDN, or supported attention layer"107)108109supports_two_stage = (110isinstance(compute_layer, TwoStageAttentionLayer)111and compute_layer.supports_two_stage_attention()112)113if not supports_two_stage:114raise ValueError(115f"Shortcut attention layer {paired_type!r} does not support two-stage attention"116)117moe_layer = layers[physical_index + 1]118grouped_layers.append(119ShortcutMoEBlock(120compute_layer,121moe_layer,122overlap_a2a=config.moe_shortcut_parallel,123attn_local_idx=physical_index,124moe_local_idx=physical_index + 1,125)126)127physical_index += 2

490# Launch the moe_router491with quant_context_factory(moe_config, self.moe_layer_idx):492moe_hidden_states = self._read_shortcut_hidden(493moe_hidden_states, recompute_context=attn_recompute_context494)495route_input, route_probs = self._moe_router_preprocess(496shortcut_hidden=moe_hidden_states,497padding_mask=padding_mask,498packed_seq_params=moe_packed_seq_params,499recompute_context=attn_recompute_context,500)501if overlap:502self.route_ready_event.record(torch.cuda.current_stream())

362dispatch_stream = self._get_a2a_overlap_stream()363dispatch_stream.wait_event(self.route_ready_event)364hidden_states.record_stream(dispatch_stream)365probs.record_stream(dispatch_stream)366367with torch.cuda.stream(dispatch_stream):368return self.moe_layer.mlp.dispatch(hidden_states, probs)

534# Launch the dispatch, experts, and combine535with quant_context_factory(moe_config, self.moe_layer_idx):536dispatched_input, dispatched_probs = self._launch_dispatch(537route_input, route_probs, async_op=overlap538)539if overlap:540dispatched_input, dispatched_probs = self._wait_dispatch(541dispatched_input, dispatched_probs542)543544output, _ = self.moe_layer.mlp.routed_experts_compute(545dispatched_input, dispatched_probs546)547combined_output = self._launch_combine(output, async_op=overlap)

The parallel block in the figure overlaps dispatch with attention while combine stays exposed, and the shortcut-connected MoE covers most of dispatch and combine with the layers around it. The price falls on the router, which now reads the stream from before the compute it overlaps and never sees that compute's output.

Does the model learn worse for it? Cai et al. and the LongCat-Flash team saw it learn about as well. Cai et al. put an MoE in every block of a LLaMA2-style model, wired as Pos-1, the ScMoE variant that Megatron implements, and it averaged slightly higher on zero-shot benchmarks than a standard top-2 MoE¹⁷³. In LongCat-Flash, the loss curves with and without the shortcut were virtually indistinguishable¹⁷⁴. AMD's FarSkip-Collective rewires the block so that computation goes on with a residual that does not yet hold the latest communicated output. Shipped in Instella-MoE, it matched the MoE block's benchmark average after 200B training tokens^175,176, although the tables of both AMD reports show it losing on MMLU and HumanEval+.

What you get in return is speed. ScMoE trained 1.12x to 1.14x faster on one node of 8 GPUs over NVLink, with the paper's standard top-2 MoE as the baseline¹⁷³, and Instella-MoE gained 12.7% in pre-training throughput¹⁷⁷. LongCat-Flash measured the exposed communication itself: with its shortcut-connected MoE, the share of time spent in dispatch and combine that is not overlapped fell from 25.3% to 8.4%¹⁷⁸, in expert-parallel groups of 32 accelerators that use pipelined all-gather and reduce-scatter instead of all-to-all.

In Megatron, the figure's parallel row (ScMoE's Pos-1) is --moe-shortcut-connection, which makes each MoE read the input of the attention layer just before it, so its router no longer waits for that attention. The shared expert still reads attention's output, a detail the figure's row does not draw, and the routed experts run on the main stream after attention's core and before its output projection, as the row's experts follow its attention. The rewiring alone overlaps nothing: the all-to-alls overlap only once --moe-shortcut-parallel puts them on a side CUDA stream. An optional extra norm on the merged output comes with --moe-shortcut-post-norm. All three flags are only on upstream main, and they work only in the hybrid stack, not together with FSDP, CUDA graphs, full recompute or --moe-shared-expert-overlap. Without a block change, Megatron can also hide the all-to-all behind another micro-batch (3.7 Overlap).

An exposed all-to-all makes the step slow, and a trace shows how much of it is exposed. Take an nsys timeline (--profile with --profile-step-start and --profile-step-end, see 3.7 Profiling) and compare the all-to-all kernels on the side stream with the compute on the main stream: the stretches where the all-to-all runs alone are the exposed share.

Use Megatron's shortcut block (the figure's parallel row) when dispatch is exposed and you can move to Megatron's hybrid stack: ScMoE's Pos-1 design trained 1.14x faster than a standard top-2 MoE on one 8-GPU NVLink node¹⁷³. Set both --moe-shortcut-connection and --moe-shortcut-parallel, since only the second makes the all-to-alls overlap, and leave out FSDP, CUDA graphs, full recompute and --moe-shared-expert-overlap, which do not run with it.

What to watch

With the router, the balancer, the experts and the block layout chosen, the architecture is complete, and what remains is training it. Where should training start, which optimizer should you use, and how will you know when something goes wrong?

3.6 Training: start, optimizer, stability

Training an MoE adds failure modes that a dense model does not have, and the teams behind large MoEs have run into several of them. A Switch Transformer diverged in bfloat16⁴⁰, and two of six baseline ST-MoE runs went unstable⁴¹. Routing can also collapse onto a few experts, an imbalance that feeds itself¹⁷⁹. What makes these runs fail, and how would you notice early that yours is about to?

Training choices reported by model authors. Initialization and optimizer entries highlight the choice relevant to this section; MTP entries describe training, which can differ from the released inference configuration.

The first decision is where training starts, from random weights or from the dense model we already trained in 2.4. What does starting from the dense checkpoint buy you, and what does it cost?

Upcycling

Upcycling starts the MoE from a trained dense checkpoint by copying its MLP into the experts. An expert as wide as the dense MLP gets a full copy, and a narrower one gets a slice: in our run, with a dense MLP 8192 wide and experts of 1024, Megatron splits the dense MLP into 8 slices and copies each slice into 8 of the 64 experts. Either way, the MoE starts where the dense run stopped.

Copying the weights leaves a second choice: whether the optimizer state and the learning-rate schedule carry over too. Komatsuzaki et al. continued the dense run's schedule where it left off¹⁸⁴, but reused the optimizer state only for vision models, because on language it gave no improvement¹⁸⁵. Megatron's upcycling step saves the MoE at iteration 1 without optimizer state (training.py#L3069-L3083), so the MoE run starts with a fresh optimizer and a new schedule.

The head start is not free: it is the compute already spent on the dense run, and a run from scratch eventually catches up with the upcycled run. How soon varies a lot. In OLMoE, a run from scratch caught up after 500B tokens, 25% of the 2T tokens the dense model had been trained on¹⁸⁶, while in Komatsuzaki et al.'s sparse upcycling runs it took 120% of the dense checkpoint's compute¹⁸⁷. The figure below redraws OLMoE's comparison from its public training logs: an MoE upcycled from the dense OLMo-1B after 2T tokens of training, against the same MoE trained from scratch, each for 610B tokens.

In the figure, the upcycled run starts ahead on all four metrics, and the run from scratch narrows the gap only gradually. Its HellaSwag and MMLU accuracies come level with the upcycled run's around 500B tokens, where the authors place the catch-up, while its training and validation losses stay slightly higher until about 600B tokens.

A trained MoE can grow the same way, by copying its own experts instead of a dense MLP. The figure below turns an 8-expert MoE into a 16-expert one and shows what the copies do to routing.

With top-2 routing and each copy inheriting its original's router row, the figure's exact copies tie: every token picks an expert and its copy at equal weight, so it runs one distinct expert instead of two. Perturbing the copy gives the router a gradient between the pair, and random new experts avoid the tie but put routing weight on untrained experts. All three have been used to grow MoEs: Expert Upcycling copies each expert with its router row, going from 32 to 64 experts¹⁸⁸, K-EXAONE 2.0 adds a random rotation to each copy, going from 128 to 256¹⁸⁹, Expert Upcycling also tested Drop-Upcycling, which re-initializes part of each copy's weights, half in the setting its authors found best^190,191, and EMO draws the new experts and router rows from a Gaussian¹⁹². In Megatron, upcycling from a dense checkpoint is --moe-use-upcycling with --moe-upcycling-granularity, and at the pin it needs the two local fixes described in 2.5.

3043if args.moe_use_upcycling:3044torch.distributed.barrier()3045assert not checkpoint_exists(args.save), (3046"The upcycling destination directory already exists. "3047"Please check if --moe-use-upcycling is mistakenly enabled. "3048"Upcycling should only be set for the first run when converting the dense model. "3049"All subsequent runs should remove this flag. "3050)3051# before changing moe related global args, save them in local variables3052num_experts = args.num_experts3053expert_model_parallel_size = args.expert_model_parallel_size3054moe_ffn_hidden_size = args.ffn_hidden_size30553056# set dense model related args in to global args before getting dense model3057args.num_experts = None3058args.expert_model_parallel_size = 13059args.ffn_hidden_size = moe_ffn_hidden_size * args.moe_upcycling_granularity30603061# get dense model3062dense_model_for_upcycling = _build_model_wrapper(wrap_with_ddp=True)30633064# recover moe upcycling related args in global args before executing upcycling3065args.num_experts = num_experts3066args.expert_model_parallel_size = expert_model_parallel_size3067args.ffn_hidden_size = moe_ffn_hidden_size30683069# execute upcycling3070_, args.num_floating_point_operations_so_far = upcycling_utils.load_and_upcycle_model(3071load_checkpoint,3072unwrapped_model,3073dense_model_for_upcycling,3074load_kwargs={3075'model': dense_model_for_upcycling,3076'optimizer': None,3077'opt_param_scheduler': None,3078},3079)3080args.iteration = 13081save_checkpoint(3082args.iteration, model, None, None, args.num_floating_point_operations_so_far3083)

145def _get_moe_activation_scale():146"""147Calc moe activation scale factor relative to dense activation.148For more detail please refer to https://arxiv.org/abs/2410.07524.149"""150if moe_router_pre_softmax:151moe_activation_scale = (expansion_rate * granularity * granularity) / moe_router_topk152else:153moe_activation_scale = granularity154return moe_activation_scale155156def _get_weight_scale():157moe_activation_scale = _get_moe_activation_scale()158if gated_linear_unit == True:159scale = moe_activation_scale ** (1 / 3)160elif activation_func_name == ActivationFuncName.squared_relu:161scale = moe_activation_scale ** (1 / 3)162else:163scale = moe_activation_scale ** (1 / 2)164return scale165166def _process_fc1_weight_param(param):167param = param.clone()168scale = _get_weight_scale()169param = param * scale170if activation_func_name == ActivationFuncName.silu and gated_linear_unit == True:171param_1, param_2 = torch.chunk(param, 2, dim=0)172params_1 = torch.tensor_split(param_1, granularity, dim=0)173params_2 = torch.tensor_split(param_2, granularity, dim=0)174params = [torch.cat([params_1[i], params_2[i]], dim=0) for i in range(granularity)]175else:176params = torch.tensor_split(param, granularity, dim=0)177params = params * expansion_rate178return params

2632def load_checkpoint(2633ddp_model,2634optimizer,2635opt_param_scheduler,

2917def _build_model_wrapper(wrap_with_ddp: bool):2918if cfg_container is not None and getattr(cfg_container, "model", None) is not None:2919from megatron.training.utils import start_memory_history_recording29202921start_memory_history_recording(cfg_container.profiling)29222923cfg = cfg_container2924model_config = cfg.model2925builder_cls = model_config.get_builder_cls()

An upcycle that went wrong loses its head start, so compare its loss with the dense run continued for the same tokens. Since the experts start as identical copies, also watch MaxVio per layer while they drift apart, computed offline from the routing trace (A.1).

Use it when a dense checkpoint exists and the MoE run is short, since a run from scratch caught up after 25% of the dense model's training tokens in OLMoE and 120% of its compute in Komatsuzaki et al.'s runs^186,187. Upcycling is an initialization, so set it on the first run only, and at the pin apply the two local fixes from 2.5 first.

Whether you upcycle or start from random weights only decides where the weights begin; from the first step on, the optimizer decides where they go. Each expert, though, sees only the tokens routed to it, so should your optimizer move an expert that got a handful of tokens as far as a busy one?

Optimizer

AdamW, the optimizer DeepSeek-V3 trained with¹⁶³, scales each weight's update by that weight's own gradient history, so how far an expert moves follows the gradients its tokens produced. Muon, an optimizer for weight matrices used by several recent MoEs, sizes the step another way. It orthogonalizes each update, pushing its nonzero singular values toward one, and then multiplies it by a scale set by the matrix's shape alone, so the size of the step no longer depends on how large the gradient behind it is.

Megatron orthogonalizes each expert's matrix on its own (transformer_engine.py#L2665-L2668, transformer_config.py#L956, emerging_optimizers.py#L147-L149), as DeepSeek-V4 does¹⁹³, so an expert the router rarely picks can move about as far as the busiest one, only in its own direction, and it keeps moving even in a step where it got no tokens, along its momentum. The figure below shows what that means over 49 toy experts, plotting update size against tokens per step for AdamW, for Muon applied to each expert separately, and for Muon over all experts as one matrix.

As an expert's tokens fall in the toy, its update shrinks under AdamW as the gradient turns to noise and, under Muon over all experts as one matrix, with the size of its gradient, while per-expert Muon keeps the busiest expert's step down to a few tokens per step. Below that, the per-expert step shrinks too: each token adds a single direction to the gradient, and with so few tokens even the momentum of recent steps no longer fills the expert's matrix.

The figure shows how Muon moves each expert; whether that trains a better MoE has been measured on whole models. Moonlight, an MoE with 3B active and 16B total parameters, performed significantly better with Muon than its twin trained with AdamW under the same settings, compared at about 1.2T tokens¹⁹⁴.

Once you are on Muon, one choice remains that adopters make differently: which optimizer updates the router. DeepSeek-V4 keeps it on Muon¹⁹³, while Qwen3.8-Next moved it to AdamW after Muon destabilized the router early in training¹⁹⁵. Lau and Su compared router optimizers on two small MoEs, one shaped like OLMoE-1B-7B and a downsized gpt-oss: with Muon on the hidden matrices, AdamW on the router ended slightly worse than row-normalized or left-spectral router updates¹⁹⁶.

In Megatron, Muon is --optimizer muon, and to put the router on Adam instead, add --moe-router-skip-muon, available from Megatron Core v0.20 on.

147def _is_nonlinear_or_embedding(param):148"""True for parameters that should NOT use the emerging optimizer."""149return getattr(param, 'is_embedding_or_output_parameter', False) or len(param.shape) != 2150151152def _is_muon_excluded(param):153"""True for parameters that should use the scalar optimizer instead of Muon."""154return not getattr(param, 'use_muon', True) or _is_nonlinear_or_embedding(param)

88if self.config.moe_router_skip_muon:89setattr(self.weight, 'use_muon', False)90if self.bias is not None:91setattr(self.bias, 'use_muon', False)

303if scale_shape is None:304size = [grad.size(-2), grad.size(-1)]305if partition_dim is not None:306size[partition_dim] *= get_pg_size(tp_group)307else:308# This overrides only the final Muon scalar; NS still uses grad's physical309# shape and partition metadata for its collectives.310size = scale_shape311# Only forward the kwarg when enabled; older emerging_optimizers do not312# accept it at all, and __init__ has already rejected use_syrk on those.313ns_kwargs = {"use_syrk": True} if use_syrk else {}314orth_grad = newton_schulz_tp(315grad,316steps=num_ns_steps,317coefficient_type=coefficient_type,318tp_group=tp_group,319partition_dim=partition_dim,320tp_mode="duplicated" if tp_mode_this_group == "blockwise" else tp_mode_this_group,321**ns_kwargs,322)323scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode)324return orth_grad * scale_factor * extra_scale_factor

1792assert not (1793args.use_layer_wise_distributed_optimizer and args.moe_single_grouped_weight1794), (1795"The LayerWise distributed optimizer does not support --moe-single-grouped-weight: "1796"Muon semantics for a single grouped [E, N, K] expert weight are not defined. "1797"Disable --moe-single-grouped-weight or use Adam/DistributedOptimizer."1798)

Under Muon, experts can collapse¹⁹⁷, and too large an update scale can make the run diverge¹⁹⁸. Track what each expert receives and produces: tokens per expert per step, offline from the routing trace (A.1), and, per layer, the largest and the smallest expert output norm, each divided by the layer's median, from a forward hook of ours (A.1).

Use Muon when you can afford to tune it. Start at --muon-extra-scale-factor 0.2, since Megatron's default of 1.0 diverged on a dense 7B in Mellum2's tests¹⁹⁸, and choose the router's optimizer with --moe-router-skip-muon, a call adopters make both ways. It does not run with Megatron-FSDP or Torch-FSDP2.

However the optimizer moves the weights, each step still learns only from predicting the next token. Could you get more signal out of each step by predicting further ahead, and what would that cost in step time?

MTP

Multi-token prediction adds one transformer layer per extra token to predict. That layer takes the main model's hidden state together with the next token's embedding, and predicts the token after next through the shared output head. DeepSeek-V3 found it better on most benchmarks at both of its test scales, and the extra layer can be discarded at inference¹⁹⁹. The figure below follows one position of the sequence through that wiring.

On their way in, the two inputs each go through their own RMSNorm, and the MTP layer concatenates them and projects the result from twice the hidden size back to it before its transformer block. Both predictions come out of the same output head, and the MTP loss is weighted and added to the main loss¹⁹⁹.

The extra layer costs step time, and under pipeline parallelism that cost lands in one place: unless the layout moves it, Megatron puts the MTP layer on the last stage (multi_token_prediction.py#L886-L905), which already holds the output head the MTP layer reuses. On our 256-node runs, early steps took about 37% longer with one MTP layer than without, and spreading the model over 8 pipeline stages did not close that gap. Mellum2 reports 7% extra training time for a 14B MoE²⁰⁰.

1709hidden_states = apply_module(self.hnorm)(hidden_states)1710hidden_states = make_viewless_tensor(1711inp=hidden_states, requires_grad=True, keep_graph=True1712)1713hidden_states = torch.cat((decoder_input, hidden_states), -1)1714hidden_states, _ = self.eh_proj(hidden_states)

1276mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers

932# without custom PP layout, we only support put all of MTP layers on the last pipeline stage933if mtp_num_layers is not None:934if pp_size is None:935# Compatibility fallback for callers without explicit pipeline metadata.936pp_size = parallel_state.get_pipeline_model_parallel_world_size()937mtp_on_this_rank = pp_rank == pp_size - 1

2598if self.config.mtp_detach_heads:2599hidden_states = hidden_states.detach()

Much of that gap may come from where the layer sits. A pipeline step waits for its slowest stage, so a layer added to one stage slows the step by its share of that stage's work, not by its much smaller share of the whole model. Ling 2.0's team ran into the same imbalance: they measured an MTP layer at about 1.7 times the compute of a standard MoE layer, and repartitioning the pipeline around it gained 40% end to end²⁰¹.

Megatron turns MTP on with --mtp-num-layers, the number of extra depths to add, and --mtp-loss-scaling-factor weights the MTP loss: DeepSeek-V3 used 0.3 for its first 10T tokens and 0.1 after¹⁶³, and Mellum2 used 0.1²⁰⁰. To choose which stage holds the MTP layer, written m, use --pipeline-model-parallel-layout. If you want the MTP layer only as a draft model for speculative decoding, --mtp-detach-heads detaches its inputs, so the MTP loss trains the MTP layer and leaves the main model alone.

Megatron logs an mtp_1 loss next to the lm loss whenever MTP is on, while the MTP layer's cost shows elsewhere, in step time and in the memory of the stage that holds it. To find the slow stage, log forward and backward compute time per pipeline rank with --timing-log-level 2 --timing-log-option all (A.1); the stage holding the MTP layer should show up as the slowest rank.

Use it when the denser training signal DeepSeek-V3 reports¹⁹⁹ is worth the step time. Under pipeline parallelism, place the MTP layer with --pipeline-model-parallel-layout, or the last stage carries it.

Each part so far has told you what to watch for its own choice. Some failures, though, belong to no single choice and can start anywhere in the model, so what should every run have in place to catch them?

Stability

An MoE run can fail loudly or quietly. A loud failure shows in the loss, which spikes or turns into NaN, while a quiet one, such as routing collapse⁷⁶ or a rogue expert¹⁵, can leave the loss flat.

The loudest failure is also the easiest to catch, since Megatron checks every step's loss and gradients for a NaN by default (3.0). What you want to know next is whether to blame the hardware, so when the check fires, Megatron reruns the step (--rerun-mode validate_results, also on by default). The first rerun, on the same GPU, tells a transient error from one that repeats. If the NaN repeats, Megatron exits with code 16 so that the job restarts from a checkpoint on another GPU, where the second rerun tells a faulty GPU from a NaN that the data or the model really produced.

Bad hardware and bad batches do not always end in a NaN, because they can also return a finite loss or gradient with a far too large exponent. Megatron catches those with --check-for-spiky-loss and --check-for-large-grads: each flags a value 10x the largest of the first 100 it has seen, reruns the step once on the same GPU, logs whether the value came back, and keeps training.

All of these checks react to a step that has already gone wrong. To stop the loss from breaking in the first place, 3.1 already protects the router, with fp32 math and a z-loss that holds its logits down, and recent large MoEs add guards of their own. Some cap a value before it can grow: DeepSeek-V3 and Step 3.5 Flash clip the gradient norm^163,183, DeepSeek-V4.1-Flash clamps the experts' SwiGLU activations¹⁸⁰, and Step 3.5 Flash clips expert activations as well¹⁸³. Others keep the most sensitive operations in higher precision: DeepSeek-V3 kept its embedding, output head, MoE gating, normalization and attention in BF16 or FP32 during FP8 training¹⁶³, and Kimi K3 keeps the attention output in FP32¹⁸¹. LongCat-Flash carries the z-loss idea over to the hidden states¹⁸².

The quiet failures may not show in the loss at all, so you have to look for them elsewhere. The figure below simulates three of them next to a healthy run: routing collapse, a dead expert and a rogue expert.

Since the simulation keeps the loss flat through all three failures, only the other signals tell them apart from the healthy run. Routing collapse drives MaxVio up. A dead expert pulls the smallest expert output norm far below the layer's median, and a rogue expert pushes the largest far above it, from about 14 to about 5,700 times the median as in Step 3.5 Flash's layer 45¹⁵. The loss alarm, 0.1 above the lowest loss so far, borrows the margin of Qwen3.8-Next's spike criterion²⁰².

263# Check individual rank losses are not NaN prior to DP all-reduce.264rerun_state_machine = get_rerun_state_machine()265if args.check_for_nan_in_loss_and_grad:266rerun_state_machine.validate_result(267result=loss,268rejection_func=torch.isnan,269message="found NaN in local forward loss calculation",270tolerance=0.0,  # forward pass calculations are deterministic271fatal=True,272)273rerun_state_machine.validate_result(274result=loss,275rejection_func=torch.isinf,276message="found Inf in local forward loss calculation",277tolerance=0.0,  # forward pass calculations are deterministic278fatal=True,279)280# Check for spiky loss281if args.check_for_spiky_loss:282rerun_state_machine.validate_result(283result=loss,284rejection_func=partial(285rerun_state_machine.is_unexpectedly_large,286threshold=SPIKY_LOSS_FACTOR,287context="loss",288),289message="Spiky loss",290tolerance=0.0,  # forward pass calculations are deterministic291fatal=False,292)

389def check_grads(self, check_for_nan_or_inf, check_for_large):390"""391Make sure norm of grads in bucket are not NaN prior to data-parallel392all-reduce / reduce-scatter.393"""394rerun_state_machine = get_rerun_state_machine()395for i in range(len(self.buckets)):396grad_norm = self.buckets[i].grad_data.norm(p=2)397# check for NaN, Inf and unexpectedly large grads398if check_for_nan_or_inf:399rerun_state_machine.validate_result(400result=grad_norm,401rejection_func=torch.isnan,402message=f"found NaN in local grad norm for bucket #{i} "403f"in backward pass before data-parallel communication collective",404tolerance=0.001,  # 0.1% tolerance to account for non-deterministic FA backward405fatal=True,406)407rerun_state_machine.validate_result(408result=grad_norm,409rejection_func=torch.isinf,410message=f"found Inf in local grad norm for bucket #{i} "411f"in backward pass before data-parallel communication collective",412tolerance=0.001,  # 0.1% tolerance to account for non-deterministic FA backward413fatal=True,414)415if check_for_large:416rerun_state_machine.validate_result(417result=grad_norm,418rejection_func=partial(419rerun_state_machine.is_unexpectedly_large, threshold=10, context="grads"420),421message=f"found unexpected large grads in bucket #{i} "422f"in backward pass before data-parallel communication collective",423tolerance=0.001,  # 0.1% tolerance to account for non-deterministic FA backward424fatal=False,425)

752value: float = math.fabs(result.item())753# Ignore NaNs and Infs. They should be checked separately.754if math.isnan(value) or math.isinf(value):755return False756757if resample or context not in self.large_value_counts:758self.large_value_counts[context] = 0759if self.large_value_counts[context] < num_samples:760self.large_value_counts[context] += 1761self.max_values[context] = max(self.max_values.get(context, 0.0), value)762if self.large_value_counts[context] == num_samples:763logger.warning(f"Max value for {context}: {self.max_values[context]}")764return False765766return value >= self.max_values[context] * threshold

622if self.state == RerunState.RERUNNING_IN_PLACE:623if comparison > tolerance:624if not fatal:625self.continue_requested = True626log_failure(627"First rerun: unexpected result is not reproducible within the tolerance "628f"({result} != {self.initial_result})",629fatal=fatal,630)631self._log_validation_error_to_file(632status=RerunValidationStatus.FIRST_RERUN_NOT_REPRODUCIBLE,633result=result,634message=message,635)636log_failure("Possible transient error!", fatal=fatal)637638else:639if fatal:640self.checkpoint_requested = True641else:642self.continue_requested = True643# Remember the node and device we're running on so that we can check we're not644# rerunning on the same GPU when we resume from the checkpoint.645self.suspicious_node = os.uname()[1]646self.suspicious_device = torch.cuda.current_device()647self._log_validation_error_to_file(648status=RerunValidationStatus.FIRST_RERUN_REPRODUCIBLE,649result=result,650message=message,651)652log_failure(653"First rerun: unexpected result is reproducible within the tolerance "654f"({result} = {self.initial_result}). "655"Need to rerun on a different GPU to verify correctness.",656fatal=fatal,657)658# This is the second re-run.659elif self.state == RerunState.RERUNNING_FROM_CHECKPOINT:660# Ensure we're not on the same GPU as the first rerun.661node = os.uname()[1]662device = torch.cuda.current_device()663if node == self.suspicious_node and device == self.suspicious_device:664logger.error(665f"Got rescheduled on the same GPU. Need to resume again from the same "666f"checkpoint (node: {self.suspicious_node}, gpu: {self.suspicious_device})"667)668self.restart_again_requested = True669elif comparison > tolerance:670self._log_validation_error_to_file(671status=RerunValidationStatus.SECOND_RERUN_NOT_REPRODUCIBLE,672result=result,673message=message,674)675logger.warning(676"Second rerun: unexpected result is not reproducible on a different GPU, "677f"therefore was likely incorrect ({result} != {self.initial_result})"678)679log_failure("Possible persistent error!")680if not fatal:681self.continue_requested = True... 13 more lines

For spikes, watch the router's logits through the z-loss (3.1). The failures that leave the loss flat need signals of their own: for routing collapse, MaxVio per layer, computed offline from the routing trace (A.1), next to the load-balancing loss, and for dead and rogue experts, the largest and the smallest expert output norm in each layer over the layer's median, from a hook of ours (A.1).

Use it always: keep the default NaN check and --rerun-mode validate_results.

What to watch

A run that trains cleanly can still spend a large share of each step waiting on communication rather than computing. Where does the time in an MoE step go, and how much of it can you get back?

3.7 Systems: where the time goes and how to get it back

Every MoE layer sends its tokens to their experts and brings them back, and on a real run that traffic can take much of the step. In the one step we traced on one node, half the time was exposed all-to-all, while the router was still unbalanced. The figure below zooms in on those two transfers, dispatch out to the experts and combine back, which run over NVLink inside a node and over the fabric across nodes.

You can place the experts so that less traffic crosses the slower network between nodes, and hide the traffic that remains behind other compute. You can keep the memory that costs in check, cap how many nodes each token reaches, and trade dropped tokens or precision for speed. Each part comes from its own setup, so do not add their gains together; test them as alternatives, one at a time.

Systems choices documented in model reports. Parallelism and overlap describe training; post-training quantization is labeled separately. These are model-specific implementations, not interchangeable Megatron flag settings.

Before reaching for any of these levers, you need to see where your own step spends its time, since each kind of waste calls for a different fix. What does one profiled step on a single node show?

Profiling

Megatron can record a trace of any step for you. With --profile alone it records an nsys trace, and --use-pytorch-profiler switches it to a torch profiler trace (argument_utils.py#L706-L709). In either case, --profile-step-start, --profile-step-end and --profile-ranks pick which steps and ranks to trace. The step we traced comes from a 7B-A1B MoE, with 7B parameters in total and 1B active per token. Its 16 layers hold 64 experts each with top-8 routing, and it trained on one node of 8 H100s. Tracing it cost one short job, which traced iteration 4 and had moved past it within about six minutes of launch. The figures below break that step down and annotate its trace.

The all-to-all took half of that step, but how much of that half was the network actually moving data? When one GPU's experts get more tokens than the rest, that GPU finishes its expert compute late, and NVIDIA's Megatron Core team reports that the other ranks then wait for it at synchronization points²⁰⁸. In an MoE layer the all-to-all can be one of those points, so part of its time in a trace can be waiting rather than moving data. Our step had a load-balancing loss of 3.2, where 1 is balanced, so its router was still unbalanced. A separate run of the same config, with routing forced even, ran that same iteration 40% faster.

4955if (4956cfg.profiling.use_pytorch_profiler4957and (len(cfg.profiling.profile_ranks) == 0 or4958torch.distributed.get_rank() in cfg.profiling.profile_ranks)4959):4960if cfg.profiling.pytorch_profiler_collect_chakra:4961et_dir = Path(f"{cfg.logger.tensorboard_dir}/../chakra")4962et_dir.mkdir(parents=True, exist_ok=True)4963et = torch.profiler.ExecutionTraceObserver().register_callback(f"{et_dir}/rank-{torch.distributed.get_rank()}.json.gz")4964else:4965et = None4966def trace_handler(p):4967profile_dir = Path(f"{cfg.logger.tensorboard_dir}/../torch_profile")4968profile_dir.mkdir(parents=True, exist_ok=True)4969p.export_chrome_trace(f"{profile_dir}/rank-{torch.distributed.get_rank()}.json.gz")4970prof = torch.profiler.profile(4971schedule=torch.profiler.schedule(4972wait=max(cfg.profiling.profile_step_start - 1, 0),4973warmup=1 if cfg.profiling.profile_step_start > 0 else 0,4974active=cfg.profiling.profile_step_end - cfg.profiling.profile_step_start,4975repeat=1,4976),4977on_trace_ready=trace_handler,4978record_shapes=cfg.profiling.pytorch_profiler_collect_shapes,4979with_stack=cfg.profiling.pytorch_profiler_collect_callstack,4980execution_trace_observer=et,4981)4982prof.start()

5041if ((cfg.profiling.use_nsys_profiler or cfg.profiling.use_pytorch_profiler)5042and (len(cfg.profiling.profile_ranks) == 0 or5043torch.distributed.get_rank() in cfg.profiling.profile_ranks)):5044# Enable NVTX range when profiling starts and nvtx_ranges is set.5045if iteration == cfg.profiling.profile_step_start and cfg.profiling.nvtx_ranges:5046configure_nvtx_profiling(True)5047if cfg.profiling.use_pytorch_profiler:5048prof.step()5049elif iteration == cfg.profiling.profile_step_start:5050torch.cuda.check_error(torch.cuda.cudart().cudaProfilerStart())5051if cfg.profiling.record_shapes:5052nsys_nvtx_context = torch.autograd.profiler.emit_nvtx(record_shapes=cfg.profiling.record_shapes)5053nsys_nvtx_context.__enter__()

On your own run, start with the balance signals of 3.2. MaxVio, computed from the token counts that --save-tokens-per-expert-interval writes, shows how uneven the load is across experts. Summing those counts over each GPU's experts shows which rank receives the most tokens, and so which rank the others most likely wait for. To see how much of the all-to-all that imbalance costs, re-trace with --moe-router-force-load-balancing, which swaps the router's logits for random ones so that every expert gets about the same load (router.py#L1185-L1187). If the all-to-all shrinks under forced balance, imbalance is costing you time that better balancing (3.3) can win back, and what remains is mostly the transfer, which the rest of this section works on.

Traces also show which MoE kernels are worth fusing. Before dispatch, Megatron reorders the tokens so that each expert's tokens sit together (token_dispatcher.py#L668-L682), and after combine it puts them back in order (#L893-L900). With --moe-permute-fusion, the permute and copy kernels that do this fall from 11.45% to 7.14% of the step, and step time drops 3.3%.

Judge each change by step time or TFLOP/s. A trace then tells you why the number moved: read the share of the step that the all-to-all leaves exposed.

Use it once for every new config: trace one step at the node count you will run. If the all-to-all takes a large share of it, check MaxVio and re-trace with --moe-router-force-load-balancing: fix balancing first (3.3) if the all-to-all shrinks, and work on Interconnect and Overlap below if it does not. If permute and copy kernels take a large share, try --moe-permute-fusion.

The step we traced kept its all-to-all inside one node, on NVLink, and a larger run cannot always do that. What happens to the step when an expert-parallel group spans two nodes and its all-to-all has to cross the network between them?

Interconnect

Between nodes, the all-to-all leaves NVLink for the cluster's network fabric, InfiniBand or AWS's EFA for instance, so the first thing to check is that NCCL uses that fabric at all. It can fall back to plain TCP sockets even when the fabric is requested: two of our runs set FI_PROVIDER=efa and named an OFI plugin, NCCL still logged "Using network Socket", and they ran about 25x slower than our EP16 run on EFA. The next figure shows the settings we tried, in order.

In the figure, NCCL reached EFA only once the host's RDMA libraries came first on the library path; naming the plugin, by path or by name, left it on sockets. With the fabric working, what is left to choose is the expert-parallel size, --expert-model-parallel-size, which decides how far the all-to-all has to travel. With expert tensor parallelism at 1 and the default rank order, Megatron forms each EP group from consecutive ranks (parallel_state.py#L921-L929). An EP size that divides the number of GPUs in a node then keeps the all-to-all inside that node, on NVLink, and a larger one spreads it over several nodes. Our nodes hold eight H100s each, so the figure below compares EP8 in one node with EP16 across two.

Spread over two nodes, EP16 ran 59% below EP8 in one node, even with EFA working, although the two runs also differ in their nodes and in how many micro-batches each step accumulates.

920expert_order = _inject_gtp_remat_axis(order, after="ep")921expert_decoder_rank_generator = RankGenerator(922tp=expert_tensor_parallel_size,923ep=expert_model_parallel_size,924dp=expert_data_parallel_size,925pp=pipeline_model_parallel_size,926cp=1,927order=expert_order,928rank_offset=rank_offset,929gtp_remat=expert_gtp_remat_size,930)

1419for ranks in expert_decoder_rank_generator.get_ranks('ep'):1420group = create_group(1421ranks,1422timeout=timeout,1423pg_options=get_nccl_options("ep", nccl_comm_cfgs),1424group_desc="EXPERT_MODEL_PARALLEL_GROUP",1425)1426if rank in ranks:1427_EXPERT_MODEL_PARALLEL_GROUP = group1428_EXPERT_MODEL_PARALLEL_RANKS = ranks

A fallback to sockets slowed our two-node run by the factor above, so check the startup log first. With NCCL_DEBUG=INFO, NCCL prints "Using network" and the transport it picked. You want your fabric's transport there, such as Libfabric for the OFI plugin on EFA, while "Using network Socket" means the fallback. Once the fabric is confirmed, compare alltoall bus bandwidth inside and across nodes from nccl-tests (2.1). Then measure TFLOP/s per GPU at an EP size that fits in one node and at one that spans two.

Use expert parallelism inside a node when the model fits that way: start there, add nodes as data-parallel replicas, and time both placements at your node count before you commit.

Even on a working interconnect, the tokens that cross it still sit on the critical path while the GPUs wait. Is there other work the GPUs could do while the tokens are in flight?

Overlap

We've already seen in 3.5 Block wiring how changing the architecture hides the all-to-all inside one micro-batch. A parallel block can run dispatch while attention computes, and a shortcut-connected MoE can also hide combine behind the dense compute that runs beside it. Rewiring changes what the router sees, however, and you may not want to touch an architecture you have already settled on. Without rewiring, the exposed all-to-all is hard to hide inside one micro-batch. A training step, however, runs many micro-batches, and the work of one does not wait on another until their gradients are summed.

Megatron borrows that work: --overlap-moe-expert-parallel-comm runs one micro-batch's forward beside another's backward, so the all-to-all of one hides behind the compute of the other (model_parallel_config.py#L349-L352). Pipeline schedules also pair one micro-batch's forward with another's backward, as the Ultra-Scale Playbook describes, and DeepSeek-V3's DualPipe uses that pairing to hide its all-to-alls²⁰⁹. The figure below draws one MoE layer with and without that interleaving.

Drawn with one micro-batch after the other, the layer leaves all four all-to-alls exposed. The interleaved drawing also sets --delay-wgrad-compute, which Megatron's MoE guide pairs with the overlap (README.md#L434). That flag computes the weight gradients later, so the backward dispatch waits only for the experts' input gradients and runs under that micro-batch's own delayed expert weight gradients (model_chunk_schedule_plan.py#L266-L268), while the other three all-to-alls run under compute from the other micro-batch. On our runs, which set the overlap flag alone, the overlap raised throughput by up to 35% on two nodes. It also costs memory: rank 0 reserved 5 to 10 GB more, while its peak allocated memory barely moved.

54"""Scheduler for 1f1b with no pipelining.5556This function schedules micro-batches in a way that the forward pass of Transformer layers57for one micro-batch runs in parallel with the backward pass of another.58Each layer's forward and backward operations are co-scheduled to maximize the overlap of59their computations and communications.60EP A2A in forward step is hidden by the attention/mlp computation in the backward step,61and vice versa.62Assuming we have 4 microbatches, the schedule is as follows:63Phases 0: 1st microbatch forward64Phases 1: 1st microbatch backward + 2nd microbatch forward65Phases 2: 2nd microbatch backward + 3rd microbatch forward66Phases 3: 3rd microbatch backward + 4th microbatch forward67Phases 4: 4th microbatch backward

226def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False):227"""Schedule one-forward-one-backward operations for a single transformer layer.228229This function interleaves forward and backward operations, overlapping the communications230(dispatch or combine) of one with the computations (pre_dispatch or mlp) of the other231to maximize parallelism and efficiency.232233When f_layer and b_layer are not None, forward and backward pass are overlapped as follows:234comm_stream: combine_bwd       | dispatch_fwd->dispatch_bwd  | combine_fwd235comp_stream: pre_dispatch_fwd  | mlp_bwd->mlp_bwd_dw->mlp_fwd| pre_dispatch_bwd236For MTP, mtp_post_process_fwd is executed after the combine_fwd in the comp_stream,237and mtp_post_process_bwd is executed before the combine_bwd in the comp_stream.238239Args:240f_layer (TransformerLayerSchedulePlan): Forward layer (for current microbatch)241b_layer (TransformerLayerSchedulePlan): Backward layer (for previous microbatch)242f_input (Tensor): Input for forward computation243b_grad (Tensor): Gradient for backward computation244is_last_layer_in_bwd (bool):245Whether the current layer is the last layer in the backward pass.246247Returns:248Functions or values for next iteration's computation249"""250251if b_layer is not None:252b_grad = b_layer.mtp_post_process.backward(b_grad)253b_grad = b_layer.moe_combine.backward(b_grad)254255if f_layer is not None:256with f_layer.get_low_precision_context():257f_input = f_layer.pre_dispatch_computation.forward(f_input)258259if b_layer is not None:260b_grad = b_layer.mlp.backward(b_grad)261262if f_layer is not None:263with f_layer.get_low_precision_context():264f_input = f_layer.moe_dispatch.forward(f_input)265266if b_layer is not None:267b_layer.mlp.backward_dw()268b_grad = b_layer.moe_dispatch.backward(b_grad)269270if b_layer is not None and b_layer.config.ep_overlap_early_attn_memory_release:271b_grad = b_layer.pre_dispatch_computation.backward(b_grad)272273if f_layer is not None:274with f_layer.get_low_precision_context():275f_input = f_layer.mlp.forward(f_input)276277if f_layer is not None:278with f_layer.get_low_precision_context():279f_input = f_layer.moe_combine.forward(f_input)280281if b_layer is not None and not b_layer.config.ep_overlap_early_attn_memory_release:282b_grad = b_layer.pre_dispatch_computation.backward(b_grad)283284if f_layer is not None:285with f_layer.get_low_precision_context():... 9 more lines

Log TFLOP/s per GPU (--log-throughput) before and after turning the overlap on, and keep it only if the step gets faster. Then look at a trace to see which all-to-alls are overlapped with compute.

Use it when the all-to-all is exposed and memory allows: add --overlap-moe-expert-parallel-comm --delay-wgrad-compute, the pair Megatron's MoE guide recommends (README.md#L434), instead of --moe-shared-expert-overlap. --delay-wgrad-compute runs only with the overlap on and with Transformer Engine layers (transformer_config.py#L3631-L3635, arguments.py#L1956-L1958). With pipeline parallelism, the overlap also needs virtual pipeline stages (transformer_config.py#L3565-L3599).

The overlap wins back time but costs memory. How can you keep a run's memory in check without giving that time away again?

Memory

An MoE changes memory use in two ways, through its weights and through its activations. Its weights hold every expert, although each token uses only k of them, and expert parallelism spreads those experts so that each GPU keeps only its own. Its activations, the intermediate results a layer keeps for the backward pass, grow with top-k. The dispatcher copies each token's hidden state once per selected expert (token_dispatcher.py#L559), and every copy keeps its own expert activations, so a GPU whose experts receive more tokens also holds more of them.

The usual answer is activation recomputation, which discards activations after the forward pass and recreates them when the backward pass needs them. The Ultra-Scale Playbook covers its two classic forms, full recompute of every layer and selective recompute of attention. Full recompute runs each layer's forward twice, which at micro-batch 2 cost us 24% of throughput, and the overlap from 3.7 Overlap needs it turned off (transformer_config.py#L3579-L3581).

A finer approach starts from what each module of a layer saves for its backward. Some of those tensors are cheap to recreate, such as the outputs of the normalizations and of the experts' activation function, which take a few elementwise operations. Others are large and costly to recompute: the query, key and value that enter attention, and the dispatched tokens that enter the experts' first matrix multiply, k copies of each token. Megatron treats the two kinds differently²¹⁰, and the next figure marks which saved tensor each of its options acts on.

The blue labels are recomputation options. Megatron's guide recomputes the norms and the expert activation (--recompute-granularity selective --recompute-modules layernorm moe_act): it discards their cheap outputs and recreates them just before the backward that reads them (random.py#L1042-L1064). The outlined labels belong to fine-grained activation offloading, and the guide offloads attention's inputs and output and the experts' first input (--fine-grained-activation-offloading --offload-modules core_attn attn_proj expert_fc1).

Offloading copies the inputs of those large modules to pinned CPU memory on a separate CUDA stream while later modules compute (fine_grained_activation_offload.py#L1262-L1269). Each time the backward finishes an offloaded module, it starts copying back, on a second stream, the tensors of the next one, so they are on the GPU before they are needed (#L1332-L1342). Recomputation pays in compute, while offloading pays in bandwidth between host and GPU, which these copies hide behind compute. The figure below runs three layers through a forward and a backward pass with every saved input kept, with full recompute, and with that split, and tracks the saved inputs each setting holds on the GPU.

With only three layers, one micro-batch and illustrative durations, full recompute reaches the lowest peak but adds a second forward pass, while the fine-grained split holds more and adds only a few small recomputes, because its copies run under compute. The last layer keeps its large inputs on the GPU, since its backward starts as soon as its forward ends and a reload would stall it (fine_grained_activation_offload.py#L630-L637). Slow the host link in the simulation and the copies to the CPU finish late, so the peak rises; slow it further and the backward waits for its reloads.

319self.activation_recompute = (320self.config.recompute_granularity == 'selective'321and "moe_act" in self.config.recompute_modules322)

1130if self.activation_recompute:1131self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput()1132with moe_act_manager as fc1_output:1133bias_act_output = self.activation_checkpoint.checkpoint(1134bias_act_func, fc1_output, bias_parallel, permuted_probs1135)1136else:1137with moe_act_manager as fc1_output:1138bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs)1139output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert)1140if self.activation_recompute:1141self.activation_checkpoint.discard_output_and_register_recompute(output)

1262def on_group_commit_forward(self, name, forced_released_tensors):1263"""Called at the end of a layer group's forward pass to trigger offloading."""1264if not self.do_offload:1265return1266debug_rank(f"--on_group_commit_forward {name}")1267# Wait for compute to finish before starting offload1268self.d2h_stream.wait_stream(torch.cuda.current_stream())1269self.bulk_offload(name, forced_released_tensors)

1332def on_group_start_backward(self):1333"""1334Called at the start of a layer group's backward pass.1335Triggers reloading of tensors from CPU.1336"""1337if not self.do_offload:1338return1339debug_rank(f"--on_group_start_backward {self}")1340# Wait for compute to finish before starting reload1341self.h2d_stream.wait_stream(torch.cuda.current_stream())