MindSpeed-LLM 昇腾大语言模型分布式训练框架:架构解析与实战指南
MindSpeed-LLM 昇腾大语言模型分布式训练框架:架构解析与实战指南
MindSpeed-LLM 是面向华为昇腾硬件设计的大语言模型(LLM)端到端分布式训练框架,覆盖分布式预训练、分布式指令微调以及权重转换、在线推理、基线评估等完整开发工具链。本文以官方 Introduction 文档 为核心骨架,结合仓库源码与真实示例脚本,系统讲解其四层架构、双训练后端(Megatron 与 FSDP2)、多维并行策略、模型支持范围与核心特性,帮助读者快速理解框架全貌并上手训练任务。
MindSpeed-LLM 是什么
MindSpeed-LLM 是 MindSpeed 生态中的大模型分布式训练框架,旨在为华为昇腾芯片生态合作伙伴提供端到端的大语言模型训练方案。它原生集成了 MindSpeed Core 训练加速库,并围绕昇腾硬件从并行优化、内存优化、通信优化、计算优化四个维度对 LLM 训练进行系统加速。
框架基于 Transformer 支持 Dense(稠密)与 MoE(混合专家)两大类模型,内置 100 余个主流开源模型及开箱即用的训练脚本,覆盖模型预训练、微调与高效分布式训练全流程。硬件侧,模型支持清单 与 安装指南 明确说明仓库支持 Ascend 950 系列、Atlas A3 训练系列与 Atlas A2 训练系列产品,且单张 NPU 要求至少 64 GB 显存(可使用 npu-smi info 查看)。
MindSpeed-LLM 四层架构
MindSpeed-LLM 架构由四个层次组成,如图所示。
第一层:MindSpeed-LLM 功能模块
框架提供完整的功能模块集合,主要包括:
- Megatron 与 Hugging Face 权重互转:支撑 Hugging Face 权重在 Megatron-core 格式下的加载与转换,涉及 TP/PP/DPP/VPP/CP/EP 等多种切分维度;
- 分布式评估与推理:支持模型在分布式环境下的在线推理与基于公开数据集的基线评估;
- 性能采集与确定性计算:提供基于昇腾芯片的 profiling 数据采集工具,以及保证结果可复现的确定性计算能力;
- 断点续训与权重保存:基于 checkpoint 实现训练中断后的恢复,并在训练结束时保存最终权重。
第二层:MindSpeed-LLM 训练后端
框架提供两个并行训练后端,二者可独立选型:
- Megatron 训练后端:基于 Megatron-LM 构建,提供以 ModelSpec 为模板的增量式模型开发方案,开发一个新模型通常只需数周甚至数天;
- FSDP2 训练后端:基于 MindSpeed-FSDP 构建,直接对接第三方 Transformers 库,提供增量式模型开发方案,开发一个新模型通常只需数天。
从源码结构看,mindspeed_llm/ 目录下的 core(并行策略、模型定义、Transformer 组件、数据集处理、高可用等底层能力)、features_manager(通过 patch 机制向训练流程注入 Flash Attention、MoE 路由、量化感知训练等特性)、training(参数解析、分布式初始化、checkpoint 管理与训练循环控制)、tasks(评估、推理、微调、DPO、权重转换等任务执行层)共同支撑 Megatron 后端,而 fsdp2 则作为独立后端并行存在,包含模型定义、数据处理、分布式策略与推理引擎的完整训练流水线。
第三层:MindSpeed-LLM 训练工作流
- 分布式预训练:支持端到端分布式预训练,包含数据预处理方案与主流模型、数据切分方案;
- 分布式指令微调:支持业界主流微调算法(全参微调、LoRA 等),并提供微调性能与内存优化手段。
第四层:MindSpeed-LLM 模型支持
- 100 余个预置开源模型:覆盖主流 Dense、MoE 与 State Space Model(SSM)系列,提供开箱即用的高性能训练脚本;
- 兼容多种主流 LLM 架构:支持基于 Transformer 与 SSM 架构的大模型。
核心特性总览
- 主流 LLM 支持:支持 100 余个主流 LLM 模型,包括 Qwen3、DeepSeek、Mamba2 等系列,覆盖 Dense、MoE、SSM 架构,并提供面向昇腾架构定制的高性能训练脚本。
- 分布式预训练:提供端到端分布式预训练,包含数据预处理方案与多维并行策略——张量并行(TP)、流水线并行(PP)、数据并行(DP)、上下文并行(CP)、专家并行(EP)。
- 分布式指令微调:支持全参微调、LoRA 等业界主流微调算法,并提供微调性能与内存优化方法。
- 模型权重转换:支持 Megatron 与 Hugging Face 格式互转,以及 LoRA 微调权重的独立与合并转换。
- 在线推理与评估:支持模型分布式在线推理,以及基于公开数据集与标准基准数据集的在线评估。
双训练后端深度解析
Megatron 后端:ModelSpec 增量式模型开发
Megatron 后端以 ModelSpec 为模板,开发者只需定义模型结构规范即可快速接入新模型。以 Qwen3-8B 预训练脚本 pretrain_qwen3_8b_4K_ptd.sh 为例,其核心参数可分为以下几组:
# 模型结构(ModelSpec 驱动)
--use-mcore-models \
--spec mindspeed_llm.tasks.models.spec.qwen3_spec layer_spec \
--num-layers 36 --hidden-size 4096 --ffn-hidden-size 12288 \
--num-attention-heads 32 --num-query-groups 8 --group-query-attention \
--padded-vocab-size 151936 --make-vocab-size-divisible-by 1 \
--position-embedding-type rope --rotary-base 1000000 \
--normalization RMSNorm --swiglu --qk-layernorm
# 并行配置
--tensor-model-parallel-size 1 --pipeline-model-parallel-size 2
# 加速与优化特性
--use-flash-attn --use-fused-rotary-pos-emb --use-fused-swiglu \
--use-fused-rmsnorm --use-distributed-optimizer --reuse-fp32-param \
--overlap-grad-reduce --overlap-param-gather --use-ascend-coc
# 训练超参
--micro-batch-size 1 --global-batch-size 64 --seq-length 4096 \
--lr 1.25e-6 --lr-decay-style cosine --weight-decay 1e-1 \
--bf16 --train-iters 2000
其中 --spec 参数指向仓库内的 Qwen3 模型规范定义(mindspeed_llm/tasks/models/spec/qwen3_spec),--use-ascend-coc、--overlap-grad-reduce 等开关对应 features/README.md 中列出的通信与计算重叠优化特性。
FSDP2 后端:直接对接 Transformers 生态
FSDP2 后端基于 PyTorch 官方推荐的 Fully Sharded Data Parallel(FSDP2)实现,可直接复用第三方 Transformers 库的模型定义,显著降低新模型接入成本。仓库在 examples/fsdp2 下提供了 Qwen3、Qwen3-MoE 等模型的 YAML 配置文件与 .sh 启动脚本,完整参数参考 FSDP2 全量参数说明,性能优化方法见 FSDP2 后端性能优化指南。
训练工作流实战
分布式预训练:数据预处理与训练拉起
预训练流程分为数据预处理与模型训练两步。数据预处理使用仓库根目录的 preprocess_data.py,以 Qwen3 预训练数据处理脚本 data_convert_qwen3_pretrain.sh 为例:
python ./preprocess_data.py \
--input ./dataset/train-00000-of-00042-xxx.parquet \
--tokenizer-name-or-path ./model_from_hf/qwen3_hf/ \
--tokenizer-type PretrainedFromHF \
--handler-name GeneralPretrainHandler \
--output-prefix ./dataset/enwiki \
--json-keys text \
--workers 4 --log-interval 1000
--handler-name 指定数据处理 Handler(预训练使用 GeneralPretrainHandler,指令微调使用 AlpacaStyleInstructionHandler 等,具体见 mindspeed_llm/tasks/preprocess)。随后即可通过 torchrun 拉起 pretrain_gpt.py 开始训练,预训练方案(多样本数据集预训练、多样本 Pack 模式预训练)详见 pretrain 文档。
分布式指令微调:全参微调与 LoRA
指令微调通过仓库根目录的 posttrain_gpt.py 拉起,以 Qwen3-8B 全参微调脚本 tune_qwen3_8b_4K_full_ptd.sh 为例,其微调专属参数包括:
--finetune \
--stage sft \
--is-instruction-dataset \
--prompt-type qwen3 \
--no-pad-to-seq-lengths \
--handler-name AlpacaStyleInstructionHandler \
--enable-thinking true
框架支持的微调方案包括单样本微调、多样本 Pack 微调、多轮对话微调、LoRA、LU-LoRA、离线 DPO、层间分离训练(LDT)等,完整列表见 分布式微调方案与特性。
权重转换:HF 与 Megatron 互转
权重转换使用仓库根目录的 convert_ckpt_v2.py,以 Qwen3 的 HF 转 MCore 脚本 ckpt_convert_qwen3_hf2mcore.sh 为例:
python convert_ckpt_v2.py \
--load-model-type hf \
--save-model-type mg \
--target-tensor-parallel-size 1 \
--target-pipeline-parallel-size 8 \
--load-dir ./model_from_hf/qwen3_hf/ \
--save-dir ./model_weights/qwen3_mcore/ \
--model-type-hf qwen3
转换支持 TP、PP、DPP、VPP、CP、EP、loop layer 等多种切分维度,并支持 LoRA 权重合并转换,详细参数说明见 权重转换文档。
在线推理与基线评估
在线推理使用 inference.py,评估使用 evaluation.py。以 Qwen3-8B 为例,推理脚本 generate_qwen3_8b_ptd.sh 通过 torchrun 拉起推理并支持 --max-new-tokens 256 控制生成长度;评估脚本 evaluate_qwen3_8b_ptd.sh 通过 --task mmlu 指定评测任务(支持 MMLU、CEval、BoolQ、BBH、AGIEval、HumanEval、CMMLU、GSM8k、HellaSwag、NeedleBench 等标准数据集),--no-chat-template 控制是否使用对话模板,详见 评估指南。
源码结构佐证
从 项目导读 可以看出,mindspeed_llm/ 核心代码目录包含五个关键子模块,协作关系为 core/ → features_manager/ → training/ → tasks/,fsdp2/ 作为独立后端并行存在:
core/:底层核心能力,含并行策略(CP/TP/PP)、模型定义、Transformer 组件、数据集处理与高可用;features_manager/:特性注册与管理,通过 patch 机制按需组合激活 Flash Attention、MoE 路由、QAT 等特性;training/:训练框架初始化与主循环,负责训练生命周期管理;tasks/:训练任务入口与服务逻辑,是训练过程的执行层;fsdp2/:独立 FSDP2 后端完整实现。
仓库根目录还提供 preprocess_data.py、pretrain_gpt.py、posttrain_gpt.py、convert_ckpt_v2.py、evaluation.py、inference.py、train_fsdp2.py 等工具与训练流程脚本,覆盖从数据处理到训练、评估、在线推理的完整链路。
快速上手路径
- 环境安装:参考 安装指南,支持镜像安装(内置 CANN 与 TorchNPU)与源码安装两种方式;
- 选择后端与模型:Megatron 后端参考 快速入门(MCore),FSDP2 后端参考 快速入门(FSDP2);模型清单见 支持模型列表;
- 运行示例脚本:按需选择 examples/mcore 或 examples/fsdp2 下的数据处理、权重转换、预训练、微调、推理、评估脚本;
- 查看特性与调优:性能/显存优化特性清单见 训练方案与特性,性能采集工具见 profiling 文档,版本配套与更新说明见 版本说明。
总结
MindSpeed-LLM 以四层架构为骨架,通过 Megatron 与 FSDP2 双训练后端、多维并行策略、100 余个预置模型及完整的工具链,为昇腾硬件提供了从数据预处理到预训练、微调、推理、评估的端到端 LLM 训练解决方案。开发者既可直接复用开箱即用的示例脚本快速拉起训练,也可基于 ModelSpec 模板或 FSDP2 的 Transformers 生态对接能力在数天至数周内完成新模型的接入与训练。
