Expert Upcycling (GitHub repository)

Expert Upcycling is a technique that, instead of training a MoE model large from scratch, starts with a small E-expert model and expands it to mE experts partway through training. The core idea is to duplicate already-trained experts to increase capacity, then use continued pre-training (CPT) to drive role differentiation among the duplicates.

The operation is organized into three stages.

• Expert replication: gradient-based importance scores are used to duplicate more useful experts.
• Router extension: router weights are copied into new slots, and small perturbations are added to the bias to create routing diversity.
• CPT: stochastic gradient diversity and load balancing break the symmetry between duplicates and promote specialization.

Since Top-K routing is kept unchanged, the per-token inference cost does not change. Instead, quality is improved by increasing the total number of experts.

In experiments, an interleaved MoE with 7B→13B total parameters and 1B active was trained on 380B tokens. The model upcycled from 32→64 experts came close to the fixed 64-expert baseline, with 56.4 vs. 56.7 average accuracy across 11 downstream benchmarks and 1.263 vs. 1.267 validation loss.

The cost savings are also significant. GPU hours dropped by about 32%, from 27,888 vs. 41,328, and when a pre-trained checkpoint is already available, only CPT needs to be performed, allowing savings of up to about 67%. Even in larger full MoE architectures with a 256→512 experts, TopK=8 configuration, gap closure of 93~95% was observed across the 154M~1B scale range.

The repository can be run directly based on the NeMo container, either by applying runtime monkey-patching to Megatron-LM/NeMo classes via apply_patches(), or by calling it at the model level via the perform_expert_upcycling() API. After modifying configs/upcycle.yaml, it can be run on a single GPU or multi-GPU via torchrun, and it also includes CPU integration tests and GPU end-to-end tests. The license is CC-BY-NC-4.0, and since this is released for paper reproduction purposes, PRs are not accepted.