大模型训练：保存mindir 格式模型

摘要MindIRMindSpore Intermediate Representation是昇思 MindSpore 原生中间表示模型格式注意 MindIR 不是训练过程保存断点的格式大模型训练阶段输出 ckpt 检查点训练完成后导出 MindIR 用于推理部署。MindIR 同时保存计算图结构、算子定义与权重参数屏蔽不同硬件后端差异实现 “一次训练多端部署”。大模型分布式训练完成后从 ckpt 权重导出 MindIR可交付 MindSpore 原生推理、MindSpore‑Lite也可经由 ATC 工具编译生成 om 离线模型部署到昇腾 NPU。本文讲解 MindIR 基本概念、和 ckpt 的区别、大模型训练后导出 MindIR 完整流程、分布式场景处理、代码实践、约束限制与问题排查帮助开发者打通大模型训练到推理部署链路。1 背景介绍昇思 MindSpore 大模型训练预训练、微调过程中使用CheckPointckpt保存权重断点ckpt 只存储参数值不保存完整网络图结构主要用于断点续训、继续微调不能直接用于跨硬件推理部署。很多开发者容易混淆 ckpt 与 MindIR训练循环迭代保存 ckpt训练收敛完成加载 ckpt 权重到推理网络导出 MindIR。大模型场景还会遇到分布式切分权重导出 MindIR 需要处理并行策略文件。MindIR 作为统一中间表示向上承接训练产出向下对接昇腾推理全栈。实际开发痛点分布式大模型直接导出 MindIR 失败动态 shape 配置错误控制流算子导致导出报错导出后的 MindIR 推理精度与训练结果不一致。掌握 MindIR 导出约束是大模型训推打通的关键。2 MindIR 格式核心概念与和 CKPT 对比2.1 MindIR 核心能力存储完整计算图算子、张量连接关系、网络拓扑嵌入模型权重参数文件自包含硬件无关 IR 描述支持 Ascend/GPU/CPU 多后端输出产物用于推理不可用于恢复训练、继续微调。2.2 CKPT vs MindIR 对比格式存储内容主要用途大模型训练阶段ckpt仅网络权重参数断点续训、继续微调训练过程定期保存输出MindIR计算图 权重跨平台推理部署训练结束后导出得到重要提示训练循环不会直接生成 MindIRMindIR 是训练完成后导出产物不能拿 MindIR 回到训练做 resume 继续训练。3 大模型训练‑MindIR 导出‑部署完整链路大模型分布式训练MindSpore 执行预训练 / 微调迭代输出分布式 ckpt 权重、并行策略 strategy_ckpt构造推理网络加载网络定义设置为 eval 推理模式加载分布式 ckpt 权重结合训练策略文件使用load_distributed_checkpoint把权重灌入推理网络export 导出 MindIR构造示例输入 tensor调用 export 接口生成.mindir文件推理验证MindSpore 直接加载 MindIR 验证输出离线部署ATC 工具将 MindIR 编译为 om 模型用于 Ascend NPU 推理服务。分布式特殊说明大模型 TP/PP 并行训练产出分片 ckpt导出 MindIR 时需要读取训练保存的strategy_ckpt生成推理 layout 策略再加载分片权重否则权重加载失败导出报错。4 代码实践4.1 环境准备 env.shsource /usr/local/Ascend/ascend‑toolkit/set_env.sh pip install mindspore‑ascend2.6.0 export ASCEND_GLOBAL_LOG_LEVEL34.2 大模型训练后导出 MindIR 示例 export_mindir.pyimport numpy as np from mindspore import Tensor, export, context from mindspore import load_distributed_checkpoint, set_auto_parallel_context # 以简易大模型网络举例替换为你的大模型网络定义 from large_model import LargeLLMNet def main(): # 设置静态图模式导出MindIR必须GRAPH_MODE context.set_context(modecontext.GRAPH_MODE, device_targetAscend, device_id0) # 分布式场景加载训练阶段保存的并行策略文件 set_auto_parallel_context( strategy_ckpt_load_file./output/train_strategy.ckpt ) # 1.初始化推理网络eval模式关闭dropout等训练组件 net LargeLLMNet(vocab_size32000, hidden_size512) net.set_train(False) # 构造推理输入样例shape需要匹配模型输入 input_ids Tensor(np.ones((1,128), dtypenp.int32)) # 2.生成推理layout策略用于分布式权重加载 predict_layout net.infer_predict_layout(input_ids) # 分布式分片ckpt文件列表训练输出的多个分片权重 ckpt_file_list [ ./output/ckpt/rank_0‑llm.ckpt, ./output/ckpt/rank_1‑llm.ckpt ] # 加载分布式训练ckpt权重灌入网络 load_distributed_checkpoint(net, ckpt_file_list, predict_layout) # 3.导出MindIRfile_format指定MINDIR export( net, input_ids, file_name./llm_model, file_formatMINDIR ) print(MindIR导出完成输出文件 llm_model.mindir) if __name__ __main__: main()4.3 MindIR 加载推理验证 infer_mindir.pyfrom mindspore import context, Model import numpy as np from mindspore import Tensor context.set_context(modecontext.GRAPH_MODE, device_targetAscend) def run_mindir_infer(): model Model(./llm_model.mindir) input_data Tensor(np.ones((1,128), dtypenp.int32)) output model.predict(input_data) print(MindIR推理输出shape, output[0].shape) if __name__ __main__: run_mindir_infer()4.4 ATC 工具MindIR 编译昇腾 om 离线模型source env.sh atc --model./llm_model.mindir \ --output./llm_model.om \ --soc_versionAscend910B \ --input_shapeinput_ids:1,1285 FAQ 常见问题Q1MindIR 可以用来恢复大模型继续训练吗A不可以。MindIR 是推理 IR不包含优化器状态、训练循环信息断点续训只能使用 ckpt 文件。Q2大模型分布式训练导出 MindIR 报权重加载失败A确认加载训练输出的train_strategy.ckpt调用infer_predict_layout生成推理 layoutckpt_file_list 路径与 rank 分片文件一一对应。Q3export 导出 MindIR 报错控制流不支持AMindIR 对部分动态控制流支持有限需要简化网络推理分支把训练阶段独有的控制逻辑在 eval 模式去除。Q4MindIR 推理结果和 ckpt 直接推理结果不一致A检查网络是否没有切换 set_train (False)确认 dropout 等随机模块是否关闭核对 export 输入样例 dtype 与维度。Q5MindIR 和 ONNX 怎么选择A昇思全栈优先 MindIR原生适配 MindSpore、ATC、MindSpore‑Lite跨其他框架交互场景选择 ONNX。6 总结昇思 MindSpore 大模型训练阶段输出 ckpt 检查点用于断点续训、微调MindIR 不是训练产物是训练收敛后导出的推理中间表示同时保存网络图结构与权重屏蔽硬件差异。大模型 TP/PP 分布式场景导出 MindIR需要带上训练并行策略文件加载分片 ckpt 权重到 eval 推理网络调用 export 接口生成 mindir 文件得到 MindIR 后既可以在 MindSpore 直接推理验证也可以经由 ATC 编译得到 om 离线模型部署昇腾 NPU。