エキスパートアップサイクル（GitHubリポジトリ）
Amazon の研究チームが提案した「Expert Upcycling」手法は、既存の MoE モデルをトレーニング中に専門家の数を増やすことで、計算コストを最大 67% 削減しつつ性能を維持する画期的な技術である。
キーポイント
中盤での専門家拡張メカニズム
「Expert replication（専門家の複製）」、「Router extension（ルーターの拡張）」、「Continued pre-training（継続的事前学習）」の 3 つのステップにより、トレーニング途中で MoE のパラメータ数を増やす手法を確立した。
劇的な計算コストの削減
7B から 13B モデルへの拡張実験において、GPU 使用時間を約 32% 削減し、事前学習済みチェックポイントからの継続学習では最大 67% のコスト削減を実現した。
性能とスケーラビリティの維持
拡張後のモデルは固定サイズのベースラインと同等の精度（11 のベンチマークで平均誤差率 56.4% vs 56.7%）を達成し、154M から 1B パラメータ規模まで効果が発揮されることを確認した。
推論コストの不変性
Top-K ルーティング方式は変更せず、トークンあたりの推論コストを増加させることなく、モデルの容量と性能のみを向上させた。
トレーニング中の移行ポイントでの呼び出し
.upcycle_experts() と .upcycle_router() は、学習プロセスの所望の転換点で呼び出す必要があります。
モデルラップからのアンワップが必要
モデルが通常ラップされているため、デコーダにアクセスするには明示的にアンワップ処理を行う必要があります。
柔軟なアップサイクル戦略
Expert duplication と Router expansion の両方で、勾配ベースの重要度スコアに基づく推奨戦略や、ノイズ追加、SVD 分解など多様な初期化手法が用意されています。
影響分析・編集コメントを表示
影響分析
この手法は、大規模言語モデルのトレーニングにおける計算資源のボトルネックである「メモリ使用量」と「通信コスト」を解決する実用的な解決策を提供します。特に、既存の高性能モデルをより大規模化する必要があるケースにおいて、膨大な計算リソースと時間を節約しながら性能を最大化できるため、業界全体の MoE 採用ハードルを大幅に下げる可能性があります。
編集コメント
MoE モデルのトレーニング効率化において、既存モデルを「アップサイクル」する発想は非常に革新的です。リソース制約のある環境でも大規模モデルの開発が可能になるため、実務現場での導入加速が期待されます。
継続事前学習におけるMixture-of-Experts（MoE）モデルのCapacity Expansion（容量拡張）。
Dwivedi 他、「Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts」（プレプリント）。
スケーリング法則によれば、MoEの品質は、アクティブな計算量を固定した状態で総エキスパート数が増加すると予測可能に向上しますが、大規模なMoEをゼロから学習するのは高コストです。メモリ、勾配、すべての-to-all通信はすべて総パラメータ数に比例してスケールします。Expert Upcycling（エキスパートのアップサイクル）は、この問題を回避するために、学習をより小さなEエキスパートモデルで開始し、アップサイクル演算子を通じて中間の学習段階でmEエキスパートに拡張します。
- エキスパート複製：各エキスパートが複製されます（高有用性のエキスパートには、勾配に基づく重要度スコアに応じてより多くのコピーが割り当てられます）。
- ルーター拡張：ルーターの重みが新しいスロットにコピーされ、ルーティングの多様性を seeded するために小さなバイアス摂動が加えられます。
- 継続事前学習（CPT）：確率的勾配の多様性と損失のない負荷平衡により、複製品間の対称性が打破され、専門化が促進されます。
Top-Kルーティングは全体を通して固定されたままなので、トークンあたりの推論コストは変更されません。
*図1：エキスパートアップサイクル手順の概要。*
380Bトークンで事前学習された、7B→13Bの総パラメータ（アクティブ1B）インターリーブMoEにおける主要な結果：
- アップサイクルされたモデル（32→64エキスパート）は、11のダウンストリームベンチマーク（平均精度56.4対56.7）および検証損失（1.263対1.267）において、固定サイズ64エキスパートのベースラインと同等のパフォーマンスを示します。
- GPU時間におけるトレーニングコストは約32%削減されます（41,328時間から27,888時間へ）。事前学習済みチェックポイントが既に存在する場合（例えば、以前のトレーニング実行や公開リリースから取得したものであれば）、事前学習のコストは既に支払われており、CPTフェーズのみが必要となるため、節約率は約67%に達します。
- 結果は、154Mから1Bの総パラメータ数にわたる規模において93〜95%のギャップ解消を示す、フルMoEアーキテクチャ（256→512エクスパート、TopK=8）にも一般化します。
*図2: 7B→13BアップサイクルモデルとベースラインのGPU時間、検証損失、および下流タスクでの精度比較。*
インストール
推奨環境: NeMo 2.x コンテナ
公式のNeMoコンテナから開始してください。PyTorch、Megatron-LM、Transformer Engine、NeMo、Lightning、omegaconfがすべて事前にインストールされています。
docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
-v /path/to/expert-upcycling:/workspace/expert-upcycling \
-it nvcr.io/nvidia/nemo:24.09 bash
コンテナ内:
cd /workspace/expert-upcycling
pip install -e .
pip install dacite
注意: コンテナ内で pip install -e ".[nemo]" を使用しないでください。これにより、コンテナに事前にインストールされているNeMoと競合します。
新規構築（NeMoコンテナ不使用）
依存関係を手動でインストールし、関連するエクストラオプションを指定してパッケージをインストールします。
コアのみ（torch + numpy）:
pip install -e .
pip install dacite
Megatron-LM統合あり:
pip install -e ".[megatron]"
完全なNeMoエントリポイント（NeMo、Lightning、omegaconfをインストール）:
pip install -e ".[nemo]"
クイックスタート
オプション A: NeMo エントリポイント（推奨）
configs/upcycle.yaml を編集してモデルの次元を設定し、リポジトリのルートから実行してください：
単一 GPU
/workspace/expert-upcycling に移動
python -m expert_upcycling.entrypoint \
--config-path=configs --config-name=upcycle \
resume.restore_config.path=/path/to/base/checkpoint
複数 GPU（例：テンソル並列処理を用いた 8 個の GPU）
torchrun --nproc_per_node=8 -m expert_upcycling.entrypoint \
--config-path=configs --config-name=upcycle \
resume.restore_config.path=/path/to/base/checkpoint \
strategy.tensor_model_parallel_size=8
このコールバックは最初のオプティマイザステップで発火し、エキスパートの数を倍増させ、アップサイクルされたチェックポイントを保存して終了します。出力パスのデフォルトは <input_checkpoint>-upcycled です。
オプション B: 既存のトレーニングスクリプトにパッチを当てる
import expert_upcycling
expert_upcycling.apply_patches()
現在、TEGroupedMLP には .upcycle_experts() が、TopKRouter には .upcycle_router() が存在します
トレーニング中の希望する遷移ポイントでこれらを呼び出してください。
注意：モデルは通常ラップされているため、デコーダにアクセスするにはアンラップする必要があります：
inner = model
for attr in ("module", "module"):
if hasattr(inner, attr):
inner = getattr(inner, attr)
for i, layer in enumerate(inner.decoder.layers):
if hasattr(layer.mlp, 'experts'):
selected = layer.mlp.experts.upcycle_experts(optimizer, i, expert_cfg)
if hasattr(layer.mlp, 'router'):
layer.mlp.router.upcycle_router(router_cfg, selected)
オプション C: モデルレベルの API を使用する
from expert_upcycling import perform_expert_upcycling
perform_expert_upcycling(
model, optimizer,
expert_cfg={"usefulness_metric": "gradient_norm", "selection_strategy": "greedy"},
router_cfg={"method": "bias_only", "bias_noise_scale": 0.01},
)
Upcycling Strategies
Expert duplication
Strategy
Description
Utility-based (recommended)
Duplicate high-importance experts using gradient-based scores (weight norm, saliency, gradient squared, approx Fisher)
copy
Exact duplication (baseline)
copy_noise
Duplication + Gaussian noise
drop_upcycle
Re-initialize a fraction of columns
svd_perturb
SVD decomposition + perturbation
+ 6 more
See expert_upcycling.config.UpcycleMethod
Router expansion
Strategy
Description
bias_only (recommended)
Keep weights identical, add noise to bias
copy
Exact duplication
copy_noise
Duplication + noise
+ 7 more
See expert_upcycling.config.RouterUpcycleMethod
Architecture
This package treats Megatron-LM and NeMo as third-party dependencies — no fork required. Upcycling methods are injected at runtime via monkey-patching:
expert-upcycling/ # pip install -e .
├── expert_upcycling/
│ ├── config.py # 全列挙型 + データクラス（依存関係なし）
│ ├── expert_upcycler.py # ヒューリスティック戦略（torch のみ）
│ ├── expert_selector.py # ユーティリティベースの選択（torch + numpy）
│ ├── router_upcycler.py # ルーター戦略（torch のみ）
│ ├── optimizer_utils.py # オプティマイザ状態の処理（torch のみ）
│ ├── patch.py # Megatron-LM クラスへのモンキーパッチ
│ ├── upcycle_model.py # モデルの走査
│ └── entrypoint.py # NeMo 起動スクリプト
├── configs/
│ └── upcycle.yaml # 例：設定ファイル
└── scripts/
└── run_upcycle.sh # 例：起動スクリプト
テストの実行
CPU テスト（GPU なし、Megatron インストール不要）
python tests/test_comprehensive.py # 91 テスト：全メソッド、全戦略
pytest tests/test_integration.py -v # 7 件のエンドツーエンド統合テスト
GPU テスト（NeMo コンテナ + GPU が必要）
python tests/test_entrypoint_gpu.py # 実際の TEGroupedMLP + TopKRouter、32->64 エキスパート
引用
@article{dwivedi2025expertupcycling,
title={Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts},
author={Dwivedi, Chaitanya and Gupta, Himanshu and Varshney, Neeraj and Jayarao, Pratik and Yin, Bing and Chilimbi, Trishul and Huang, Binxuan},
year={2026}
}
ライセンス
CC-BY-NC-4.0
このコードは、関連する論文で記述されている手法と知見を支持するため、学術的および科学的な再現性の目的のみに公開されています。論文で使用されたコードと完全に一致させるため、プルリクエストは受け付けておりません。
原文を表示
Capacity expansion for Mixture-of-Experts models during continued pre-training.
Dwivedi et al., "Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts" (preprint).
Scaling laws show that MoE quality improves predictably with total expert count at fixed active computation, but training large MoEs from scratch is expensive — memory, gradients, and all-to-all communication all scale with total parameters. Expert upcycling sidesteps this by starting training with a smaller E-expert model and expanding to mE experts mid-training via the upcycling operator:
- Expert replication — each expert is duplicated (high-utility experts receive more copies via gradient-based importance scores).
- Router extension — router weights are copied to new slots with small bias perturbations to seed routing diversity.
- Continued pre-training (CPT) — stochastic gradient diversity and loss-free load balancing break symmetry among duplicates, driving specialization.
Top-K routing is held fixed throughout, so per-token inference cost is unchanged.
*Figure 1: Overview of the expert upcycling procedure.*
Key results on a 7B→13B total parameter (1B active) interleaved MoE, pre-trained on 380B tokens:
- The upcycled model (32→64 experts) matches the fixed-size 64-expert baseline across 11 downstream benchmarks (56.4 vs. 56.7 avg accuracy) and validation loss (1.263 vs. 1.267).
- Training cost is reduced by ~32% of GPU hours (27,888 vs. 41,328 hours). When a pre-trained checkpoint already exists (e.g., from a prior training run or a public release), the pre-training cost is already paid and only the CPT phase is needed, bringing savings to ~67%.
- Results generalize to full MoE architectures (256→512 experts, TopK=8) with 93–95% gap closure across scales from 154M to 1B total parameters.
*Figure 2: GPU hours, validation loss, and downstream accuracy for the 7B→13B upcycled model vs. baselines.*
Installation
Recommended: NeMo 2.x container
Start from the official NeMo container — PyTorch, Megatron-LM, Transformer Engine, NeMo, Lightning, and omegaconf are all pre-installed.
docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
-v /path/to/expert-upcycling:/workspace/expert-upcycling \
-it nvcr.io/nvidia/nemo:24.09 bash
# Inside the container:
cd /workspace/expert-upcycling
pip install -e .
pip install daciteDo not use pip install -e ".[nemo]" inside the container — it would conflict with the container's pre-installed NeMo.
From scratch (no NeMo container)
Install dependencies manually, then install the package with the relevant extras:
# Core only (torch + numpy):
pip install -e .
pip install dacite
# With Megatron-LM integration:
pip install -e ".[megatron]"
# Full NeMo entrypoint (installs NeMo, Lightning, omegaconf):
pip install -e ".[nemo]"Quick Start
Option A: NeMo entrypoint (recommended)
Edit configs/upcycle.yaml to set your model dimensions, then run from the repo root:
# Single GPU
cd /workspace/expert-upcycling
python -m expert_upcycling.entrypoint \
--config-path=configs --config-name=upcycle \
resume.restore_config.path=/path/to/base/checkpoint
# Multi-GPU (e.g. 8 GPUs with tensor parallelism)
torchrun --nproc_per_node=8 -m expert_upcycling.entrypoint \
--config-path=configs --config-name=upcycle \
resume.restore_config.path=/path/to/base/checkpoint \
strategy.tensor_model_parallel_size=8The callback fires on the first optimizer step, doubles the expert count, saves the upcycled checkpoint, and exits. The output path defaults to <input_checkpoint>-upcycled.
Option B: Patch existing training script
import expert_upcycling
expert_upcycling.apply_patches()
# Now TEGroupedMLP has .upcycle_experts() and TopKRouter has .upcycle_router()
# Call them during training at the desired transition point.
# Note: model is typically wrapped — unwrap to reach the decoder:
inner = model
for attr in ("module", "module"):
if hasattr(inner, attr):
inner = getattr(inner, attr)
for i, layer in enumerate(inner.decoder.layers):
if hasattr(layer.mlp, 'experts'):
selected = layer.mlp.experts.upcycle_experts(optimizer, i, expert_cfg)
if hasattr(layer.mlp, 'router'):
layer.mlp.router.upcycle_router(router_cfg, selected)Option C: Use the model-level API
from expert_upcycling import perform_expert_upcycling
perform_expert_upcycling(
model, optimizer,
expert_cfg={"usefulness_metric": "gradient_norm", "selection_strategy": "greedy"},
router_cfg={"method": "bias_only", "bias_noise_scale": 0.01},
)Upcycling Strategies
Expert duplication
Strategy
Description
Utility-based (recommended)
Duplicate high-importance experts using gradient-based scores (weight norm, saliency, gradient squared, approx Fisher)
copy
Exact duplication (baseline)
copy_noise
Duplication + Gaussian noise
drop_upcycle
Re-initialize a fraction of columns
svd_perturb
SVD decomposition + perturbation
+ 6 more
See expert_upcycling.config.UpcycleMethod
Router expansion
Strategy
Description
bias_only (recommended)
Keep weights identical, add noise to bias
copy
Exact duplication
copy_noise
Duplication + noise
+ 7 more
See expert_upcycling.config.RouterUpcycleMethod
Architecture
This package treats Megatron-LM and NeMo as third-party dependencies — no fork required. Upcycling methods are injected at runtime via monkey-patching:
expert-upcycling/ # pip install -e .
├── expert_upcycling/
│ ├── config.py # All enums + dataclasses (no deps)
│ ├── expert_upcycler.py # Heuristic strategies (torch only)
│ ├── expert_selector.py # Utility-based selection (torch + numpy)
│ ├── router_upcycler.py # Router strategies (torch only)
│ ├── optimizer_utils.py # Optimizer state handling (torch only)
│ ├── patch.py # Monkey-patches onto Megatron-LM classes
│ ├── upcycle_model.py # Model traversal
│ └── entrypoint.py # NeMo launch script
├── configs/
│ └── upcycle.yaml # Example config
└── scripts/
└── run_upcycle.sh # Example launch script
Running Tests
# CPU tests (no GPU, no Megatron install required)
python tests/test_comprehensive.py # 91 tests: all methods, all strategies
pytest tests/test_integration.py -v # 7 end-to-end integration tests
# GPU test (requires NeMo container + GPU)
python tests/test_entrypoint_gpu.py # real TEGroupedMLP + TopKRouter, 32->64 expertsCitation
@article{dwivedi2025expertupcycling,
title={Expert Upcycling: Shifting the Compute-Efficient Frontier of Mixture-of-Experts},
author={Dwivedi, Chaitanya and Gupta, Himanshu and Varshney, Neeraj and Jayarao, Pratik and Yin, Bing and Chilimbi, Trishul and Huang, Binxuan},
year={2026}
}License
CC-BY-NC-4.0
This code is being released solely for academic and scientific reproducibility purposes, in support of the methods and findings described in the associated publication. Pull requests are not being accepted in order to maintain the code exactly as it was used in the paper.
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み