如何实现LLM训练脚本编写与平台启动配置

作为一名有着编程和项目经验的开发者，我也参与到了类似MCP架构的大模型落地项目中，其中涉及到在云平台训练模型的工作。在这个过程中，如何编写训练脚本以及进行平台启动配置是关键环节，下面我就把自己的经验分享给大家。

一、背景与训练参数经验

在云平台训练模型时，我主要负责模型微调（fine – tuning）的工作。在微调过程中，有两个超参数频繁出现，那就是学习率（learning rate）和epoch数量。一开始，我直接照搬全量训练的参数，结果不是训练过程出现问题，就是模型效果不稳定。

就拿用LLaMA2预训练一个通用大语言模型来说，这个模型本身已经具备很多知识，但我们想让它更贴合业务，比如回答问题时带上品牌口吻、遵循业务流程。然而，大模型的参数动辄几十亿，哪怕只对参数做一点改动，输出结果都会有很大变化。经过多次尝试和总结，我发现微调时，把学习率压低到1e - 5 ~ 5e - 5比较合适。

另外，大模型预训练时使用的语料量通常以TB为单位，而微调语料可能只有几百MB。如果训练的epoch太多，模型就会过度记住训练集里的每一句话，这样一来，它的泛化能力反而会下降。所以，一般把epoch控制在2 ~ 4之间就足够了。

这里给大家推荐一个微调训练参数模板，搭配transformers.Trainer使用，在NPU场景下配合MindSpeed也能很好地适配，代码如下：

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

二、配置训练脚本与平台任务

（一）编写训练脚本

接下来讲讲如何以PyTorch + HuggingFace Transformers为基础，编写一个能在常见平台（支持NPU/GPU）上运行的最小可运行微调训练脚本train.py 。代码如下：

# 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()

在这段代码中，parse_args函数用于解析命令行参数，方便我们在运行脚本时灵活配置训练参数。main函数则负责整个训练流程，包括设备选择、模型和分词器的加载、数据集的处理、训练参数的设置以及模型的训练。

（二）平台任务配置

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

/workspace ├── train.py ├── config.json ├── start_train.sh # 启动脚本 └── /model # 预训练模型权重（可选） /data └── train.jsonl # 你的微调数据 /output └── ... # 模型训练输出

其中，start_train.sh是启动脚本，内容如下：

#!/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目录下。注意，还需要给这个脚本赋予可执行权限，命令如下：

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小时，视数据大小定
这里设置的挂载路径在容器中会变成文件系统路径，所以`train.py`中使用绝对路径来指定数据和输出目录。任务启动后，我们可以查看日志来了解训练进度，例如：

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

三、微调优化方案——LoRA

直接微调大模型不仅会占用大量显存、成本高昂，而且训练时间很长。有没有更轻量级的方案呢？答案是肯定的，那就是LoRA。

LoRA的原理就是在不改动原模型参数的基础上，插入一个“小模块”，只训练这个小模块就可以了。比如说对于一个linear层，LoRA会进行这样的操作：

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

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

要使用LoRA，首先需要安装HuggingFace官方的PEFT工具包，命令如下：

pip install peft

然后在前面的train.py中加入以下代码：

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会自动接管参数训练。

四、模型验证

如果我们有validation数据（比如业务问答对），就可以对模型进行验证。代码如下：

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通常下降得会慢一点，这也是正常现象。

五、模型保存与合并

训练完成后，我们需要保存模型。使用trainer的save_model方法就可以实现，代码如下：

trainer.save_model("/output/lora_model")

如果想要把模型拿去部署（比如转成ONNX、TorchScript、推理服务等），则需要合并LoRA权重，代码如下：

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格式模型了，可以用于后续的部署工作。