PyTorch TorchTune 端到端大语言模型工作流实战指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00657/article/details/148506174

PyTorch TorchTune 端到端大语言模型工作流实战指南

torchtune A Native-PyTorch Library for LLM Fine-tuning 项目地址: https://gitcode.com/gh_mirrors/to/torchtune

概述

本文将详细介绍如何使用 PyTorch TorchTune 项目完成大语言模型(LLM)的端到端工作流程，包括模型微调、评估、量化和生成等关键环节。TorchTune 是一个专注于大语言模型微调和推理的 PyTorch 库，提供了简单易用的工具链。

准备工作

在开始之前，请确保：

已熟悉 TorchTune 的基本概念和架构
已完成 TorchTune 的安装
了解配置文件(checkpoint)和检查点的基本概念

模型微调

下载预训练模型

首先我们需要获取基础模型。TorchTune 提供了便捷的命令行工具来下载 Hugging Face Hub 上的模型：

tune download meta-llama/Llama-3.2-3B-Instruct --ignore-patterns "original/consolidated.00.pth"

这条命令会下载 Llama3.2 3B Instruct 模型并保存到本地文件系统。

LoRA 微调

对于大多数开发者来说，使用 LoRA (Low-Rank Adaptation)进行参数高效微调是最实用的选择。LoRA 通过冻结基础模型并添加少量可训练参数，显著降低了内存需求。

使用 TorchTune，你可以在 RTX 3090/4090 这样的消费级 GPU 上(16GB 显存)完成 Llama-3.2-3B-Instruct 模型的 bfloat16 精度微调。

查看可用的配置：

tune ls

我们将使用单设备 LoRA 微调配方：

tune run lora_finetune_single_device --config llama3_2/3B_lora_single_device

典型配置参数：

batch_size=4
dtype=bfloat16
每个 epoch 训练时间约 2-3 小时
峰值内存使用约 16GB

输出文件解析

训练完成后，输出目录包含以下重要文件：

adapter_model.safetensors - LoRA 适配器权重
model-{}-of-{}.safetensors - 完整模型权重(如果设置了保存完整模型)
adapter_config.json - Hugging Face PEFT 加载适配器所需的配置
model.safetensors.index.json - Hugging Face from_pretrained() 使用的索引文件

模型评估

使用 EleutherAI 评估套件

TorchTune 集成了 EleutherAI 的评估工具链。我们可以使用 truthfulqa_mc2 任务来评估模型回答问题的真实性。

首先安装依赖：

pip install lm_eval>=0.4.5

然后复制并修改评估配置文件：

tune cp eleuther_evaluation ./custom_eval_config.yaml

关键配置项包括：

output_dir: 指向训练好的模型目录
tokenizer.path: 分词器模型路径
tasks: ["truthfulqa_mc2"] 指定评估任务

运行评估：

tune run eleuther_eval --config ./custom_eval_config.yaml

生成测试

除了标准化评估，直接生成文本观察模型表现也很重要。我们可以使用生成配方：

tune cp generation ./custom_generation_config.yaml

修改配置文件后运行生成：

tune run generate --config ./custom_generation_config.yaml prompt.user="Tell me a joke. "

常用生成参数：

temperature=0.6 (控制生成多样性)
top_k=300 (限制候选词数量)
max_new_tokens=300 (最大生成长度)

模型量化

TorchTune 使用 torchao 进行训练后量化。量化可以显著减少模型大小并提升推理速度。

支持多种量化技术：

int4_weight_only (4位权重量化)
int8_weight_only (8位权重量化)
int8_dynamic_activation_int8_weight (动态8位激活和权重量化)

量化示例代码：

from torchao.quantization.quant_api import quantize_, int4_weight_only
quantize_(model, int4_weight_only())

量化后可以使用 torch.compile 进一步优化性能。对于 Llama 系列模型，torchao 提供了专门的生成脚本用于量化模型推理。

模型部署

与 Hugging Face 生态集成

方案1: 基础模型+适配器

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

trained_model_path = "path/to/trained/model"
original_model_name = "meta-llama/Llama-3.2-1B-Instruct"

model = AutoModelForCausalLM.from_pretrained(original_model_name)
peft_model = PeftModel.from_pretrained(model, trained_model_path)
tokenizer = AutoTokenizer.from_pretrained(original_model_name)

方案2: 使用合并权重

from transformers import AutoModelForCausalLM, AutoTokenizer

trained_model_path = "path/to/trained/model"
model = AutoModelForCausalLM.from_pretrained(trained_model_path)