Mistral-7B 是开源轻量高效大模型,具备滑动窗口注意力、Rotary 位置编码,推理吞吐高,适合昇腾 910 系列 NPU 部署。昇思 MindSpore 结合 MindFormers 套件,提供 LoRA 轻量化微调、分布式并行、MindRecord 数据集流水线、断点续训全链路能力。相比全参微调,LoRA 仅在 QKV 注意力层注入低秩矩阵,冻结基座权重,显存开销大幅降低,适配单卡 / 多卡昇腾集群,是 Mistral 业务 SFT 首选方案。本文以昇腾 Ascend 后端、MindSpore 2.6、MindFormers 为基础,给出可直接运行的 LoRA 微调主代码,配套数据预处理、训练循环、权重合并推理,同时说明并行配置与踩坑要点。

# mistral_lora_sft.py 昇思MindSpore Mistral-7B LoRA微调,昇腾NPU
import mindspore as ms
import mindspore.nn as nn
from mindspore.dataset import GeneratorDataset
from mindformers import MistralForCausalLM, MistralTokenizer, LoRAConfig, get_peft_model
from mindformers.core.lr import CosineWithWarmUpLR
from mindformers.tools.utils import set_seed
import json
import os

# ===================== 1. 昇腾环境初始化 =====================
ms.set_context(device_target="Ascend", mode=ms.GRAPH_MODE)
set_seed(42)
ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.DATA_PARALLEL, gradients_mean=True)

# ===================== 2. 超参配置 =====================
SEQ_LEN = 1024
BATCH_SIZE = 2
EPOCHS = 3
LR = 2e-4
WARMUP_STEPS = 50
LORA_R = 16
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
MODEL_PATH = "./mistral-7b-ms"
DATA_PATH = "./alpaca_zh.json"
SAVE_LORA_PATH = "./lora_mistral_ckpt"

# ===================== 3. 数据集加载与预处理 =====================
class AlpacaDataset:
    def __init__(self, json_path, tokenizer):
        self.tokenizer = tokenizer
        with open(json_path, "r", encoding="utf-8") as f:
            self.data = json.load(f)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        prompt = f"<s>[INST]{sample['instruction']}\n{sample['input']}[/INST]{sample['output']}</s>"
        token = self.tokenizer(prompt, max_length=SEQ_LEN, padding="max_length", truncation=True, return_tensors="ms")
        input_ids = token["input_ids"][0]
        attention_mask = token["attention_mask"][0]
        labels = input_ids.copy()
        return input_ids, attention_mask, labels

tokenizer = MistralTokenizer.from_pretrained(MODEL_PATH)
train_ds = GeneratorDataset(AlpacaDataset(DATA_PATH, tokenizer), column_names=["input_ids", "attention_mask", "labels"])
train_ds = train_ds.batch(BATCH_SIZE, drop_remainder=True)

# ===================== 4. 加载Mistral基座+LoRA注入 =====================
model = MistralForCausalLM.from_pretrained(MODEL_PATH, ms_dtype=ms.float16)
lora_config = LoRAConfig(
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    lora_dropout=LORA_DROPOUT,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 打印可训练参数量,一般<1%

# ===================== 5. 损失函数、优化器、学习率 =====================
loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
steps_per_epoch = train_ds.get_dataset_size()
total_steps = steps_per_epoch * EPOCHS
lr_scheduler = CosineWithWarmUpLR(learning_rate=LR, warmup_steps=WARMUP_STEPS, total_steps=total_steps)
optimizer = nn.AdamWeightDecay(model.trainable_params(), learning_rate=lr_scheduler, weight_decay=0.01)

# 正向+梯度封装
def forward_fn(input_ids, attn_mask, labels):
    logits = model(input_ids, attention_mask=attn_mask)
    shift_logits = logits[:, :-1, :]
    shift_labels = labels[:, 1:]
    loss = loss_fn(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1))
    return loss, logits

grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

# ===================== 6. 训练循环 =====================
def train():
    model.set_train()
    for epoch in range(EPOCHS):
        for step, (input_ids, attn_mask, labels) in enumerate(train_ds.create_tuple_iterator()):
            (loss, _), grads = grad_fn(input_ids, attn_mask, labels)
            optimizer(grads)
            if step % 10 == 0:
                print(f"Epoch:{epoch+1}, Step:{step}, Loss:{loss.asnumpy():.4f}")
        # 保存LoRA适配器,不保存基座大权重
        model.save_pretrained(os.path.join(SAVE_LORA_PATH, f"epoch_{epoch}"))
    print("LoRA微调完成,适配器权重已保存")

# ===================== 7. 微调后推理函数 =====================
def infer(prompt_text):
    model.set_train(False)
    inputs = tokenizer(f"<s>[INST]{prompt_text}[/INST]", return_tensors="ms")
    output = model.generate(**inputs, max_new_tokens=512, temperature=0.7, top_p=0.9)
    return tokenizer.decode(output[0], skip_special_tokens=True)

if __name__ == "__main__":
    train()
    res = infer("解释昇腾NPU的CANN栈")
    print(res)

代码解析与架构说明

Mistral 在昇思中适配了 CANN 算子栈,滑动窗口注意力sliding_window=4096,长序列时只计算局部窗口,显著降低 KV Cache 显存占用。代码中 Graph 模式编译,昇腾自动做算子融合、通信优化;数据并行模式下 msrun 拉起多卡,梯度 AllReduce 聚合,适合 910B/910C 集群。LoRA 只在自注意力 QKV 投影层插入低秩矩阵,基座权重冻结,训练只更新少量适配器参数,单张昇腾 910B 32G 显存可跑 Mistral-7B LoRA,无需全参所需上百 GB 显存。

数据集部分采用 Alpaca 格式指令数据,遵循 Mistral 官方[INST]指令模板,必须严格匹配,否则模型输出会乱序。数据转为 MindSpore GeneratorDataset,也可预处理为 MindRecord 格式,IO 吞吐更高,适合百万级样本。学习率使用余弦退火预热,LoRA 微调学习率一般 1e-4~3e-4,远大于基座预训练学习率,避免破坏原有基础能力。

分布式启动与权重合并

多卡分布式启动脚本 msrun,单机 8 卡示例:

msrun --worker_num=8 --local_worker_num=8 --master_port=8118 python mistral_lora_sft.py

微调完成后,可将 LoRA 适配器合并入 Mistral 基座,导出单份 ckpt 用于 MindIE 推理部署,合并后模型可直接转 OM 离线模型,在昇腾卡上高吞吐部署。

踩坑复盘与调优方向

第一,Mistral 的 Rotary 位置编码和滑动窗口注意力在昇腾老版本 CANN 存在精度异常,建议使用 CANN 7.0.RC1 以上版本,MindSpore>=2.5;第二,输入序列长度超过滑动窗口 4096 会出现注意力失效,超长场景需要修改配置开启全局注意力;第三,昇腾 FP16 容易溢出,优先使用 bfloat16 精度;第四,LoRA target_modules 不要包含 MLP 层,会大幅增加参数量,显存压力上升;第五,数据集必须做严格清洗,指令模板不匹配极易出现过拟合、输出拒绝。

落地方案

这套方案面向昇腾集群 Mistral 指令微调,落地路径为:原始指令数据清洗→MindRecord 数据构建→LoRA SFT 微调→验证集评估→LoRA 权重合并→MindIE 推理服务化。若算力充足、数据集规模十万级以上,可尝试全参微调,开启模型并行切分 Mistral 权重;小样本业务场景优先 LoRA,迭代速度快,成本可控。Mistral 本身推理速度快,微调后适合知识库问答、智能客服、结构化输出等边缘与中心侧昇腾业务场景。

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐