先交代一个背景这两年大模型相关的工具链、框架、教程层出不穷但大部分资料都是围绕 PyTorch 生态展开的。昇思 MindSpore 作为另一条完整的技术栈在国产算力适配、自动并行、动静统一等方面做了很多差异化的设计尤其适合有昇腾硬件、或者希望绕开 CUDA 依赖的团队。我大概从去年开始把一部分预训练任务从 PyTorch 迁移到昇思上跑中间踩了不少坑也总结出了一套相对通用的操作路径。这篇文章不吹不黑纯粹从实操角度把昇思上做大模型预训练的通用流程、关键设计、常见问题梳理一遍给想入坑或者正在迁移的同行做一个参考。1. 整体流程设计与方案选型1.1 昇思预训练到底在做什么预训练这件事本质上就是在大规模无标注文本或者图文、语音上通过自监督目标函数让模型学会语言的统计规律和通用表示。昇思上的预训练流程和你在 PyTorch 里做的事情没有本质区别但因为底层实现不同整个流程的组织方式有明显差异。一个标准的昇思预训练任务至少包含以下几个环节数据集准备与流式加载、模型结构定义与初始化、并行策略配置如果你有多卡或集群、优化器与学习率调度、训练循环前向、反向、梯度累积、参数更新、checkpoint 保存与恢复、日志与指标监控、最后是评估和推理验证。我在 PyTorch 时代习惯于把数据集全部加载到内存再 shuffle到了昇思这里很快发现行不通。因为昇思的数据处理管线是基于GeneratorDatasetmapbatch的流式架构数据不会一次性全部进内存而是按需生成和变换。这个设计对超大语料几十 TB 甚至上百 TB非常友好但需要你调整自己的编程习惯。1.2 为什么选择昇思而不是 PyTorch这是每个迁移者都会问的问题。我当时的出发点很朴素团队里有昇腾卡而昇思在昇腾上的算子适配和显存管理确实做得更细。PyTorch 虽然可以通过各种方式在昇腾上跑但总有一些算子要自己适配分布式通信的效率也差一些。从框架设计的角度看昇思有几个对预训练特别友好的特性一是动静统一。你可以先用动态图模式快速调试模型结构和训练逻辑等确认无误后用set_context(modeGRAPH_MODE)切到静态图模式做高性能训练。PyTorch 里torch.compile也能做到类似效果但昇思的动静切换在工程上更透明调试期和生产期共用同一套代码。二是自动并行。昇思的AutoParallel可以根据模型大小、数据量、集群拓扑自动搜索最优的并行策略包括数据并行、算子级模型并行、流水线并行。你不需要像 Megatron-LM 那样手动切分 Transformer 层框架会在计算图编译阶段帮你完成切分和通信插入。三是统一的 Checkpoint 管理。昇思把模型参数、优化器状态、RNG 状态统一保存恢复训练时能精确回到中断点这对于动不动训练几周的大模型来说非常关键。当然PyTorch 生态更丰富社区资源多HuggingFace 上的模型基本能直接加载。昇思也有 MindSpore Transformers 仓库和 ModelZoo很多主流模型结构都能直接拿到预训练权重或直接跑通训练脚本。如果你主要工作在昇腾硬件上昇思的综合体验是优于 PyTorch 的。1.3 预训练方法选择的三个层次我把昇思上做预训练的方法分为三个层次你可以根据自己的硬件资源、模型规模和技术储备来选择。第一层使用现成脚本和模型仓库。昇思提供了mindspore/models仓库里面有 GPT、BERT、Llama 等模型的完整预训练脚本包括数据处理、模型定义、训练循环、并行策略。你的工作是准备数据、调整超参数、启动训练。这个层次适合第一次接触昇思、或者只想快速验证硬件的团队。第二层在现有模型结构上做定制。比如你要改注意力机制、换位置编码、或者调整 MoE 结构那么可以在 MindSpore Transformers 的模型代码基础上修改训练脚本复用现成的Trainer接口。这个层次需要你对模型实现和昇思 API 有一定理解但不需要从零写分布式逻辑。第三层完全自定义模型和训练流程。你要自己定义nn.Cell自己写TrainOneStepCell自己配置并行策略。这个层次灵活性最大但工作量也最大适合做前沿算法研究的团队。我个人建议无论你处于哪个层次先把官方仓库的标准训练脚本跑通一遍再动手改自己的东西。这样你能确认环境、数据格式、训练流程都已经正确后续排错的范围会小很多。2. 环境准备与数据管线搭建2.1 安装与版本匹配昇思的版本和硬件驱动、固件、Python 版本有严格对应关系这是新手最容易卡住的第一关。我建议先确定你的硬件型号和昇腾驱动版本再根据官方版本匹配表选择 MindSpore 版本最后创建对应的 Python 虚拟环境。这里给出一个实际可用的组合示例昇腾 910B 芯片驱动版本 24.1.rc1Python 3.9MindSpore 2.3.0CANN 8.0.RC1mindspore-transformers 0.3.0安装命令一般是这样conda create -n ms230 python3.9 -y conda activate ms230 pip install mindspore2.3.0 pip install mindspore-transformers0.3.0 pip install datasets tokenizers numpy sentencepiece注意MindSpore 的 pip 包已经从 PyPI 移到了昇思官方源所以需要加-i https://ms-release.obs.cn-north-4.myhuaweicloud.com/...这类地址。这个细节官方文档有写但很多人没注意导致pip install mindspore装到了旧版本或者直接报找不到包。安装完成后跑一个最简单的验证脚本import mindspore as ms from mindspore import Tensor ms.set_context(device_targetAscend) x Tensor([1.0, 2.0, 3.0]) print(x.sum())能正常输出结果说明环境基本没问题。2.2 数据集的准备与流式加载数据准备是预训练里最耗时、最容易出错的部分。昇思的mindspore.dataset提供了GeneratorDataset、MindDataset、TFRecordDataset等接口。对于超大规模语料我建议把原始文本转换成MindRecord格式也就是昇思自己的二进制数据格式加载速度和随机读性能都远好于直接读文本。转换流程大概是这样的先把原始文本清洗、去重、过滤低质量内容然后用 Tokenizer 把文本切成 token ids再加上特殊 token、padding、truncation最后写入 MindRecord。这一步可以在离线阶段用多进程并行处理我当时是用 32 个进程跑了大概 6 个小时把 5TB 的原始文本转成了约 3TB 的 MindRecord 文件。有一个细节需要特别留意即使你用MindDataset读取训练时也经常需要在map操作里做动态 padding 或者 mask 构造。昇思的map操作支持 Python 函数但函数会在数据管线的 worker 进程里执行所以不要在函数里引用训练进程的大对象否则序列化会很慢。一个典型的预处理函数例子import mindspore.dataset as ds import numpy as np def preprocess_for_gpt(example): input_ids example[input_ids] # 构造标签向右偏移一位 labels input_ids[1:] [tokenizer.pad_token_id] # 对 padding 部分做 mask attention_mask [1] * len(input_ids) return input_ids, labels, attention_mask dataset ds.MindDataset(dataset_filesdata/*.mindrecord, num_shardsrank_size, shard_idrank_id, shuffleTrue) dataset dataset.map(operationspreprocess_for_gpt, input_columns[input_ids], output_columns[input_ids, labels, attention_mask]) dataset dataset.batch(batch_size8, drop_remainderTrue)注意num_shards和shard_id这个参数在多卡训练时每个卡读取不同的数据分片避免数据重复。这个机制是昇思数据集模块自带的分布式采样比 PyTorch 的DistributedSampler更直接。2.3 Tokenizer 的一致性问题Tokenizer 是预训练中一个容易被忽略但极其关键的组件。训练时用的 tokenizer 和推理时用的 tokenizer 必须完全一致否则模型效果会大打折扣。昇思这边没有自己的 tokenizer 库通常做法是直接用 HuggingFace 的tokenizers库或者transformers里的 tokenizer在预处理阶段把文本转换成 token ids 存到 MindRecord 里。这样做的好处是训练阶段完全不依赖 tokenizer 库加快数据加载坏处是一旦需要在线处理数据比如某些动态任务你得在训练进程里也加载 tokenizer那就要保证 tokenizer 库的版本一致。我的建议是预训练阶段全部离线转成 ids不要在线 tokenize能省不少事。3. 模型定义与并行策略核心细节3.1 用 nn.Cell 搭模型昇思的模型定义继承mindspore.nn.Cell和 PyTorch 的nn.Module很像但有几个差异点需要注意。比如Cell的__init__里必须调用super().__init__()而construct方法相当于 PyTorch 的forward。一个最小化的 GPT 模型结构示例import mindspore as ms import mindspore.nn as nn from mindspore.common.initializer import Normal class GPTBlock(nn.Cell): def __init__(self, hidden_size, num_heads, intermediate_size): super().__init__() self.layernorm1 nn.LayerNorm((hidden_size,)) self.attention nn.MultiheadAttention(hidden_size, num_heads) self.layernorm2 nn.LayerNorm((hidden_size,)) self.mlp nn.SequentialCell([ nn.Dense(hidden_size, intermediate_size), nn.GELU(), nn.Dense(intermediate_size, hidden_size) ]) def construct(self, x, maskNone): x x self.attention(self.layernorm1(x), mask) x x self.mlp(self.layernorm2(x)) return x class GPTModel(nn.Cell): def __init__(self, vocab_size, hidden_size, num_layers, num_heads): super().__init__() self.embedding nn.Embedding(vocab_size, hidden_size) self.blocks nn.CellList([GPTBlock(hidden_size, num_heads, 4*hidden_size) for _ in range(num_layers)]) self.layernorm nn.LayerNorm((hidden_size,)) self.lm_head nn.Dense(hidden_size, vocab_size) def construct(self, input_ids): x self.embedding(input_ids) for block in self.blocks: x block(x) x self.layernorm(x) logits self.lm_head(x) return logits我在实际使用中把nn.MultiheadAttention换成了自定义的 Attention 实现因为预训练通常需要加 attention mask、position bias 等额外处理。如果你需要精细控制建议直接看 MindSpore Transformers 仓库里的mindspore_transformers/models/gpt实现里面有完整的 RMSNorm、旋转位置编码、GQA 等实现直接改参数字典就能用。3.2 并行策略数据并行、模型并行、流水线并行怎么选昇思的自动并行是一个非常值得了解的机制。它有两种用法第一种是ms.set_auto_parallel_context(parallel_modems.ParallelMode.DATA_PARALLEL)也就是纯数据并行。每张卡上有完整的模型副本数据按卡切分梯度进行 all-reduce。这个模式适合模型单卡能放下、但数据量很大的场景比如 7B 或者 13B 级别。第二种是SEMI_AUTO_PARALLEL或者AUTO_PARALLEL让框架自动插入通信集合通信原语支持算子级模型并行、流水线并行、优化器状态切分。13B 以上模型或者单卡显存撑不下的场景必须用这个模式。实际配置示例import mindspore as ms ms.set_context(modems.GRAPH_MODE, device_targetAscend) ms.set_auto_parallel_context( parallel_modems.ParallelMode.SEMI_AUTO_PARALLEL, gradients_meanTrue, comm_fusion{gradient: 8}, # 梯度通信融合 pipeline_stages2, # 流水线并行阶段数 )这里pipeline_stages2的含义是把 Transformer 层分成两组分别放在不同的设备上层与层之间通过点到点通信传递激活值和梯度。更大的模型可以把流水线阶段数调大。你可能注意到我没有在模型代码里手动添加任何通信操作这是因为昇思的编译图会自动分析哪些算子需要跨设备通信并在图上插入合适的通信原语。这个机制非常强大但也意味着你不能在construct里写不规范的 Python 控制流否则自动并行分析可能失败。调试期建议先在PYNATIVE_MODE下跑通再切到GRAPH_MODE。3.3 优化器混合精度与梯度累积大模型预训练几乎必用混合精度。昇思里用ms.amp模块来管理from mindspore import amp model GPTModel(...) optimizer nn.AdamWeightDecay(model.trainable_params(), learning_rate3e-4) loss_scaler amp.DynamicLossScaler(scale_value2**16, scale_factor2, scale_window2000) model amp.auto_mixed_precision(model, amp_levelO2)amp_levelO2表示大部分算子用 FP16 计算部分对精度敏感的算子如 LayerNorm、Softmax保持 FP32。这个设置对预训练效果影响很大如果你发现损失在训练中不稳定可以把O2降到O1或者保守地自定义哪些算子保持 FP32。梯度累积在昇思里是通过nn.TrainOneStepCell加上一个累积器实现的一个好的方式是直接用mindspore.nn.wrap.GradAccumulationfrom mindspore.nn.wrap import GradAccumulation accum_steps 16 net_with_loss nn.WithLossCell(model, loss_fn) train_net GradAccumulation(net_with_loss, optimizer, accum_steps)梯度累积的意义在于当单卡 batch size 受显存限制比较小时你可以通过累积多个 mini-batch 的梯度来等效一个大 batch。比如单卡 batch size 只能设为 8累积 16 步就相当于 128 的全局 batch size。4. 训练实操从单卡到多卡4.1 单卡调试与动态图模式正式大规模训练之前务必先在小数据上做单卡调试。我通常会用 1000 条样本跑 5 到 10 个 step验证损失是否能正常下降、显存占用是否在预期范围内、梯度是否为正常数值。调试时使用动态图模式ms.set_context(modems.PYNATIVE_MODE, device_targetAscend)动态图模式下报错信息更直观可以像 PyTorch 一样打断点、打印张量。但注意动态图模式下不要做大规模并行否则性能极差。调试完再切回GRAPH_MODE。4.2 多卡启动rank 和卡号的坑多卡训练需要正确设置设备 ID 和 rank 信息。最简单的启动方式是用昇思的modelarts或rank_table方式但如果你在裸机环境可以用mpirunmpirun -n 8 python train.py --config configs/gpt_7b.yaml在train.py里你需要读取rank_id和rank_size通过环境变量RANK_ID、RANK_SIZE获取。注意昇思的设备 ID物理卡号和 rank ID逻辑编号可能不一致需要通过如下方式绑定import os import mindspore as ms rank_id int(os.getenv(RANK_ID, 0)) rank_size int(os.getenv(RANK_SIZE, 1)) device_id int(os.getenv(DEVICE_ID, rank_id % 8)) ms.set_context(device_iddevice_id)如果你用的是mpirun且没有设置RANK_ID有可能所有进程都读到默认值 0导致所有进程都使用 0 号卡这是新手最常犯的错误。我的做法是在启动脚本里显式传入mpirun -n 8 -H localhost:8 bash -c export RANK_ID$OMPI_COMM_WORLD_RANK python train.py4.3 损失函数与训练循环预训练分类任务的损失函数就是交叉熵但在昇思上直接使用nn.CrossEntropyLoss需要留意标签形状和 ignore index。一个实用写法import mindspore.nn as nn class PretrainLoss(nn.Cell): def __init__(self, vocab_size): super().__init__() self.loss_fn nn.CrossEntropyLoss(ignore_index-100) self.vocab_size vocab_size def construct(self, logits, labels): # logits: [batch, seq_len, vocab_size] # labels: [batch, seq_len] logits logits.view(-1, self.vocab_size) labels labels.view(-1) return self.loss_fn(logits, labels)训练循环可以不用手写直接用mindspore.Model的高级 APIfrom mindspore import Model, TrainOneStepCell model GPTModel(...) loss_fn PretrainLoss(...) optimizer nn.AdamWeightDecay(...) net_with_loss nn.WithLossCell(model, loss_fn) train_net TrainOneStepCell(net_with_loss, optimizer) for epoch in range(num_epochs): for step, (input_ids, labels, attention_mask) in enumerate(dataset): loss train_net(input_ids, labels, attention_mask) if step % 100 0: print(fstep: {step}, loss: {loss})4.4 Checkpoint 保存与断点续训训练几周突然因为机器故障中断如果没有保存 checkpoint前面的时间就全浪费了。昇思的 checkpoint 管理有几个接口要熟悉ms.save_checkpoint(network, path)保存模型参数ms.load_checkpoint(path, network)加载模型参数ms.load_param_into_net将 checkpoint 中的参数更新到网络中对于大模型我建议保存以下内容模型参数、优化器参数、step 计数、RNG 状态、学习率调度器状态。昇思的CheckpointConfig可以自动定期保存from mindspore.train.callback import CheckpointConfig, ModelCheckpoint ckpt_config CheckpointConfig( save_checkpoint_steps1000, keep_checkpoint_max5, save_checkpoint_seconds3600 ) ckpt_callback ModelCheckpoint(prefixgpt7b, directory./ckpt, configckpt_config)但注意ModelCheckpoint默认只保存network参数优化器状态不会自动保存。等你恢复训练时如果优化器的动量、方差信息缺失前期训练的节奏会被打断。我的做法是自定义一个 callback在保存 checkpoint 时同时保存 optimizer 的状态class FullCheckpointCallback(Callback): def __init__(self, network, optimizer, config): self.network network self.optimizer optimizer self.config config def step_end(self, run_context): cb_params run_context.original_args() step cb_params.cur_step_num if step % self.config.save_steps 0: ckpt_path f./ckpt/gpt7b_step_{step}.ckpt ms.save_checkpoint(self.network, ckpt_path) opt_path f./ckpt/gpt7b_step_{step}_optimizer.ckpt ms.save_checkpoint(self.optimizer, opt_path)恢复训练时先加载模型参数再加载优化器参数继续从 step 处开始即可。5. 推理部署与模型导出5.1 从训练态到推理态的转换预训练完成后需要把模型从训练状态切到推理状态。这里有三个关键操作一是去掉损失函数和优化器只保留model本身二是把model.set_train(False)确保 dropout 等训练相关的层被关闭三是如果你想用 MindSpore Lite 做推理需要把模型导出为 MindIR 格式import numpy as np import mindspore as ms from mindspore import Tensor model.set_train(False) input_ids Tensor(np.ones((1, 128), dtypenp.int32)) ms.export(model, input_ids, file_namegpt7b, file_formatMINDIR)导出时注意ms.export需要传入一个样例输入来确定输入 shape 和 dtype而且这个 shape 要和你在推理时用的 shape 一致否则导出后推理可能报 shape 不匹配的错。5.2 自回归生成推理示例如果你只是做一个简单的在线推理 demo不需要导出 MindIR用 Python 直接跑construct也行。这里给出一个最基础的自回归生成流程def generate(model, tokenizer, prompt, max_new_tokens64): model.set_train(False) input_ids tokenizer.encode(prompt, return_tensorsnp) for _ in range(max_new_tokens): logits model(Tensor(input_ids, ms.int32)) next_token_logits logits[0, -1, :] next_token int(next_token_logits.argmax()) input_ids np.concatenate([input_ids, [[next_token]]], axis1) if next_token tokenizer.eos_token_id: break return tokenizer.decode(input_ids[0])注意这里的 logits 形状是[batch, seq_len, vocab_size]每次迭代都要重新跑一遍完整的前向。如果序列很长性能会比较差实际生产建议用 KV Cache 优化。昇思的nn.MultiheadAttention没有自动暴露 KV cache需要自己实现官方 MindSpore Transformers 仓库里 GPT 模型有带 KV cache 的 generate 接口可以直接参考。5.3 部署推理时的一个精度细节如果你在训练时用了混合精度amp_levelO2推理时要保持一致。很多模型在训练时精度正常部署时损失飙升、生成质量差就是因为推理时没有开启同样的混合精度设置。建议在推理代码里也调用model amp.auto_mixed_precision(model, amp_levelO2) model.set_train(False)6. 热门工具链与生态协同6.1 MindSpore Transformers 与 ModelZoo如果你不想从零搭模型建议直接使用mindspore-transformers这个库它在mindspore.models目录下维护了大量预训练模型的实现包括GPT、BERT、Llama、Bloom等。安装后可以这样加载模型from mindformers import GPT2LMHeadModel, GPT2Config config GPT2Config.from_pretrained(gpt2) model GPT2LMHeadModel(config)也可以直接用官方训练好的权重做推理或微调节省大量训练时间。但注意从 HuggingFace 下载的权重不能直接加载到昇思模型里需要通过转换脚本把pytorch_model.bin转成昇思的 ckpt 格式。这类转换脚本在mindformers的tools目录下都有用convert_weight命令就可以。6.2 与 HuggingFace tokenizer 的配合昇思的模型输入输出格式和 HuggingFace 不完全一致。HuggingFace 的BertTokenizer返回的是dict而昇思的模型construct常常接受位置参数。所以我在写数据管线时总是先手动把 tokenizer 的返回值拆开按模型需要的参数顺序传参。这里有一个最佳实践预处理阶段不要保存 tokenizer 的对象直接保存 token ids 和 attention mask 到 MindRecord这样训练代码完全不需要 import transformers也避免了 tokenizer 库版本升级导致的不一致问题。6.3 模型微调与 LoRA 进阶预训练完成后通常还要做领域微调或指令微调。昇思也支持 LoRA 这类参数高效微调方法。mindformers里的PeftPretrainedModel就是为 LoRA 设计的你只需要在加载模型时指定lora_rank和lora_alphafrom mindformers import GPT2LMHeadModel model GPT2LMHeadModel.from_pretrained( gpt2, lora_rank8, lora_alpha16, lora_dropout0.1 )LoRA 微调时冻结大部分参数只训练低秩矩阵显存占用和训练速度都比全参数微调理想得多。我通常先用 LoRA 跑一个快速实验确认效果方向没问题再决定是否做全参数微调。6.4 常见的模型切换与权重迁移从 PyTorch 迁移一个已经训练好的模型到昇思主要工作是权重 key 的映射。HuggingFace 的模型参数名一般带transformer.h.0.attn.c_attn.weight这种前缀而昇思模型的参数名可能是blocks.0.attention.query.weight需要写一个映射字典来转换。mindformers提供了自动转换功能但遇到自定义结构还是要自己写脚本。一个简单的映射示例def convert_pytorch_to_ms(state_dict): new_state_dict {} for key, value in state_dict.items(): new_key key.replace(transformer.h, blocks) new_key new_key.replace(attn.c_attn, attention.qkv) new_state_dict[new_key] value return new_state_dict看起来简单但实际操作中多层嵌套的 key 替换很容易出错。建议转换后用一个随机输入做前向对齐测试对比 PyTorch 和昇思的输出是否一致。7. 常见问题排查与性能调优实录7.1 启动训练时卡住或报错常见情况一runtime error: ascend device init failed。这个问题八成是设备 ID 设置错误或者ASCEND_VISIBLE_DEVICES环境变量没设置。检查你的物理卡是否被其他进程占用用npu-smi info查看。常见情况二多卡启动后只有 0 号卡在跑。这是因为你没有正确区分 rank 和 device或者数据集shard_id没有按 rank 分配。7.2 训练 loss 不下降或出现 NaNloss 不下降先检查学习率。大模型通常要用带 warmup 的余弦学习率初始学习率一般在1e-4到3e-4warmup 步数占训练总步数的 1% 到 2%。如果学习率没问题再检查数据标签是否正确偏移mask 是否正确有没有把 padding 位置的标签也计算损失了出现 NaN大概率是混合精度问题。先切到纯 FP32 试一次如果能跑通说明 FP16 下某些算子溢出了需要给网络添加手动loss_scaler或者修改loss_scale策略。7.3 显存不够怎么办显存不够时按优先级依次尝试减小 batch size、开启梯度累积、开启激活重计算recompute、使用模型并行、切换混合精度等级到更激进。昇思的激活重计算配置很简单from mindspore.nn import recompute for block in model.blocks: recompute(block)这会在前向时保留部分中间激活反向时重新计算典型用显存换时间的方案。一般来说开启重计算后显存占用能下降 40% 以上。7.4 性能调优经验单卡训练时先用npu-smi info看算力利用率是否打满。如果利用率低问题通常出在数据加载上。检查 MindRecord 文件是否有足够多的分片文件、num_parallel_workers是否设置合理、shuffle buffer 大小是否合适。我一般把num_parallel_workers8shuffle buffer 设为训练数据一个 epoch 的 0.1% 左右。多卡训练时关注梯度通信耗时。如果通信占比过高增大comm_fusion的融合值把多个梯度打包一次通信发送。还有一种技巧是开启梯度压缩昇思支持optimizer的gradient_compression_typeGradientCompressType.NONE如果带宽瓶颈你还可以试GradientCompressType.ON。另外AMP level 的选择也直接影响收敛速度和质量。不要贪图 O2 的加速而忽略了精度风险可以从 O1 开始训观察几个 epoch 的 loss 曲线无异常后再换 O2。7.5 断点续训中的隐蔽坑断点续训最隐蔽的问题是学习率调度器状态没恢复。如果你保存了模型参数和优化器参数但把global_step重置为 0学习率会重新从最高点开始相当于浪费了前面的 warmup 过程还会导致 loss 突然拉升。所以我在保存 checkpoint 时总是额外保存一个 JSON 文件记录当前 step、epoch、学习率、RNG 状态。恢复训练时把这个 step 传给学习率调度器确保学习率曲线是连续的。7.6 关于模型评测的一个建议预训练阶段就要提前想好评测方案。很多团队把模型训完才发现评测代码和环境没准备此时回滚成本极高。我的习惯是在训练脚本里就内置一个轻量评测函数每隔 N 步用一个小验证集跑一下 perplexity同时保存一组固定 prompt 的生成样例直观观察训练效果。这样你能尽早发现模型是否过拟合、是否出现重复生成等问题。8. 总结与长期迭代的一些体会从 PyTorch 迁移到昇思做预训练初期确实有学习成本尤其是数据管线和并行策略的配置方式跟之前很不相同。但是一旦跑通昇思在昇腾硬件上的性能表现是实打实的尤其大规模自动并行这块省去了很多手工切分模型的痛苦。如果让我给一个实施顺序的建议那就是先用mindformers现成脚本跑通一个小模型再逐步替换成自己的数据和自定义模型最后再引入并行策略。不要一上来就追求 70B 甚至更大规模在小模型上建立正确的工作流比什么都重要。最后分享一个小技巧昇思的set_context(modeGRAPH_MODE)编译时间会比较长特别是大模型第一次编译可能要等十几分钟。如果只是小改动调参尽量保留PYNATIVE_MODE去跑等确认代码逻辑没有问题了再切图模式做正式训练。这个小习惯能帮你节省大量调试等待时间。大模型预训练这条路没有捷径但工具链选对了、流程打通了后面就是耐心调参和等待 loss 曲线慢慢下降的事了。希望这篇分享能帮你在昇思上少走一些弯路顺利把预训练任务跑起来。
昇思MindSpore大模型预训练实战：流程、并行与调优全解析
先交代一个背景这两年大模型相关的工具链、框架、教程层出不穷但大部分资料都是围绕 PyTorch 生态展开的。昇思 MindSpore 作为另一条完整的技术栈在国产算力适配、自动并行、动静统一等方面做了很多差异化的设计尤其适合有昇腾硬件、或者希望绕开 CUDA 依赖的团队。我大概从去年开始把一部分预训练任务从 PyTorch 迁移到昇思上跑中间踩了不少坑也总结出了一套相对通用的操作路径。这篇文章不吹不黑纯粹从实操角度把昇思上做大模型预训练的通用流程、关键设计、常见问题梳理一遍给想入坑或者正在迁移的同行做一个参考。1. 整体流程设计与方案选型1.1 昇思预训练到底在做什么预训练这件事本质上就是在大规模无标注文本或者图文、语音上通过自监督目标函数让模型学会语言的统计规律和通用表示。昇思上的预训练流程和你在 PyTorch 里做的事情没有本质区别但因为底层实现不同整个流程的组织方式有明显差异。一个标准的昇思预训练任务至少包含以下几个环节数据集准备与流式加载、模型结构定义与初始化、并行策略配置如果你有多卡或集群、优化器与学习率调度、训练循环前向、反向、梯度累积、参数更新、checkpoint 保存与恢复、日志与指标监控、最后是评估和推理验证。我在 PyTorch 时代习惯于把数据集全部加载到内存再 shuffle到了昇思这里很快发现行不通。因为昇思的数据处理管线是基于GeneratorDatasetmapbatch的流式架构数据不会一次性全部进内存而是按需生成和变换。这个设计对超大语料几十 TB 甚至上百 TB非常友好但需要你调整自己的编程习惯。1.2 为什么选择昇思而不是 PyTorch这是每个迁移者都会问的问题。我当时的出发点很朴素团队里有昇腾卡而昇思在昇腾上的算子适配和显存管理确实做得更细。PyTorch 虽然可以通过各种方式在昇腾上跑但总有一些算子要自己适配分布式通信的效率也差一些。从框架设计的角度看昇思有几个对预训练特别友好的特性一是动静统一。你可以先用动态图模式快速调试模型结构和训练逻辑等确认无误后用set_context(modeGRAPH_MODE)切到静态图模式做高性能训练。PyTorch 里torch.compile也能做到类似效果但昇思的动静切换在工程上更透明调试期和生产期共用同一套代码。二是自动并行。昇思的AutoParallel可以根据模型大小、数据量、集群拓扑自动搜索最优的并行策略包括数据并行、算子级模型并行、流水线并行。你不需要像 Megatron-LM 那样手动切分 Transformer 层框架会在计算图编译阶段帮你完成切分和通信插入。三是统一的 Checkpoint 管理。昇思把模型参数、优化器状态、RNG 状态统一保存恢复训练时能精确回到中断点这对于动不动训练几周的大模型来说非常关键。当然PyTorch 生态更丰富社区资源多HuggingFace 上的模型基本能直接加载。昇思也有 MindSpore Transformers 仓库和 ModelZoo很多主流模型结构都能直接拿到预训练权重或直接跑通训练脚本。如果你主要工作在昇腾硬件上昇思的综合体验是优于 PyTorch 的。1.3 预训练方法选择的三个层次我把昇思上做预训练的方法分为三个层次你可以根据自己的硬件资源、模型规模和技术储备来选择。第一层使用现成脚本和模型仓库。昇思提供了mindspore/models仓库里面有 GPT、BERT、Llama 等模型的完整预训练脚本包括数据处理、模型定义、训练循环、并行策略。你的工作是准备数据、调整超参数、启动训练。这个层次适合第一次接触昇思、或者只想快速验证硬件的团队。第二层在现有模型结构上做定制。比如你要改注意力机制、换位置编码、或者调整 MoE 结构那么可以在 MindSpore Transformers 的模型代码基础上修改训练脚本复用现成的Trainer接口。这个层次需要你对模型实现和昇思 API 有一定理解但不需要从零写分布式逻辑。第三层完全自定义模型和训练流程。你要自己定义nn.Cell自己写TrainOneStepCell自己配置并行策略。这个层次灵活性最大但工作量也最大适合做前沿算法研究的团队。我个人建议无论你处于哪个层次先把官方仓库的标准训练脚本跑通一遍再动手改自己的东西。这样你能确认环境、数据格式、训练流程都已经正确后续排错的范围会小很多。2. 环境准备与数据管线搭建2.1 安装与版本匹配昇思的版本和硬件驱动、固件、Python 版本有严格对应关系这是新手最容易卡住的第一关。我建议先确定你的硬件型号和昇腾驱动版本再根据官方版本匹配表选择 MindSpore 版本最后创建对应的 Python 虚拟环境。这里给出一个实际可用的组合示例昇腾 910B 芯片驱动版本 24.1.rc1Python 3.9MindSpore 2.3.0CANN 8.0.RC1mindspore-transformers 0.3.0安装命令一般是这样conda create -n ms230 python3.9 -y conda activate ms230 pip install mindspore2.3.0 pip install mindspore-transformers0.3.0 pip install datasets tokenizers numpy sentencepiece注意MindSpore 的 pip 包已经从 PyPI 移到了昇思官方源所以需要加-i https://ms-release.obs.cn-north-4.myhuaweicloud.com/...这类地址。这个细节官方文档有写但很多人没注意导致pip install mindspore装到了旧版本或者直接报找不到包。安装完成后跑一个最简单的验证脚本import mindspore as ms from mindspore import Tensor ms.set_context(device_targetAscend) x Tensor([1.0, 2.0, 3.0]) print(x.sum())能正常输出结果说明环境基本没问题。2.2 数据集的准备与流式加载数据准备是预训练里最耗时、最容易出错的部分。昇思的mindspore.dataset提供了GeneratorDataset、MindDataset、TFRecordDataset等接口。对于超大规模语料我建议把原始文本转换成MindRecord格式也就是昇思自己的二进制数据格式加载速度和随机读性能都远好于直接读文本。转换流程大概是这样的先把原始文本清洗、去重、过滤低质量内容然后用 Tokenizer 把文本切成 token ids再加上特殊 token、padding、truncation最后写入 MindRecord。这一步可以在离线阶段用多进程并行处理我当时是用 32 个进程跑了大概 6 个小时把 5TB 的原始文本转成了约 3TB 的 MindRecord 文件。有一个细节需要特别留意即使你用MindDataset读取训练时也经常需要在map操作里做动态 padding 或者 mask 构造。昇思的map操作支持 Python 函数但函数会在数据管线的 worker 进程里执行所以不要在函数里引用训练进程的大对象否则序列化会很慢。一个典型的预处理函数例子import mindspore.dataset as ds import numpy as np def preprocess_for_gpt(example): input_ids example[input_ids] # 构造标签向右偏移一位 labels input_ids[1:] [tokenizer.pad_token_id] # 对 padding 部分做 mask attention_mask [1] * len(input_ids) return input_ids, labels, attention_mask dataset ds.MindDataset(dataset_filesdata/*.mindrecord, num_shardsrank_size, shard_idrank_id, shuffleTrue) dataset dataset.map(operationspreprocess_for_gpt, input_columns[input_ids], output_columns[input_ids, labels, attention_mask]) dataset dataset.batch(batch_size8, drop_remainderTrue)注意num_shards和shard_id这个参数在多卡训练时每个卡读取不同的数据分片避免数据重复。这个机制是昇思数据集模块自带的分布式采样比 PyTorch 的DistributedSampler更直接。2.3 Tokenizer 的一致性问题Tokenizer 是预训练中一个容易被忽略但极其关键的组件。训练时用的 tokenizer 和推理时用的 tokenizer 必须完全一致否则模型效果会大打折扣。昇思这边没有自己的 tokenizer 库通常做法是直接用 HuggingFace 的tokenizers库或者transformers里的 tokenizer在预处理阶段把文本转换成 token ids 存到 MindRecord 里。这样做的好处是训练阶段完全不依赖 tokenizer 库加快数据加载坏处是一旦需要在线处理数据比如某些动态任务你得在训练进程里也加载 tokenizer那就要保证 tokenizer 库的版本一致。我的建议是预训练阶段全部离线转成 ids不要在线 tokenize能省不少事。3. 模型定义与并行策略核心细节3.1 用 nn.Cell 搭模型昇思的模型定义继承mindspore.nn.Cell和 PyTorch 的nn.Module很像但有几个差异点需要注意。比如Cell的__init__里必须调用super().__init__()而construct方法相当于 PyTorch 的forward。一个最小化的 GPT 模型结构示例import mindspore as ms import mindspore.nn as nn from mindspore.common.initializer import Normal class GPTBlock(nn.Cell): def __init__(self, hidden_size, num_heads, intermediate_size): super().__init__() self.layernorm1 nn.LayerNorm((hidden_size,)) self.attention nn.MultiheadAttention(hidden_size, num_heads) self.layernorm2 nn.LayerNorm((hidden_size,)) self.mlp nn.SequentialCell([ nn.Dense(hidden_size, intermediate_size), nn.GELU(), nn.Dense(intermediate_size, hidden_size) ]) def construct(self, x, maskNone): x x self.attention(self.layernorm1(x), mask) x x self.mlp(self.layernorm2(x)) return x class GPTModel(nn.Cell): def __init__(self, vocab_size, hidden_size, num_layers, num_heads): super().__init__() self.embedding nn.Embedding(vocab_size, hidden_size) self.blocks nn.CellList([GPTBlock(hidden_size, num_heads, 4*hidden_size) for _ in range(num_layers)]) self.layernorm nn.LayerNorm((hidden_size,)) self.lm_head nn.Dense(hidden_size, vocab_size) def construct(self, input_ids): x self.embedding(input_ids) for block in self.blocks: x block(x) x self.layernorm(x) logits self.lm_head(x) return logits我在实际使用中把nn.MultiheadAttention换成了自定义的 Attention 实现因为预训练通常需要加 attention mask、position bias 等额外处理。如果你需要精细控制建议直接看 MindSpore Transformers 仓库里的mindspore_transformers/models/gpt实现里面有完整的 RMSNorm、旋转位置编码、GQA 等实现直接改参数字典就能用。3.2 并行策略数据并行、模型并行、流水线并行怎么选昇思的自动并行是一个非常值得了解的机制。它有两种用法第一种是ms.set_auto_parallel_context(parallel_modems.ParallelMode.DATA_PARALLEL)也就是纯数据并行。每张卡上有完整的模型副本数据按卡切分梯度进行 all-reduce。这个模式适合模型单卡能放下、但数据量很大的场景比如 7B 或者 13B 级别。第二种是SEMI_AUTO_PARALLEL或者AUTO_PARALLEL让框架自动插入通信集合通信原语支持算子级模型并行、流水线并行、优化器状态切分。13B 以上模型或者单卡显存撑不下的场景必须用这个模式。实际配置示例import mindspore as ms ms.set_context(modems.GRAPH_MODE, device_targetAscend) ms.set_auto_parallel_context( parallel_modems.ParallelMode.SEMI_AUTO_PARALLEL, gradients_meanTrue, comm_fusion{gradient: 8}, # 梯度通信融合 pipeline_stages2, # 流水线并行阶段数 )这里pipeline_stages2的含义是把 Transformer 层分成两组分别放在不同的设备上层与层之间通过点到点通信传递激活值和梯度。更大的模型可以把流水线阶段数调大。你可能注意到我没有在模型代码里手动添加任何通信操作这是因为昇思的编译图会自动分析哪些算子需要跨设备通信并在图上插入合适的通信原语。这个机制非常强大但也意味着你不能在construct里写不规范的 Python 控制流否则自动并行分析可能失败。调试期建议先在PYNATIVE_MODE下跑通再切到GRAPH_MODE。3.3 优化器混合精度与梯度累积大模型预训练几乎必用混合精度。昇思里用ms.amp模块来管理from mindspore import amp model GPTModel(...) optimizer nn.AdamWeightDecay(model.trainable_params(), learning_rate3e-4) loss_scaler amp.DynamicLossScaler(scale_value2**16, scale_factor2, scale_window2000) model amp.auto_mixed_precision(model, amp_levelO2)amp_levelO2表示大部分算子用 FP16 计算部分对精度敏感的算子如 LayerNorm、Softmax保持 FP32。这个设置对预训练效果影响很大如果你发现损失在训练中不稳定可以把O2降到O1或者保守地自定义哪些算子保持 FP32。梯度累积在昇思里是通过nn.TrainOneStepCell加上一个累积器实现的一个好的方式是直接用mindspore.nn.wrap.GradAccumulationfrom mindspore.nn.wrap import GradAccumulation accum_steps 16 net_with_loss nn.WithLossCell(model, loss_fn) train_net GradAccumulation(net_with_loss, optimizer, accum_steps)梯度累积的意义在于当单卡 batch size 受显存限制比较小时你可以通过累积多个 mini-batch 的梯度来等效一个大 batch。比如单卡 batch size 只能设为 8累积 16 步就相当于 128 的全局 batch size。4. 训练实操从单卡到多卡4.1 单卡调试与动态图模式正式大规模训练之前务必先在小数据上做单卡调试。我通常会用 1000 条样本跑 5 到 10 个 step验证损失是否能正常下降、显存占用是否在预期范围内、梯度是否为正常数值。调试时使用动态图模式ms.set_context(modems.PYNATIVE_MODE, device_targetAscend)动态图模式下报错信息更直观可以像 PyTorch 一样打断点、打印张量。但注意动态图模式下不要做大规模并行否则性能极差。调试完再切回GRAPH_MODE。4.2 多卡启动rank 和卡号的坑多卡训练需要正确设置设备 ID 和 rank 信息。最简单的启动方式是用昇思的modelarts或rank_table方式但如果你在裸机环境可以用mpirunmpirun -n 8 python train.py --config configs/gpt_7b.yaml在train.py里你需要读取rank_id和rank_size通过环境变量RANK_ID、RANK_SIZE获取。注意昇思的设备 ID物理卡号和 rank ID逻辑编号可能不一致需要通过如下方式绑定import os import mindspore as ms rank_id int(os.getenv(RANK_ID, 0)) rank_size int(os.getenv(RANK_SIZE, 1)) device_id int(os.getenv(DEVICE_ID, rank_id % 8)) ms.set_context(device_iddevice_id)如果你用的是mpirun且没有设置RANK_ID有可能所有进程都读到默认值 0导致所有进程都使用 0 号卡这是新手最常犯的错误。我的做法是在启动脚本里显式传入mpirun -n 8 -H localhost:8 bash -c export RANK_ID$OMPI_COMM_WORLD_RANK python train.py4.3 损失函数与训练循环预训练分类任务的损失函数就是交叉熵但在昇思上直接使用nn.CrossEntropyLoss需要留意标签形状和 ignore index。一个实用写法import mindspore.nn as nn class PretrainLoss(nn.Cell): def __init__(self, vocab_size): super().__init__() self.loss_fn nn.CrossEntropyLoss(ignore_index-100) self.vocab_size vocab_size def construct(self, logits, labels): # logits: [batch, seq_len, vocab_size] # labels: [batch, seq_len] logits logits.view(-1, self.vocab_size) labels labels.view(-1) return self.loss_fn(logits, labels)训练循环可以不用手写直接用mindspore.Model的高级 APIfrom mindspore import Model, TrainOneStepCell model GPTModel(...) loss_fn PretrainLoss(...) optimizer nn.AdamWeightDecay(...) net_with_loss nn.WithLossCell(model, loss_fn) train_net TrainOneStepCell(net_with_loss, optimizer) for epoch in range(num_epochs): for step, (input_ids, labels, attention_mask) in enumerate(dataset): loss train_net(input_ids, labels, attention_mask) if step % 100 0: print(fstep: {step}, loss: {loss})4.4 Checkpoint 保存与断点续训训练几周突然因为机器故障中断如果没有保存 checkpoint前面的时间就全浪费了。昇思的 checkpoint 管理有几个接口要熟悉ms.save_checkpoint(network, path)保存模型参数ms.load_checkpoint(path, network)加载模型参数ms.load_param_into_net将 checkpoint 中的参数更新到网络中对于大模型我建议保存以下内容模型参数、优化器参数、step 计数、RNG 状态、学习率调度器状态。昇思的CheckpointConfig可以自动定期保存from mindspore.train.callback import CheckpointConfig, ModelCheckpoint ckpt_config CheckpointConfig( save_checkpoint_steps1000, keep_checkpoint_max5, save_checkpoint_seconds3600 ) ckpt_callback ModelCheckpoint(prefixgpt7b, directory./ckpt, configckpt_config)但注意ModelCheckpoint默认只保存network参数优化器状态不会自动保存。等你恢复训练时如果优化器的动量、方差信息缺失前期训练的节奏会被打断。我的做法是自定义一个 callback在保存 checkpoint 时同时保存 optimizer 的状态class FullCheckpointCallback(Callback): def __init__(self, network, optimizer, config): self.network network self.optimizer optimizer self.config config def step_end(self, run_context): cb_params run_context.original_args() step cb_params.cur_step_num if step % self.config.save_steps 0: ckpt_path f./ckpt/gpt7b_step_{step}.ckpt ms.save_checkpoint(self.network, ckpt_path) opt_path f./ckpt/gpt7b_step_{step}_optimizer.ckpt ms.save_checkpoint(self.optimizer, opt_path)恢复训练时先加载模型参数再加载优化器参数继续从 step 处开始即可。5. 推理部署与模型导出5.1 从训练态到推理态的转换预训练完成后需要把模型从训练状态切到推理状态。这里有三个关键操作一是去掉损失函数和优化器只保留model本身二是把model.set_train(False)确保 dropout 等训练相关的层被关闭三是如果你想用 MindSpore Lite 做推理需要把模型导出为 MindIR 格式import numpy as np import mindspore as ms from mindspore import Tensor model.set_train(False) input_ids Tensor(np.ones((1, 128), dtypenp.int32)) ms.export(model, input_ids, file_namegpt7b, file_formatMINDIR)导出时注意ms.export需要传入一个样例输入来确定输入 shape 和 dtype而且这个 shape 要和你在推理时用的 shape 一致否则导出后推理可能报 shape 不匹配的错。5.2 自回归生成推理示例如果你只是做一个简单的在线推理 demo不需要导出 MindIR用 Python 直接跑construct也行。这里给出一个最基础的自回归生成流程def generate(model, tokenizer, prompt, max_new_tokens64): model.set_train(False) input_ids tokenizer.encode(prompt, return_tensorsnp) for _ in range(max_new_tokens): logits model(Tensor(input_ids, ms.int32)) next_token_logits logits[0, -1, :] next_token int(next_token_logits.argmax()) input_ids np.concatenate([input_ids, [[next_token]]], axis1) if next_token tokenizer.eos_token_id: break return tokenizer.decode(input_ids[0])注意这里的 logits 形状是[batch, seq_len, vocab_size]每次迭代都要重新跑一遍完整的前向。如果序列很长性能会比较差实际生产建议用 KV Cache 优化。昇思的nn.MultiheadAttention没有自动暴露 KV cache需要自己实现官方 MindSpore Transformers 仓库里 GPT 模型有带 KV cache 的 generate 接口可以直接参考。5.3 部署推理时的一个精度细节如果你在训练时用了混合精度amp_levelO2推理时要保持一致。很多模型在训练时精度正常部署时损失飙升、生成质量差就是因为推理时没有开启同样的混合精度设置。建议在推理代码里也调用model amp.auto_mixed_precision(model, amp_levelO2) model.set_train(False)6. 热门工具链与生态协同6.1 MindSpore Transformers 与 ModelZoo如果你不想从零搭模型建议直接使用mindspore-transformers这个库它在mindspore.models目录下维护了大量预训练模型的实现包括GPT、BERT、Llama、Bloom等。安装后可以这样加载模型from mindformers import GPT2LMHeadModel, GPT2Config config GPT2Config.from_pretrained(gpt2) model GPT2LMHeadModel(config)也可以直接用官方训练好的权重做推理或微调节省大量训练时间。但注意从 HuggingFace 下载的权重不能直接加载到昇思模型里需要通过转换脚本把pytorch_model.bin转成昇思的 ckpt 格式。这类转换脚本在mindformers的tools目录下都有用convert_weight命令就可以。6.2 与 HuggingFace tokenizer 的配合昇思的模型输入输出格式和 HuggingFace 不完全一致。HuggingFace 的BertTokenizer返回的是dict而昇思的模型construct常常接受位置参数。所以我在写数据管线时总是先手动把 tokenizer 的返回值拆开按模型需要的参数顺序传参。这里有一个最佳实践预处理阶段不要保存 tokenizer 的对象直接保存 token ids 和 attention mask 到 MindRecord这样训练代码完全不需要 import transformers也避免了 tokenizer 库版本升级导致的不一致问题。6.3 模型微调与 LoRA 进阶预训练完成后通常还要做领域微调或指令微调。昇思也支持 LoRA 这类参数高效微调方法。mindformers里的PeftPretrainedModel就是为 LoRA 设计的你只需要在加载模型时指定lora_rank和lora_alphafrom mindformers import GPT2LMHeadModel model GPT2LMHeadModel.from_pretrained( gpt2, lora_rank8, lora_alpha16, lora_dropout0.1 )LoRA 微调时冻结大部分参数只训练低秩矩阵显存占用和训练速度都比全参数微调理想得多。我通常先用 LoRA 跑一个快速实验确认效果方向没问题再决定是否做全参数微调。6.4 常见的模型切换与权重迁移从 PyTorch 迁移一个已经训练好的模型到昇思主要工作是权重 key 的映射。HuggingFace 的模型参数名一般带transformer.h.0.attn.c_attn.weight这种前缀而昇思模型的参数名可能是blocks.0.attention.query.weight需要写一个映射字典来转换。mindformers提供了自动转换功能但遇到自定义结构还是要自己写脚本。一个简单的映射示例def convert_pytorch_to_ms(state_dict): new_state_dict {} for key, value in state_dict.items(): new_key key.replace(transformer.h, blocks) new_key new_key.replace(attn.c_attn, attention.qkv) new_state_dict[new_key] value return new_state_dict看起来简单但实际操作中多层嵌套的 key 替换很容易出错。建议转换后用一个随机输入做前向对齐测试对比 PyTorch 和昇思的输出是否一致。7. 常见问题排查与性能调优实录7.1 启动训练时卡住或报错常见情况一runtime error: ascend device init failed。这个问题八成是设备 ID 设置错误或者ASCEND_VISIBLE_DEVICES环境变量没设置。检查你的物理卡是否被其他进程占用用npu-smi info查看。常见情况二多卡启动后只有 0 号卡在跑。这是因为你没有正确区分 rank 和 device或者数据集shard_id没有按 rank 分配。7.2 训练 loss 不下降或出现 NaNloss 不下降先检查学习率。大模型通常要用带 warmup 的余弦学习率初始学习率一般在1e-4到3e-4warmup 步数占训练总步数的 1% 到 2%。如果学习率没问题再检查数据标签是否正确偏移mask 是否正确有没有把 padding 位置的标签也计算损失了出现 NaN大概率是混合精度问题。先切到纯 FP32 试一次如果能跑通说明 FP16 下某些算子溢出了需要给网络添加手动loss_scaler或者修改loss_scale策略。7.3 显存不够怎么办显存不够时按优先级依次尝试减小 batch size、开启梯度累积、开启激活重计算recompute、使用模型并行、切换混合精度等级到更激进。昇思的激活重计算配置很简单from mindspore.nn import recompute for block in model.blocks: recompute(block)这会在前向时保留部分中间激活反向时重新计算典型用显存换时间的方案。一般来说开启重计算后显存占用能下降 40% 以上。7.4 性能调优经验单卡训练时先用npu-smi info看算力利用率是否打满。如果利用率低问题通常出在数据加载上。检查 MindRecord 文件是否有足够多的分片文件、num_parallel_workers是否设置合理、shuffle buffer 大小是否合适。我一般把num_parallel_workers8shuffle buffer 设为训练数据一个 epoch 的 0.1% 左右。多卡训练时关注梯度通信耗时。如果通信占比过高增大comm_fusion的融合值把多个梯度打包一次通信发送。还有一种技巧是开启梯度压缩昇思支持optimizer的gradient_compression_typeGradientCompressType.NONE如果带宽瓶颈你还可以试GradientCompressType.ON。另外AMP level 的选择也直接影响收敛速度和质量。不要贪图 O2 的加速而忽略了精度风险可以从 O1 开始训观察几个 epoch 的 loss 曲线无异常后再换 O2。7.5 断点续训中的隐蔽坑断点续训最隐蔽的问题是学习率调度器状态没恢复。如果你保存了模型参数和优化器参数但把global_step重置为 0学习率会重新从最高点开始相当于浪费了前面的 warmup 过程还会导致 loss 突然拉升。所以我在保存 checkpoint 时总是额外保存一个 JSON 文件记录当前 step、epoch、学习率、RNG 状态。恢复训练时把这个 step 传给学习率调度器确保学习率曲线是连续的。7.6 关于模型评测的一个建议预训练阶段就要提前想好评测方案。很多团队把模型训完才发现评测代码和环境没准备此时回滚成本极高。我的习惯是在训练脚本里就内置一个轻量评测函数每隔 N 步用一个小验证集跑一下 perplexity同时保存一组固定 prompt 的生成样例直观观察训练效果。这样你能尽早发现模型是否过拟合、是否出现重复生成等问题。8. 总结与长期迭代的一些体会从 PyTorch 迁移到昇思做预训练初期确实有学习成本尤其是数据管线和并行策略的配置方式跟之前很不相同。但是一旦跑通昇思在昇腾硬件上的性能表现是实打实的尤其大规模自动并行这块省去了很多手工切分模型的痛苦。如果让我给一个实施顺序的建议那就是先用mindformers现成脚本跑通一个小模型再逐步替换成自己的数据和自定义模型最后再引入并行策略。不要一上来就追求 70B 甚至更大规模在小模型上建立正确的工作流比什么都重要。最后分享一个小技巧昇思的set_context(modeGRAPH_MODE)编译时间会比较长特别是大模型第一次编译可能要等十几分钟。如果只是小改动调参尽量保留PYNATIVE_MODE去跑等确认代码逻辑没有问题了再切图模式做正式训练。这个小习惯能帮你节省大量调试等待时间。大模型预训练这条路没有捷径但工具链选对了、流程打通了后面就是耐心调参和等待 loss 曲线慢慢下降的事了。希望这篇分享能帮你在昇思上少走一些弯路顺利把预训练任务跑起来。