llm 训练脚本怎么写？平台启动怎么配？一文搞定

1、背景

咱也是好起来了，作为兼具编程经验与项目经验的选手，开始着手类似 MCP 架构的大模型融入公司业务的落地项目。

其中，也涉及在云平台训练模型：

我在微调（fine-tuning）过程中，发现有两个高频出现的超参数：学习率（learning rate）和 epoch 数量

一开始直接从全量训练的经验里拷参数，结果不是训练崩了就是效果不稳。

比如用 LLaMA2 预训练一个通用大语言模型，它知道很多东西，想让它更像业务助手，回答时带品牌口吻、遵守业务流程，在大模型结构里，参数动辄几十亿，哪怕只改一点，也会有很强的输出偏移。

经验总结：所以微调时，建议把学习率压低到 1e-5 ~ 5e-5。

大模型的预训练语料量以 TB 计，而微调语料可能只有几百 MB，训练太多 epoch，模型会记住训练集里的每一句话，泛化能力反而下降。

经验总结：通常把 epoch 控制在 2 ~ 4 之间就足够。

📦 推荐微调训练参数模板，稳定实用：

css 复制代码

--learning_rate 2e-5 \
--num_train_epochs 3 \
--per_device_train_batch_size 4 \
--save_steps 500 \
--save_total_limit 2 \
--logging_steps 100

搭配 transformers.Trainer 使用即可，NPU 场景下配合 MindSpeed 使用同样适配。

2、配置

接下来，介绍如何写一个train.py 跑起来，并且在平台上配置任务：

以 PyTorch + HuggingFace Transformers 为基础，适配常见平台（支持 NPU/GPU 的），写一个最小可运行的微调训练脚本：

ini 复制代码

# train.py
import os
import argparse
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_path", type=str, default="/data/train.jsonl")
    parser.add_argument("--output_dir", type=str, default="/output")
    parser.add_argument("--model_name_or_path", type=str, default="facebook/opt-1.3b")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--device", type=str, default="npu")
    parser.add_argument("--device_id", type=int, default=0)
    return parser.parse_args()

def main():
    args = parse_args()

    if args.device == "npu":
        os.environ["ASCEND_DEVICE_ID"] = str(args.device_id)
        torch.npu.set_device(args.device_id)
        device = torch.device("npu")
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
    model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path).to(device)

    dataset = load_dataset("json", data_files=args.data_path, split="train")
    dataset = dataset.map(lambda ex: tokenizer(ex["text"], truncation=True, padding="max_length"), batched=True)

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        per_device_train_batch_size=args.batch_size,
        num_train_epochs=args.epochs,
        learning_rate=args.lr,
        save_steps=500,
        save_total_limit=2,
        logging_dir=os.path.join(args.output_dir, "logs"),
        logging_steps=100,
        report_to="none"
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    )

    trainer.train()

if __name__ == "__main__":
    main()

平台训练时，一般你会上传一个 zip 包，或者挂载一个文件夹，结构如下：

bash 复制代码

/workspace
├── train.py
├── config.json
├── start_train.sh     # 启动脚本
└── /model             # 预训练模型权重（可选）

/data
└── train.jsonl        # 你的微调数据

/output
└── ...                # 模型训练输出

然后是启动脚本：

bash 复制代码

#!/bin/bash
echo "[INFO] Starting fine-tuning..."

python3 /workspace/train.py \
  --data_path /data/train.jsonl \
  --output_dir /output \
  --model_name_or_path facebook/opt-1.3b \
  --device npu \
  --device_id 0 \
  --epochs 3 \
  --batch_size 4 \
  --lr 2e-5

保存为 start_train.sh，放在 /workspace 目录下。

别忘了赋权：

bash 复制代码

chmod +x start_train.sh

🚀 创建任务时注意以下设置：

配置项	示例
镜像	mindspeed-llm-develop:0.8.0-npu
启动方式	命令启动
启动命令	`sh /workspace/start_train.sh`
挂载路径	`/workspace`（代码）、`/data`（数据）、`/output`（结果）
资源规格	Ascend NPU（2 卡/4 卡）、或 GPU（A100）
超时时间	≥2 小时，视数据大小定

这里设置的挂载路径会在容器中变成文件系统路径，比如 /data/train.jsonl，这也是为什么 train.py 要用绝对路径。

然后，查看日志：

vbnet 复制代码

[INFO] Using device: npu
Epoch 1/3
Step 500 - loss: 2.31
...
Saving model checkpoint to /output/checkpoint-500

OK，老九它没毛病~

3、微调

直接微调一个大模型，不但显存吃紧、资源贵，而且训练时间久。有没有轻一点的方案？

有的 ------ LoRA。

不动原模型参数，只插入一个"小模块"，训练这部分就好。

比如对于一个 linear 层，LoRA 会插入：

ini 复制代码

W = W_0 + ΔW
ΔW = A × B

A 和 B 是低秩矩阵，参数量远小于原始矩阵。整个过程不需要修改原模型结构；训练速度快、显存占用低、微调稳定性更高。

如何使?

首先需要安装 HuggingFace 官方的 PEFT 工具包：

复制代码

pip install peft

然后在前面的 train.py 中加入：

ini 复制代码

from peft import get_peft_model, LoraConfig, TaskType

# 加载基础模型
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)

# 加 LoRA
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"]  # 这取决于你用的模型结构
)
model = get_peft_model(model, peft_config)

训练逻辑、Trainer 用法基本不变，LoRA 会自动接管参数。

4、验证

如果你有 validation 数据（比如业务问答对）：

ini 复制代码

val_dataset = load_dataset("json", data_files="val.jsonl", split="train")
val_dataset = val_dataset.map(tokenizer, batched=True)

eval_result = trainer.evaluate(val_dataset)
print(eval_result)

就可以得到：

eval_loss
perplexity

也可以加 BLEU、ROUGE、EM 等指标（可扩展）。。

通常来说还可以观察：

每个 save_steps 或每个 epoch， loss 是否收敛；一般 NLP 微调 loss 降到 2.x ~ 1.x 属于正常。LoRA loss 通常会下降慢一点，属于正常现象。

5、保存

模型保存：

arduino 复制代码

trainer.save_model("/output/lora_model")

想拿去部署（比如转成 ONNX、TorchScript、推理服务等），需要合并 LoRA 权重：

ini 复制代码

from peft import PeftModel
from transformers import AutoModelForCausalLM

base_model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b")
model = PeftModel.from_pretrained(base_model, "/output/lora_model")

# 合并权重
model = model.merge_and_unload()
model.save_pretrained("/output/full_model")

现在 /output/full_model 目录下就是一个完整 HuggingFace 格式模型了~

6、备忘

踩坑提示

场景	建议
想做快速微调	LoRA 是首选，速度快、精度可控
想跑大模型	推荐先跑 7B、13B，q_proj/v_proj 是高收益目标
推理部署	用 merge_and_unload 合并权重后再部署
多卡训练	LoRA 兼容 torchrun / deepspeed，可跑分布式
只做推理	可加载合并后的模型，直接 `.generate()` 生成文本