一、摘要

基于昇思 MindSpore 在单张昇腾 NPU(310P/910B)完成大模型微调 + 推理是轻量化落地常用方案。单卡流程包含:环境准备、权重加载、数据集构建、LoRA 微调、模型保存、离线推理全链路。相比于全参数微调,LoRA 低秩适配极大降低单卡显存压力,适合行业模型轻量化二次开发。

本文使用 MindSpore + MindFormers,以 Decoder-only 大模型为例,完整实现单卡 LoRA 微调、权重合并、本地推理整套自助流程,适配昇腾 CANN 环境。

运行环境:openEuler、CANN、MindSpore2.3、MindFormers、昇腾 NPU 单卡。

二、环境初始化代码(NPU 设备配置)

# env_init.py
import os
import mindspore as ms
from mindspore import context
def init_npu_env():
    # 指定昇腾NPU卡号
    os.environ["DEVICE_ID"] = "0"
    # MindSpore昇腾后端配置
    context.set_context(
        mode=context.GRAPH_MODE,
        device_target="Ascend",
        device_id=int(os.environ["DEVICE_ID"]),
        save_graphs=False
    )
    # 显存优化策略,单卡微调防OOM
    ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.STAND_ALONE)
    ms.set_context(max_call_depth=2000)
    print("昇腾NPU单卡环境初始化完成")
if __name__ == "__main__":
    init_npu_env()

三、训练数据集构建代码

采用指令微调标准 JSON 数据集,封装 MindSpore Dataset 迭代器

# dataset.py
import json
import mindspore.dataset as ds
from mindformers import PromptTokenizer
class SFTDataSet:
    def __init__(self, data_path, tokenizer_path, seq_len=512):
        self.seq_len = seq_len
        self.tokenizer = PromptTokenizer(tokenizer_path)
        with open(data_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"###指令:{sample['instruction']}\n###回答:{sample['output']}"
        token = self.tokenizer(
            prompt,
            padding="max_length",
            truncation=True,
            max_length=self.seq_len
        )
        input_ids = token["input_ids"]
        attention_mask = token["attention_mask"]
        labels = input_ids.copy()
        return input_ids, attention_mask, labels
def create_sft_dataloader(data_path, tokenizer_path, batch_size=2):
    dataset_generator = SFTDataSet(data_path, tokenizer_path)
    dataset = ds.GeneratorDataset(
        dataset_generator,
        column_names=["input_ids", "attention_mask", "labels"],
        shuffle=True
    )
    dataset = dataset.batch(batch_size, drop_remainder=True)
    return dataset

数据集 data.json 格式参考:

[
    {"instruction":"介绍昇思MindSpore","output":"MindSpore是华为开源全场景AI框架"}
]

四、单卡 LoRA 微调主训练代码

# train_lora_single_card.py
from env_init import init_npu_env
from dataset import create_sft_dataloader
import mindspore as ms
from mindformers import AutoModel, AutoConfig, LoRAConfig
from mindspore.nn import AdamWeightDecay
from mindspore.train import Model
from mindspore.train.callback import SaveCheckpoint, CheckpointConfig
init_npu_env()
# 1. LoRA配置
lora_config = LoRAConfig(
    lora_rank=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none"
)
# 2. 加载基础大模型
model_config = AutoConfig.from_pretrained("./base_model")
model_config.checkpoint_name_or_path = "./base_model/ckpt"
network = AutoModel.from_config(model_config)
# 注入LoRA层,冻结主干权重
network.freeze()
network.add_lora(lora_config)
# 3. 数据集
train_dataset = create_sft_dataloader(
    data_path="./data.json",
    tokenizer_path="./base_model",
    batch_size=2
)
# 4. 优化器与训练封装
lr = ms.nn.exponential_decay_lr(
    learning_rate=2e-4,
    decay_rate=0.9,
    total_step=1000,
    step_per_epoch=len(train_dataset),
    decay_epoch=1
)
optimizer = AdamWeightDecay(network.trainable_params(), learning_rate=lr)
# 损失函数
loss_fn = ms.nn.CrossEntropyLoss(ignore_index=0)
train_net = ms.nn.WithLossCell(network, loss_fn)
train_net = ms.nn.TrainOneStepCell(train_net, optimizer)
# 5. 训练循环与保存
ckpt_cfg = CheckpointConfig(save_checkpoint_steps=50, keep_checkpoint_max=5)
save_cb = SaveCheckpoint(config=ckpt_cfg, directory="./lora_ckpt")
epochs = 3
for epoch in range(epochs):
    for batch_data in train_dataset.create_tuple_iterator():
        input_ids, attn_mask, labels = batch_data
        loss = train_net(input_ids, attn_mask, labels)
        print(f"epoch:{epoch}, loss:{loss.asnumpy():.4f}")
print("单卡LoRA微调完成,LoRA权重已保存")

五、微调后推理代码(单卡本地推理)

# infer.py
from env_init import init_npu_env
from mindformers import AutoModel, AutoTokenizer
init_npu_env()
tokenizer = AutoTokenizer.from_pretrained("./base_model")
model = AutoModel.from_pretrained("./base_model")
# 加载训练得到的LoRA权重
model.load_lora_ckpt("./lora_ckpt/lora_rank_8.ckpt")
def predict(prompt_text):
    inputs = tokenizer(f"###指令:{prompt_text}\n###回答:", return_tensors="ms")
    output = model.generate(
        **inputs,
        max_length=256,
        temperature=0.7,
        top_p=0.9
    )
    result = tokenizer.decode(output[0], skip_special_tokens=True)
    return result
if __name__ == "__main__":
    res = predict("简单介绍MindSpore单卡微调流程")
    print("模型输出:\n", res)

六、启动脚本 shell

# run_single_card.sh
#!/bin/bash
export ASCEND_TOOLKIT_PATH=/usr/local/Ascend/ascend-toolkit/latest
source ${ASCEND_TOOLKIT_PATH}/bin/set_env.sh
export DEVICE_ID=0
python3 train_lora_single_card.py
执行:bash run_single_card.sh

七、单卡调优关键要点

显存控制:优先 LoRA 替代全参数微调;开启梯度检查点 model_config.use_recompute=True,大幅降低显存占用,避免单卡 OOM;

运行模式:GRAPH_MODE 性能远高于 PYNATIVE_MODE,正式训练统一使用图模式;

数据加载:单卡不要设置过大 batch_size,根据 NPU 显存逐级调试;

权重管理:LoRA 权重体积很小,推理时动态加载,也可执行权重合并导出完整模型用于 ATC 离线转换;

性能观测:使用 npu-smi 观测显存、算力利用率,及时发现数据加载瓶颈。

八、总结

整套自助搭建流程分为环境初始化、数据集构建、LoRA 单卡微调、权重加载推理四大环节。MindSpore 搭配 MindFormers 封装了大模型通用接口,降低了昇腾硬件上大模型开发门槛。单卡方案无需分布式集群,适合模型验证、小样本行业微调、原型验证场景。

开发流程标准统一,可快速迁移至 310P、910 系列昇腾设备。在工程实践中,可增加早停策略、验证集评估、日志保存,形成完整可自动化运行的单卡微调推理流水线。

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐