一、概述

在昇腾生态中，MindSpore Transformers 是官方支持大语言模型（LLM）开发、预训练、微调的核心框架，全面适配 LLaMA、Qwen、GPT、Bloom 等主流 decoder-only 架构，可充分利用昇腾 NPU 多核并行、高带宽特性实现高效训练。LLM 预训练对算力、显存、通信效率要求极高，传统训练方案易出现速度慢、显存溢出、利用率低等问题。MindSpore 依托静态图、自动并行、混合精度、算子融合等原生能力，让大模型预训练在单机多卡、多机多卡环境下实现高吞吐、低显存、稳定收敛。

二、MindSpore Transformers LLM 高效训练核心原理

高效训练的本质是让昇腾 NPU 持续满负荷运算，减少等待、拷贝、冗余计算。MindSpore 从四大维度实现 LLM 预训练加速：

2.1 自动并行策略

MindSpore 支持数据并行、模型并行、流水线并行、序列并行，自动对 LLM 权重、注意力、FFN 层进行切分，无需手动改造代码，即可支持 7B～13B 大模型单机多卡训练，多机扩展线性比可达 0.9 以上。

2.2 显存优化技术

采用重计算（Recomputation）、FP16/BF16 混合精度、权重累加，将显存占用降低 60% 以上，使小显存 NPU 也能跑大模型。

2.3 静态图与算子融合

MindSpore 静态图（Graph Mode）提前编译计算流，配合Transformer 专属算子融合，减少 Kernel 启动开销与内存拷贝，训练速度相比动态图提升 2～4 倍。

2.4 数据流水线优化

数据加载、预处理、NPU 传输异步执行，实现CPU 数据处理与 NPU 计算无缝衔接，避免 GPU/NPU 空转，显著提升吞吐。

三、LLM 预训练关键优化配置

高效训练必须搭配以下核心配置，是性能与稳定性的基础：

1. 自动混合精度 AMP：使用 FP16 训练，速度翻倍、显存减半；
2. 动态损失缩放：防止 LLM 训练出现 NaN 溢出；
3. 重计算激活：牺牲少量算力，大幅节省显存；
4. AdamW 优化器：LLM 预训练标准优化器；
5. 余弦学习率调度：稳定收敛，避免后期震荡；
6. 分布式并行：自动切分模型，支持多卡扩展。

四、MindSpore Transformers LLM 高效训练完整代码

以下代码基于 MindSpore 2.2 + MindSpore Transformers 1.0，适配昇腾 NPU，支持 LLaMA/GPT 类 decoder-only 模型预训练，可直接运行。

4.1 环境安装

pip install mindspore mindspore-dataset
pip install mindspore-transformers

4.2 训练启动代码（llm_pretrain.py）

import mindspore as ms
import mindspore.dataset as ds
from mindspore import nn
from mindspore.transformers import AutoModelForCausalLM, AutoTokenizer
from mindspore.transformers.optimization import AdamW, get_cosine_schedule_with_warmup

# ===================== 1. 全局配置（高效训练核心）=====================
ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")  # 静态图加速
ms.set_auto_parallel_context(parallel_mode="data_parallel")  # 多卡并行
config = {
    "model_name": "llama-7b",
    "batch_size": 4,
    "seq_len": 1024,
    "lr": 2e-5,
    "num_epochs": 2,
    "mixed_precision": "fp16",  # 混合精度
    "recompute": True,          # 重计算省显存
}

# ===================== 2. 模型与分词器加载 =====================
tokenizer = AutoTokenizer.from_pretrained("llama-7b")
model = AutoModelForCausalLM.from_pretrained(
    "llama-7b",
    mixed_precision=config["mixed_precision"],
    recompute=config["recompute"]
)

# ===================== 3. 数据集构建 =====================
def create_pretrain_dataset(data_path, batch_size, seq_len):
    data = ds.TextFileDataset(data_path)
    def tokenize(text):
        return tokenizer(
            text, max_length=seq_len, padding="max_length", truncation=True
        )
    data = data.map(tokenize)
    data = data.batch(batch_size)
    return data

dataset = create_pretrain_dataset(
    "./pretrain_data.txt", config["batch_size"], config["seq_len"]
)

# ===================== 4. 优化器与学习率 =====================
total_steps = dataset.get_dataset_size() * config["num_epochs"]
lr_schedule = get_cosine_schedule_with_warmup(
    config["lr"], total_steps, warmup_ratio=0.05
)
optimizer = AdamW(model.trainable_params(), lr=lr_schedule)

# ===================== 5. 训练模型封装 =====================
model.set_train()
loss_fn = nn.CrossEntropyLoss()
train_net = nn.TrainOneStepCell(model, optimizer)

# ===================== 6. 高效训练循环 =====================
print("===== LLM 预训练启动 =====")
for epoch in range(config["num_epochs"]):
    loss_sum = 0
    for step, batch in enumerate(dataset):
        input_ids = batch["input_ids"]
        loss = train_net(input_ids, input_ids)  # 自回归训练
        loss_sum += loss.asnumpy()
        
        if step % 10 == 0:
            print(f"Epoch {epoch} | Step {step} | Loss {loss.asnumpy():.4f}")

    print(f"=== Epoch {epoch} 完成，平均 Loss: {loss_sum/(step+1):.4f} ===")

# 保存预训练模型
ms.save_checkpoint(model, "./llm_pretrain_final.ckpt")
print("===== 预训练完成，模型已保存 =====")

4.3 多卡分布式启动命令（昇腾集群必备）

mpirun -n 4 python llm_pretrain.py

五、高效训练性能提升关键技巧

1. 开启 Graph 模式
2. 必须使用 mode=ms.GRAPH_MODE，性能提升 200%+。
3. 启用 FP16 混合精度
4. 昇腾 NPU 对 FP16 做了硬件加速，显存减半、速度翻倍，精度损失可忽略。
5. 开启重计算 Recomputation
6. LLM 预训练显存杀手是激活值，重计算可节省 50%+ 显存。
7. 使用异步数据加载
8. MindSpore Dataset 自动实现多线程预处理，让 NPU 永不等待数据。
9. 开启自动并行
10. 模型大于 7B 时，自动启用模型并行，无需手动切分。

六、训练稳定性与排障

• Loss 出现 NaN：降低学习率、开启损失缩放、调小 batch；
• 显存 OOM：开启重计算、减小 batch、降低序列长度；
• 训练速度慢：确认静态图已开启、并行模式正常；
• Loss 不下降：检查学习率、数据质量、数据集格式。

七、总结

MindSpore Transformers 为 LLM 预训练提供了开箱即用、高效稳定的昇腾原生方案，通过静态图编译、自动并行、混合精度、重计算四大核心技术，实现大模型预训练速度快、显存省、易扩展、稳收敛。整套方案无需复杂改造，只需几行配置即可完成单机 / 多卡训练，完美适配昇腾 NPU 硬件特性。