Unsloth自动化脚本编写：批量处理训练任务

Unsloth自动化脚本编写：批量处理训练任务

在实际模型微调工作中，我们常常需要反复执行相似但参数各异的训练任务——比如对同一基座模型在不同数据集上做LoRA微调、尝试多种学习率组合、对比不同量化精度下的推理效果，或是为多个业务场景分别定制专属模型。手动逐个运行Jupyter Notebook或Python脚本不仅效率低下，还极易因环境切换、路径错误、参数遗漏导致结果不可复现。此时，一套可配置、可复用、可追踪的自动化脚本体系，就不再是“锦上添花”，而是工程落地的刚需。

本文聚焦于Unsloth框架下的批量训练自动化实践，不讲抽象理论，不堆砌概念，只提供一套经过真实项目验证的、开箱即用的脚本方案。你将学会如何用纯Python+Shell构建一个轻量级但功能完整的训练调度系统：从环境准备、参数注入、任务编排，到日志归档与结果汇总。所有代码均可直接复制运行，适配CSDN星图镜像广场中预置的unsloth镜像环境。

1. 为什么需要自动化？——来自炼丹现场的真实痛点

在使用Unsloth完成DeepSeek-R1-Distill-Qwen-1.5B模型微调的过程中，我们遇到了三类高频问题：

• 重复劳动多：每次换一个数据集，就要复制粘贴整段SFTTrainer配置；调整learning_rate或batch_size，得手动改七八处代码；保存模型时还要反复确认save_method和路径名。
• 状态难追溯：一次实验跑完，训练日志散落在终端、SwanLab链接、本地outputs/目录里；想对比两次实验的loss曲线，得手动翻记录、截图、再拼图。
• 容错能力弱：某次训练因显存不足中断，重启后得重新加载模型、重设tokenizer、重新初始化trainer——而这些步骤在Notebook里是线性执行的，无法跳过已完成部分。

这些问题的本质，是把“工程任务”当成了“一次性实验”。而自动化不是为了炫技，而是把经验沉淀为可执行的规则，让模型工程师真正聚焦在数据质量、提示设计、效果评估这些高价值环节上。

自动化不是消灭人工，而是把人从机械操作中解放出来，去解决更复杂的问题。

2. 自动化脚本设计原则：轻量、清晰、可扩展

我们不追求大而全的平台化方案（如MLflow/Kubeflow），而是采用“脚手架式”设计，核心遵循三条原则：

• 零依赖：仅使用Python标准库、unsloth、transformers、datasets、swanlab等已预装包，不引入额外构建工具。
• 配置驱动：所有可变参数（模型路径、数据集、超参、保存策略）统一收口到YAML配置文件，代码逻辑与参数解耦。
• 原子化任务：每个训练任务封装为独立Python模块，支持单独调试、并行执行、失败重试。

整个系统由三部分构成：

• config/：存放YAML格式的训练配置
• scripts/：核心调度脚本与任务模板
• utils/：通用工具函数（日志管理、路径处理、SwanLab初始化）

这种结构确保了：新增一个训练任务，只需新增一个配置文件 + 复制一份模板脚本；修改全局行为（如统一启用梯度检查点），只需改一处utils/函数。

3. 核心脚本实现：从单任务到批量调度

3.1 配置文件标准化：config/train_lora.yaml

我们以LoRA微调为例，定义标准化配置。注意：所有字段均为必填，避免运行时缺省值引发歧义。

# config/train_lora.yaml experiment_name: "deepseek-r1-qwen-1.5b-lora-motor-domain" description: "电机选型领域LoRA微调，基于cleaned_dataset_v4.0.0" # 模型与分词器 model_path: "./deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" tokenizer_path: "./deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" load_in_4bit: true max_seq_length: 2048 # 数据集 dataset_path: "cleaned_dataset_v4.0.0" dataset_text_field: "text" # LoRA参数 lora_r: 16 lora_alpha: 16 lora_dropout: 0.0 target_modules: - "q_proj" - "k_proj" - "v_proj" - "o_proj" - "gate_proj" - "up_proj" - "down_proj" # 训练参数 per_device_train_batch_size: 2 gradient_accumulation_steps: 4 warmup_steps: 5 max_steps: 30 learning_rate: 2e-4 optim: "adamw_8bit" weight_decay: 0.01 lr_scheduler_type: "linear" seed: 3407 # SwanLab集成 swanlab_project: "motor-domain-finetuning" swanlab_experiment_name: "{{ experiment_name }}" swanlab_description: "{{ description }}" # 输出路径（自动补全时间戳） output_dir: "outputs/{{ experiment_name }}_{{ timestamp }}"

关键设计点：

• 使用{{ timestamp }}占位符，由脚本自动替换为20250715_142305格式，杜绝路径冲突；
• swanlab_experiment_name支持Jinja2风格模板，复用顶层配置，避免硬编码；
• 所有布尔值明确写为true/false，不依赖YAML隐式转换。

3.2 单任务执行脚本：scripts/run_lora_task.py

该脚本是自动化系统的“最小可执行单元”，接收配置路径作为唯一参数，完成端到端训练。

# scripts/run_lora_task.py import os import sys import time import yaml from datetime import datetime from pathlib import Path # 添加项目根目录到Python路径 ROOT_DIR = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT_DIR)) from unsloth import FastLanguageModel from transformers import TrainingArguments from trl import SFTTrainer, SFTConfig from datasets import load_from_disk from swanlab.integration.transformers import SwanLabCallback import swanlab def load_config(config_path): """加载并渲染YAML配置""" with open(config_path, 'r', encoding='utf-8') as f: config = yaml.safe_load(f) # 渲染时间戳 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") for key, value in config.items(): if isinstance(value, str): config[key] = value.replace("{{ timestamp }}", timestamp) config[key] = config[key].replace("{{ experiment_name }}", config.get("experiment_name", "")) config[key] = config[key].replace("{{ description }}", config.get("description", "")) return config def init_swanlab(config): """初始化SwanLab，支持API Key环境变量或配置文件""" api_key = os.getenv("SWANLAB_API_KEY") or config.get("swanlab_api_key") if api_key: swanlab.login(api_key=api_key, save=True) os.environ["SWANLAB_PROJECT"] = config["swanlab_project"] return SwanLabCallback( project=config["swanlab_project"], experiment_name=config["swanlab_experiment_name"], description=config["swanlab_description"], config={ "framework": "Unsloth+TRL", "config_file": config_path, **config # 将全部配置传入，便于SwanLab分析 } ) def main(config_path): print(f"[INFO] 开始执行训练任务：{config_path}") config = load_config(config_path) # 1. 加载模型与分词器 print("[STEP 1] 加载基座模型...") model, tokenizer = FastLanguageModel.from_pretrained( model_name=config["model_path"], max_seq_length=config["max_seq_length"], dtype=None, load_in_4bit=config["load_in_4bit"], ) # 2. 注入LoRA print("[STEP 2] 注入LoRA适配器...") model = FastLanguageModel.get_peft_model( model, r=config["lora_r"], lora_alpha=config["lora_alpha"], lora_dropout=config["lora_dropout"], target_modules=config["target_modules"], bias="none", use_gradient_checkpointing="unsloth", random_state=config["seed"], ) # 3. 加载数据集 print("[STEP 3] 加载训练数据集...") train_dataset = load_from_disk(config["dataset_path"]) # 4. 初始化SwanLab回调 print("[STEP 4] 初始化SwanLab监控...") swanlab_callback = init_swanlab(config) # 5. 构建训练器 print("[STEP 5] 构建SFTTrainer...") trainer = SFTTrainer( model=model, tokenizer=tokenizer, train_dataset=train_dataset, dataset_text_field=config["dataset_text_field"], args=SFTConfig( per_device_train_batch_size=config["per_device_train_batch_size"], gradient_accumulation_steps=config["gradient_accumulation_steps"], warmup_steps=config["warmup_steps"], max_steps=config["max_steps"], learning_rate=config["learning_rate"], optim=config["optim"], weight_decay=config["weight_decay"], lr_scheduler_type=config["lr_scheduler_type"], seed=config["seed"], logging_steps=1, output_dir=config["output_dir"], report_to="none", ), callbacks=[swanlab_callback], ) # 6. 执行训练 print("[STEP 6] 启动训练...") start_time = time.time() trainer_stats = trainer.train() end_time = time.time() # 7. 保存最终模型 print("[STEP 7] 保存微调后模型...") os.makedirs(config["output_dir"], exist_ok=True) model.save_pretrained_merged( save_directory=f"{config['output_dir']}/merged_fp16", tokenizer=tokenizer, save_method="merged_16bit" ) # 8. 记录元信息 meta_info = { "config_file": config_path, "start_time": datetime.fromtimestamp(start_time).isoformat(), "end_time": datetime.fromtimestamp(end_time).isoformat(), "duration_seconds": round(end_time - start_time, 2), "final_loss": trainer_stats.training_loss, "output_dir": config["output_dir"] } meta_path = Path(config["output_dir"]) / "meta.json" with open(meta_path, 'w', encoding='utf-8') as f: json.dump(meta_info, f, indent=2, ensure_ascii=False) print(f"[SUCCESS] 任务完成！结果保存至：{config['output_dir']}") if __name__ == "__main__": if len(sys.argv) != 2: print("用法：python scripts/run_lora_task.py <配置文件路径>") sys.exit(1) import json # 仅在此处导入，避免全局污染 main(sys.argv[1])

该脚本的关键特性：

• 强健性：每步打印明确日志，失败时抛出异常并终止，避免静默错误；
• 可调试性：支持直接运行python scripts/run_lora_task.py config/train_lora.yaml，无需任何前置构建；
• 可审计性：生成meta.json记录完整执行上下文，包括耗时、损失值、输出路径。

3.3 批量调度器：scripts/batch_runner.py

当需要同时启动多个实验时，batch_runner.py负责读取配置列表、并发执行、汇总结果。

# scripts/batch_runner.py import os import sys import subprocess import time from pathlib import Path from concurrent.futures import ThreadPoolExecutor, as_completed def run_single_task(config_path, timeout=3600): """执行单个训练任务，带超时控制""" cmd = [sys.executable, "scripts/run_lora_task.py", str(config_path)] try: result = subprocess.run( cmd, capture_output=True, text=True, timeout=timeout, cwd=str(Path(__file__).resolve().parent.parent) ) return { "config": str(config_path), "success": result.returncode == 0, "stdout": result.stdout[-500:], # 截取末尾500字符供快速诊断 "stderr": result.stderr, "returncode": result.returncode } except subprocess.TimeoutExpired: return { "config": str(config_path), "success": False, "stdout": "", "stderr": f"Timeout after {timeout}s", "returncode": -1 } def main(config_dir, max_workers=2): """批量运行指定目录下所有YAML配置""" config_files = list(Path(config_dir).glob("*.yaml")) if not config_files: print(f"[ERROR] 在 {config_dir} 下未找到YAML配置文件") return print(f"[INFO] 发现 {len(config_files)} 个配置文件，将使用 {max_workers} 个线程并发执行") results = [] with ThreadPoolExecutor(max_workers=max_workers) as executor: # 提交所有任务 future_to_config = { executor.submit(run_single_task, cf): cf for cf in config_files } # 收集结果 for future in as_completed(future_to_config): result = future.result() results.append(result) status = " 成功" if result["success"] else "❌ 失败" print(f"[{status}] {result['config']}") # 汇总报告 success_count = sum(1 for r in results if r["success"]) print(f"\n[SUMMARY] 批量执行完成：{success_count}/{len(results)} 个任务成功") if success_count < len(results): print("\n--- 失败详情 ---") for r in results: if not r["success"]: print(f"❌ {r['config']}") if r["stderr"]: print(f" 错误：{r['stderr'].strip()[:100]}...") if __name__ == "__main__": if len(sys.argv) < 2: print("用法：python scripts/batch_runner.py <配置目录路径> [线程数，默认2]") sys.exit(1) config_dir = sys.argv[1] max_workers = int(sys.argv[2]) if len(sys.argv) > 2 else 2 main(config_dir, max_workers)

使用方式极其简单：

# 运行config/下所有YAML文件（2线程并发） python scripts/batch_runner.py config/ # 指定4线程加速 python scripts/batch_runner.py config/ 4

优势在于：

• 资源可控：通过max_workers限制GPU占用，避免OOM；
• 失败隔离：单个任务失败不影响其他任务；
• 结果透明：终端实时显示每个任务状态，并生成汇总报告。

4. 实战案例：一键启动电机领域三阶段训练流水线

结合参考博文中的“继续预训练→指令微调→全量微调”流程，我们构建一个端到端的流水线配置。在config/pipeline_motor_domain.yaml中定义：

stages: - name: "cpt" type: "continual_pretraining" config_file: "config/cpt_motor.yaml" depends_on: [] # 无依赖 - name: "sft" type: "instruction_finetuning" config_file: "config/sft_motor.yaml" depends_on: ["cpt"] # 依赖上一阶段输出 - name: "fft" type: "full_finetuning" config_file: "config/fft_motor.yaml" depends_on: ["sft"] # 依赖指令微调后的模型 output_root: "pipelines/motor_domain_{{ timestamp }}"

配套的scripts/run_pipeline.py可解析此文件，按依赖顺序执行各阶段，并自动将前一阶段的merged_fp16路径注入下一阶段的model_path。这已超出本文范围，但其思想一脉相承：用声明式配置替代命令式脚本，让复杂流程变得可读、可维护、可复现。

5. 工程化建议：让自动化真正落地

光有脚本还不够，还需配套实践才能发挥最大价值：

5.1 环境一致性保障

在unsloth镜像中，我们已预装所有依赖，但仍需注意：

• 统一使用conda activate unsloth_env激活环境，避免pip与conda混用；
• 在脚本开头添加版本校验：

  import unsloth assert unsloth.__version__ >= "2025.6.0", "请升级Unsloth至2025.6.0+"

5.2 日志与可观测性增强

• 将trainer.train()的日志重定向到文件：trainer.train(log_level="info", logging_dir=f"{config['output_dir']}/logs")；
• 在SwanLab中固定experiment_name前缀，如motor-cpt-20250715，便于在Web界面按前缀筛选；
• 使用watch -n 10 nvidia-smi在终端常驻监控GPU内存变化。

5.3 故障快速定位技巧

当训练失败时，按以下顺序排查：

1. 看meta.json：确认是否卡在某一步（如start_time有值但无end_time，说明训练未结束）；
2. 查stdout末尾：batch_runner.py已截取最后500字符，通常包含关键报错；
3. 翻outputs/xxx/logs/：查看详细TensorBoard日志；
4. 进WebShell执行conda activate unsloth_env && python -m unsloth：验证环境基础功能。

5.4 安全与协作规范

• 配置文件中禁止硬编码API Key，一律通过环境变量SWANLAB_API_KEY注入；
• config/目录纳入Git版本控制，但outputs/、pipelines/等产出目录加入.gitignore；
• 团队共享配置时，使用#注释说明每个参数的实际影响（如# learning_rate: 2e-4适用于小数据集快速验证，大数据集建议降至2e-5）。

6. 总结：自动化是模型工程师的第二大脑

回看整个自动化体系，它没有发明新算法，也没有提升单次训练的性能，但它带来了三重确定性：

• 过程确定性：同样的配置，在任何机器上运行，必然得到相同的结果路径和SwanLab链接；
• 时间确定性：原本需要2小时的手动操作，现在python scripts/batch_runner.py config/一条命令，喝杯咖啡的时间就完成了10个实验；
• 认知确定性：工程师不再需要记忆“上次那个学习率是多少”，所有决策都沉淀在YAML文件里，新人上手第一天就能复现全部历史实验。

这才是技术博客最该传递的价值：不炫耀技巧的复杂，而展示如何用最朴素的工具，解决最真实的工程问题。当你把run_lora_task.py第一次成功运行起来，看到终端滚动出[SUCCESS]时，你就已经迈出了从“炼丹学徒”到“炼丹工程师”的关键一步。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。