Skip to navigation
Recipes & E2E ExamplesFine-Tune Large MoE LLMs

Fine-Tune Large MoE LLMs

View as Markdown

Introduction

Mixture-of-Experts (MoE) architectures have become the dominant design for frontier language models, activating only a fraction of their total parameters per token to deliver strong performance at reduced compute cost. This guide walks through fine-tuning four example MoE LLMs with NVIDIA NeMo AutoModel. For a full list of supported architectures, see the LLM model coverage page.

ModelHugging Face CheckpointValidated Using
GLM-5zai-org/GLM-5256 H100 GPUs (32 nodes x 8)
MiniMax-M2.5MiniMaxAI/MiniMax-M2.564 H100 GPUs (8 nodes x 8)
Step-3.5 Flashstepfun-ai/Step-3.5-Flash64 H100 GPUs (8 nodes x 8)
DeepSeek-V3.2deepseek-ai/DeepSeek-V3.2256 H100 GPUs (32 nodes x 8)

To set up your environment to run NeMo AutoModel, follow the installation guide.

Data

HellaSwag Dataset

All four recipes use the HellaSwag dataset, a commonsense natural language inference benchmark where the model must predict the most plausible continuation of a given scenario.

  • Source: rowan/hellaswag
  • Split: train (used for both training and validation in these recipes)
  • Task: Next-token prediction on commonsense sentence completions

For details on how to swap in your own dataset, see the LLM Dataset Guide and the Dataset Overview.

Recipes

MiniMax-M2.5

examples/llm_finetune/minimax_m2/minimax_m2.5_hellaswag_pp.yaml. Validated using 64 H100 GPUs (8 nodes x 8).

Key distributed settings:

distributed:
strategy: fsdp2
pp_size: 2
ep_size: 32
pipeline:
pp_schedule: interleaved1f1b
layers_per_stage: 2

GLM-5

examples/llm_finetune/glm/glm_5_hellaswag_pp.yaml. Validated using 256 H100 GPUs (32 nodes x 8).

Key distributed settings:

distributed:
strategy: fsdp2
pp_size: 4
ep_size: 64
activation_checkpointing: true
pipeline:
pp_schedule: interleaved1f1b
layers_per_stage: 2

Step-3.5 Flash (StepFun)

examples/llm_finetune/stepfun/step_3.5_flash_hellaswag_pp.yaml. Validated using 64 H100 GPUs (8 nodes x 8).

Key distributed settings:

distributed:
strategy: fsdp2
pp_size: 2
ep_size: 32
pipeline:
pp_schedule: interleaved1f1b
layers_per_stage: 2

DeepSeek-V3.2

examples/llm_finetune/deepseek_v32/deepseek_v32_hellaswag_pp.yaml. Validated using 256 H100 GPUs (32 nodes x 8).

Key distributed settings:

distributed:
strategy: fsdp2
pp_size: 4
ep_size: 64
activation_checkpointing: true
pipeline:
pp_schedule: interleaved1f1b
layers_per_stage: 2

Configure HybridEP Routing

Several large-scale MoE recipes use the HybridEP dispatcher to combine expert-parallel token permutation and communication. Select it in the model backend:

model:
backend:
experts: torch_mm
dispatcher: hybridep

HybridEP supports two routing-metadata paths:

  • Dense routing (default): Builds a Boolean routing map with shape [tokens, experts] and a matching dense probability tensor.
  • Compact routing (opt-in): Passes top-k expert indices with shape [tokens, top_k] instead of the dense Boolean routing map. NeMo AutoModel still expands the routing probabilities to [tokens, experts] to preserve the existing router-gradient layout.

Enable compact routing only after benchmarking the exact model, batch shape, GPU, and expert-parallel topology:

model:
backend:
experts: torch_mm
dispatcher: hybridep
dispatcher_hybridep_compact_routing: true

Compact routing reduces routing-map size, but its additional metadata preprocessing can increase latency for some workloads. Dense routing therefore remains the default.

The following advanced settings apply only when dispatcher: hybridep:

SettingDefaultDescription
dispatcher_hybridep_compact_routingfalsePass top-k expert indices instead of a dense Boolean routing map.
dispatcher_hybridep_permute_fusionfalseFuse token permutation into dispatch and unpermutation into combine.
dispatcher_hybridep_num_sms_preprocessingnullOverride the number of SMs used for routing-metadata preprocessing.
dispatcher_hybridep_num_blocks_permutenullOverride the number of dispatch permutation blocks.
dispatcher_hybridep_num_blocks_unpermutenullOverride the number of combine unpermutation blocks.
dispatcher_capacity_factornullCapacity mode for dynamic routing: size the dispatch buffers once from the first real batch times this factor and run later dispatches non-blocking; an overflow trips a device-side assert. Not for experts with expert_bias.
dispatcher_equal_token_countsfalseDeclare that every EP rank dispatches the same row count, so the per-dispatch pad-size all-reduce and its host sync are skipped. Keep false for variable-length or packed inputs: with genuinely unequal (aligned) row counts the HybridEP collective hangs silently instead of raising.

Leave the numeric tuning settings unset to use HybridEP defaults unless a matched benchmark demonstrates an improvement. A value that improves performance on one GPU architecture or topology might reduce performance on another.

Launch Training

NeMo AutoModel supports direct Slurm submission, local interactive runs through the NeMo AutoModel CLI, and external torchrun. For full details on Slurm batch jobs, multi-node configuration, environment variables, and other launch options, see the Run on a Cluster guide.

NeMo AutoModel CLI

uv run automodel your_single_node_config.yaml --nproc-per-node 8

The CLI takes the recipe config as its positional argument. Use this form only with a config whose distributed topology fits one node; the four configs above use multi-node PP/EP sizes and should be launched through Slurm or external torchrun as shown below.

Torchrun

export TRANSFORMERS_OFFLINE=1
export HF_HOME=your/path/to/hf_cache
export HF_DATASETS_OFFLINE=1
export WANDB_API_KEY=your_wandb_key
export MASTER_ADDR=${MASTER_ADDR:?Set MASTER_ADDR to the rank-0 hostname}
export MASTER_PORT=${MASTER_PORT:-29500}
torchrun --nproc-per-node=8 \
--nnodes=32 \
--rdzv_backend=c10d \
--rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
-m nemo_automodel.recipes.llm.train_ft \
-c examples/llm_finetune/glm/glm_5_hellaswag_pp.yaml \
--model.pretrained_model_name_or_path=/your/local/model_weights

This command uses the GLM-5 config and its model.pretrained_model_name_or_path field. For another recipe, replace the -c path and --nnodes, then update the model and tokenizer path fields defined by that YAML.

Before You Start

  • Hugging Face applies rate limits on downloads. Clone the model repository to your local filesystem beforehand.
  • Ensure your Hugging Face cache (HF_HOME) is configured and that the dataset is already cached locally.
  • To enable Weights & Biases logging, set your WANDB_API_KEY and configure the wandb section in the YAML file.