AMD-AGI/Instella-MoE
GitHub: AMD-AGI/Instella-MoE
Instella-MoE 是 AMD 推出的完全开源混合专家语言模型,在 AMD Instinct GPU 上完成端到端训练,涵盖从预训练到强化学习的完整流程。
Stars: 50 | Forks: 7
Instella-MoE 是一个最先进的、完全开源的混合专家(MoE)语言模型,拥有 160 亿总参数和 28 亿激活参数,并经历了从预训练到 RL 的端到端训练。
该模型完全从零开始,在 AMD Instinct™ MI300X 和 MI325X GPU 上使用 AMD 的 [Primus](https://github.com/AMD-AGI/Primus) 框架进行训练。Instella-MoE 将稀疏激活的 MoE 设计与 Gated Multi-head Latent Attention (Gated MLA) 和 [FarSkip-Collective](https://github.com/AMD-AGI/FarSkip-Collective) 等架构创新结合在一起。
Figure 1: Pre-trained and Post-trained Instella-MoE model performance compared with other similarly sized state-of-the-art models.
### 快速开始
要在 AMD 硬件上运行 Instella-MoE 的训练,请使用 `rocm/megatron-lm:v25.8_py310` Docker 镜像。可以通过以下命令启动:
```
IMAGE_NAME=rocm/megatron-lm:v25.8_py310
docker run \
-it \
--rm \
--network=host \
--ipc=host \
--privileged \
--device=/dev/kfd \
--device=/dev/dri \
--device=/dev/infiniband \
--group-add video \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--shm-size 200G \
-v $HOME:$HOME \
--name 'instella_moe_train' \
$IMAGE_NAME \
/bin/bash
```
要在 Docker 内部使用模拟数据测试训练,请运行:
```
git clone https://github.com/AMD-AGI/Instella-MoE.git
cd Instella-MoE/training
git submodule update --init --recursive
EXP=examples/megatron/configs_instella_moe/instella_moe-mock_pretrain.yaml \
bash examples/run_instella.sh --task pretrain
```
### 示例用法
```
from transformers import AutoModelForCausalLM, AutoTokenizer
checkpoint = "amd/Instella-MoE-16B-A3B-Think"
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", trust_remote_code=True)
prompt = [{"role": "user", "content": "What are the computational benefits of Mixture-of-Experts models?"}]
inputs = tokenizer.apply_chat_template(
prompt,
add_generation_prompt=True,
return_tensors='pt'
)
tokens = model.generate(
inputs.to(model.device),
max_new_tokens=1024,
temperature=0.6,
top_p=0.95,
do_sample=True
)
print(tokenizer.decode(tokens[0], skip_special_tokens=False))
```
## 训练
下面我们提供了分阶段的描述、训练配置和启动脚本,以重现 Instella-MoE 的训练过程。
从预训练到 DPO 阶段均基于 `training/` 目录下的 [Primus](https://github.com/AMD-AGI/Primus/tree/dev/farskip) 运行,使用统一的 `examples/run_instella.sh` 脚本(从 `training/` 目录执行)。通过 `--task` 标志选择每个阶段,并通过 `EXP` 环境变量传递阶段配置。
强化学习阶段在 `rl/` 下单独运行,它构建于 [`Miles`](https://github.com/radixark/miles) 之上,并复用了我们的 `training/` 和 `inference/` 后端。有关其详细的设置和启动说明,请参见 [`rl/README.md`](rl/README.md)。
以下阶段按顺序遵循完整的训练 pipeline,每个阶段都从上一阶段的 checkpoint 恢复。
### 预训练
预训练阶段从零开始,并在 7.1T 高质量、多样化的训练 token 上进行训练。
```
cd training
EXP=examples/megatron/configs_instella_moe/instella_moe-pretrain.yaml \
bash examples/run_instella.sh --task pretrain
```
### 中期训练
中期训练从中期训练的 checkpoint 恢复,为此我们使用 3 种数据变体进行训练,以生成 3 个不同的 checkpoint(`midtrain_v1`、`midtrain_v2`、`midtrain_v3`)。
```
EXP=examples/megatron/configs_instella_moe/instella_moe-midtrain_v1.yaml \
bash examples/run_instella.sh --task pretrain
```
#### 合并中期训练
在训练完中期训练模型变体后,我们通过合并(权重平均)这些变体的 Hugging Face checkpoint 来生成最终的中期训练 checkpoint。
使用 [merge_hf_ckpts.py](training/examples/scripts/instella_moe_process_checkpoint/merge_hf_ckpts.py):
```
python training/examples/scripts/instella_moe_process_checkpoint/merge_hf_ckpts.py \
--model_class AutoModelForCausalLM \
--out_dir /path/to/midtrain_merged \
--checkpoints /path/to/midtrain_v1 /path/to/midtrain_v2 /path/to/midtrain_v3 \
--weights 1.0 1.0 1.0 \
--dtype float32 \
--save_safetensors \
--trust_remote_code \
--copy_tokenizer
```
我们[在此](training/examples/scripts/instella_moe_process_checkpoint/README.md)分享了有关 checkpoint 处理的更多细节。
### 长上下文训练
在长上下文训练阶段,我们通过在 64k 上下文长度下进行训练,将 Instella-MoE 扩展至长上下文文档。
此阶段从中期训练合并后的 checkpoint 恢复。我们修改了模型配置并扩展了 RoPE theta 参数,以支持更长的上下文。
我们在相同的上下文长度下分两个阶段进行长上下文训练。
第一阶段使用长上下文通用混合数据,教导模型关注整个 64K 窗口内的信息。
第二阶段在数学、代码和重推理任务上进一步提升长上下文性能。
```
# Phase 1
EXP=examples/megatron/configs_instella_moe/instella_moe-long_ctx_phase1.yaml \
bash examples/run_instella.sh --task docmask
# Phase 2
EXP=examples/megatron/configs_instella_moe/instella_moe-long_ctx_phase2.yaml \
bash examples/run_instella.sh --task docmask
```
第二阶段的长上下文 checkpoint 作为我们的“基础”(Base)Instella-MoE checkpoint。
### 监督微调(SFT)
监督微调从长上下文第二阶段的 checkpoint 恢复,并在配置的 `sft_config` 块中定义的 SFT 数据混合上进行训练。
我们还将 SFT 课程划分为两个连续的阶段:阶段 1 在通用 SFT 混合数据上进行训练,阶段 2 在训练的最后阶段通过退火过渡到经过精心挑选的混合数据上。
这个阶段 2 的混合数据是一个由反馈驱动的集合,包含 512K 个样本,专门针对阶段 1 checkpoint 的弱点进行选择,以增强模型能力。
```
# Phase 1:通用 SFT 混合数据
EXP=examples/megatron/configs_instella_moe/instella_moe-sft_phase1.yaml \
bash examples/run_instella.sh --task sft
# Phase 2:反馈驱动的精选混合数据(从 phase 1 自动恢复)
EXP=examples/megatron/configs_instella_moe/instella_moe-sft_phase2.yaml \
bash examples/run_instella.sh --task sft
```
### 直接偏好优化(DPO)
直接偏好优化从最终的 SFT checkpoint 恢复,并分两个阶段运行。
阶段 1 通过纯前向传递计算并缓存参考对数概率;阶段 2 使用参考模型对数概率缓存运行实际的 DPO 训练。
在此阶段,我们禁用了路由器偏置更新和辅助负载均衡损失,因为我们发现这能提高 DPO 训练性能,同时不会损害专家负载均衡。
```
# Stage 1:计算参考 log-probs
EXP=examples/megatron/configs_instella_moe/instella_moe-dpo.yaml \
bash examples/run_instella.sh --task dpo --compute-ref-logprobs
# Stage 2:DPO 训练
EXP=examples/megatron/configs_instella_moe/instella_moe-dpo.yaml \
bash examples/run_instella.sh --task dpo
```
### RL 训练
我们使用最终的 DPO checkpoint 作为 RL 训练的基础策略。
RL 阶段在 `rl/` 中实现,它以 `miles` 为基础进行编排,并复用了我们的 `training/` 和 `inference/` 后端。
RL 环境与 SGLang 推理引擎共享同一个 Docker 镜像,该镜像与 `training/` 和 `inference/` 后端均兼容。
有关详细的设置和启动说明,请参见 [`rl/README.md`](rl/README.md)。
#### 指令遵循 RL (IF-RL)
我们首先在指令遵循任务 (IF-RL) 上运行 RL,以创建一个指令遵循专家。
设置完成后,可以通过以下命令运行 IF-RL 训练:
```
cd rl
bash examples/if_rl/run_rl_instella_if_thinkrl.sh
```
#### 多教师在线策略蒸馏 (MOPD)
指令遵循专家提升了指令遵循能力,但在数学推理等其他领域出现了性能退化。
为了缓解这一问题,我们采用了多教师在线策略蒸馏 ([MOPD](https://arxiv.org/abs/2606.30406)),它将模型同时锚定在使用 RL 训练的指令遵循专家和 DPO 之后生成的初始策略上。
设置完成后,可以通过以下命令运行 MOPD:
```
# 1. 构建 domain-tagged 混合 prompt 集
python examples/mopd/prepare_rl_dolci.py --if-frac 0.5 ... # -> opd_mixed_prompts.jsonl
# 2. 部署两个 teacher(每个对应一个 serving 部署)
# 在 host1 上
MODEL=/path/to/if_rl_teacher_hf bash examples/mopd/serve_teacher.sh
# 在 host2 上
MODEL=/path/to/dpo_teacher_hf bash examples/mopd/serve_teacher.sh
# 3. 启动 MOPD(Ray cluster 必须已经运行)
bash examples/mopd/run_rl_instella_mopd.sh
```
MOPD 训练的输出作为最终的 Instella-MoE checkpoint:Instella-MoE-16B-A3B-Think。
Instella-MoE-16B-A3B-Think 在实现强大的指令遵循性能的同时,保持了出色的数学、编程和推理能力。
## 推理
我们在 `inference/sglang` 下使用 SGLang 引擎实现了大规模推理。我们基于 [AMD-AGI/FarSkip-Collective](https://github.com/AMD-AGI/FarSkip-Collective) 应用了我们的架构创新,并将我们的推理代码作为 SGLang (v0.5.9) 之上的覆盖层呈现。
我们使用覆盖层来仅定位我们需要在 SGLang 中解决的特定更改。要开始使用我们的推理(和 RL),我们使用 `rlsys/miles:rocm7-mi300-sglang0.5.9-te2.10.0-dev-307b5e86` Docker 镜像,可以通过以下命令启动:
```
IMAGE_NAME=rlsys/miles:rocm7-mi300-sglang0.5.9-te2.10.0-dev-307b5e86
docker run \
-it \
--rm \
--network=host \
--ipc=host \
--privileged \
--device=/dev/kfd \
--device=/dev/dri \
--device=/dev/infiniband \
--group-add video \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--shm-size 200G \
-v $HOME:$HOME \
--name 'instella_moe_inference_and_rl' \
$IMAGE_NAME \
/bin/bash
```
然后在 Docker 内部运行我们的 SGLang lm-eval 评估脚本:
```
bash inference/instella_moe_scripts/instella_moe_sglang_lm_eval.sh
```
有关实现和优化的更多细节,请参见 [FarSkip-Collective README](https://github.com/AMD-AGI/FarSkip-Collective/blob/main/README.md)。
## 数据准备
数据准备代码可在 [`training/examples/scripts/instella_moe_data_preparation`](./training/examples/scripts/instella_moe_data_preparation) 中找到。
要准备用于基础模型训练的 Megatron 格式的数据集,请运行:
```
# 包含 *.jsonl 文件的目录(从 HF 下载,例如 instella_gsm8k_synthetic、'allenai/dolma3_dolmino_mix-100B-1125' 等)
JSONL_DIR=/path/to/dataset/data_jsonl
OUTPUT_PREFIX=/path/to/dataset/data_tokenized/data
PRIMUS_PATH=/path/to/Instella-MoE/training
mkdir -p "$(dirname "$OUTPUT_PREFIX")"
python "$PRIMUS_PATH/examples/megatron/preprocess_data.py" \
--input "$JSONL_DIR/*.jsonl" \
--output-prefix "$OUTPUT_PREFIX" \
--tokenizer-type DeepSeekV3Tokenizer \
--tokenizer-model deepseek-ai/DeepSeek-V3 \
--append-eod \
--workers 64 \
--partitions 8
```
有关训练数据混合的详细信息以及不同训练阶段的额外数据准备说明,请参见[训练数据准备](./docs/data_preparation.md)。
## 评估
我们在[评估说明](./docs/evaluation.md)中详细介绍了我们的评估设置。
## 附加资源
### 训练与推理框架
- Instella-MoE 构建于
- [Primus](https://github.com/AMD-AGI/Primus) 训练框架
- [FarSkip-Collective](https://github.com/AMD-AGI/FarSkip-Collective) 训练与推理实现
- 用于 RL 训练的 [Miles](https://github.com/radixark/miles/tree/main)
### Hugging Face 模型集合
我们[在此](https://huggingface.co/collections/amd/instella-moe)发布了我们的模型集合。
## 许可证
- Instella-MoE 模型根据 ResearchRAIL 许可证授权用于学术和研究目的。
- 训练代码库采用 MIT 许可证授权。
- 有关更多信息,请参阅训练代码库的 [LICENSE](./LICENSE) 文件。
## 引用
喜欢我们的工作吗?给我们点个 Star ⭐ —— 这会让我们的 GPU 愉悦地轰鸣!Instella-MoE 的技术报告即将发布,在此之前,欢迎引用:
```
@article{instella,
title={Instella: Fully Open Language Models with Stellar Performance},
author={Liu, Jiang and Wu, Jialian and Yu, Xiaodong and Su, Yusheng and Mishra, Prakamya and Ramesh, Gowtham and Ranjan, Sudhanshu and Manem, Chaitanya and Sun, Ximeng and Wang, Ze and Brahma, Pratik Prabhanjan and Liu, Zicheng and Barsoum, Emad},
journal={arXiv preprint arXiv:2511.10628},
year={2025}
}
@inproceedings{
dukler2026farskipcollective,
title={FarSkip-Collective: Unhobbling Blocking Communication in Mixture of Experts Models},
author={Yonatan Dukler and Guihong Li and Deval Shah and Jiang Liu and Vikram Appia and Emad Barsoum},
booktitle={Ninth Conference on Machine Learning and Systems},
year={2026},
url={https://openreview.net/forum?id=ruOpvLzsGV}
}
```
Figure 1: Pre-trained and Post-trained Instella-MoE model performance compared with other similarly sized state-of-the-art models.
标签:AMD ROCm, DLL 劫持, 人工智能, 凭据扫描, 大语言模型, 模型训练, 混合专家模型, 用户模式Hook绕过, 请求拦截, 逆向工具