如何使用 PyTorch 训练 LLM：分步指南

大型语言模型（LLM）是现代人工智能应用的核心组件，尤其是在自然语言处理方面。它们有能力高效处理和理解人类语言，应用场景涵盖虚拟助理、机器翻译、文本摘要以及问答系统等。

像 LangChain 这样的库有助于实现上述这类端到端的 AI 应用。我们的教程面向数据工程与数据应用的 LangChain 入门概述了您可以用 LangChain 做什么，包括 LangChain 解决的问题以及数据用例示例。

本文将详细讲解训练大型语言模型的整个过程：从搭建工作环境到使用 Pytorch 2.0.1 完成最终实现。PyTorch 是一个动态且灵活的深度学习框架，可实现简单清晰的模型实现。

前置条件

为充分理解本文内容，您应熟悉 Python 编程，具备深度学习与 Transformer 的基本概念，并了解 Pytorch 框架。完整源码可在 GitHub 获取。

在进入核心实现之前，我们需要安装并导入相关库。同时需要说明，本训练脚本受到了 Hugging Face 上该仓库的启发。

库安装

安装过程如下所示：

首先，我们在 Jupyter Notebook 中使用 %%bash 指令，将安装命令作为 bash 命令在单个单元格中运行。

Trl：用于结合强化学习训练 Transformer 语言模型。
Peft：使用参数高效微调（PEFT）方法，以高效适配预训练模型。
Torch：广泛使用的开源机器学习库。
Datasets：用于下载并加载许多常见的机器学习数据集。

Transformers：由 Hugging Face 开发的库，内置数千个预训练模型，可用于分类、摘要、翻译等多种文本任务。

%%bash
pip -q install trl
pip -q install peft
pip -q install torch
pip -q install datasets
pip -q install transformers
技术支持

现在，可以按如下方式导入这些模块：

import torch
from trl import SFTTrainer
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments技术支持

数据加载与准备

alpaca 数据集可在 Hugging Face 免费获取，我们将用其进行演示。该数据集包含三列主要字段：instructions、input 和 output。它们会被合并以生成最终的文本列。

通过提供数据集名称（即 tatsu-lab/alpaca）即可按如下方式加载：

train_dataset = load_dataset("tatsu-lab/alpaca", split="train")
print(train_dataset)技术支持

可以看到，结果数据是一个包含两个键的字典：

Features：包含数据的主要列
Num_rows：对应数据的总行数

train_dataset 的结构

可通过以下指令显示前五行。首先将字典转换为 pandas 数据框，然后显示这些行。

pandas_format = train_dataset.to_pandas()
display(pandas_format.head())技术支持

train_dataset 的前五行

为便于观察，我们打印前三行的信息。在此之前，需要安装 textwrap 库，将每行的最大字数设置为 50。第一条打印语句用 15 个短横线分隔每个块。

import textwrap
技术支持

for index in range(3):
   print("---"*15)
   print("Instruction:        
       {}".format(textwrap.fill(pandas_format.iloc[index]["instruction"],    
       width=50)))
   print("Output:  
       {}".format(textwrap.fill(pandas_format.iloc[index]["output"],  
       width=50)))
   print("Text: 
       {}".format(textwrap.fill(pandas_format.iloc[index]["text"],  
       width=50)))技术支持

前三行的详细信息

模型训练

在开始训练模型之前，我们需要设置一些前置内容：

预训练模型：我们将使用可在 Hugging Face 获取的 Salesforce/xgen-7b-8k-base 预训练模型。Salesforce 训练了这一系列名为 XGen-7B 的 7B LLM，采用标准稠密注意力，序列长度可达 8K，训练语料达 1.5T tokens。
分词器：用于对训练数据进行分词。加载预训练模型与分词器的代码如下：

pretrained_model_name = "Salesforce/xgen-7b-8k-base"
model = AutoModelForCausalLM.from_pretrained(pretrained_model_name, torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name, trust_remote_code=True)技术支持

训练配置

训练需要一些参数与配置，下面定义了两个重要的配置对象：一个 TrainingArguments 实例、一个 LoraConfig 实例，最后是 SFTTrainer 模型。

TrainingArguments

用于定义模型训练的参数。

在本例中，我们先用 output_dir 属性定义保存已训练模型的目标位置，然后再定义其他超参数，如优化方法、learning rate、训练轮数等。

model_training_args = TrainingArguments(
       output_dir="xgen-7b-8k-base-fine-tuned",
       per_device_train_batch_size=4,
       optim="adamw_torch",
       logging_steps=80,
       learning_rate=2e-4,
       warmup_ratio=0.1,
       lr_scheduler_type="linear",
       num_train_epochs=1,
       save_strategy="epoch"
   )技术支持

LoRAConfig

本场景中使用的主要参数包括 LoRA 中低秩变换矩阵的秩，设置为 16；以及 LoRA 额外参数的缩放因子，设置为 32。

此外，dropout 比例为 0.05，即训练过程中会忽略 5% 的输入单元。最后，由于我们处理的是因果语言建模任务，因此将任务类型初始化为 CAUSAL_LM。

SFTTrainer

用于结合训练数据、分词器以及上述配置来训练模型。

由于我们使用训练数据中的 text 字段，了解其分布对于设定序列的最大 token 数非常重要。

import matplotlib.pyplot as plt


pandas_format['text_length'] = pandas_format['text'].apply(len)


plt.figure(figsize=(10,6))
plt.hist(pandas_format['text_length'], bins=50, alpha=0.5, color='g')
plt.title('Distribution of Length of Text')
plt.xlabel('Length of Text')
plt.ylabel('Frequency')
plt.grid(True)
plt.show()技术支持

文本列长度的分布

根据上述观察可知，大多数文本长度在 0 到 1000 之间。如下所示，只有 4.5% 的文档长度超过 1024。

mask = pandas_format['text_length'] > 1024
percentage = (mask.sum() / pandas_format['text_length'].count()) * 100


print(f"The percentage of text documents with a length greater than 1024 is: {percentage}%")
技术支持

接着我们将序列的最大 token 数设置为 1024，长度超过该值的文本将被截断。

SFT_trainer = SFTTrainer(
       model=model,
       train_dataset=train_dataset,
       dataset_text_field="text",
       max_seq_length=1024,
       tokenizer=tokenizer,
       args=model_training_args,
       packing=True,
       peft_config=lora_peft_config,
   )技术支持

执行训练

准备就绪后，即可按如下方式启动模型训练过程：

tokenizer.pad_token = tokenizer.eos_token
model.resize_token_embeddings(len(tokenizer))
model = prepare_model_for_int8_training(model)
model = get_peft_model(model, lora_peft_config)
training_args = model_training_args
trainer = SFT_trainer
trainer.train()技术支持

需要说明的是，训练在带有 GPU 的云环境中进行，这能显著加快整体训练过程。而在本地计算机上训练则需要更长时间。

我们的博文在云端使用 LLM 与本地运行 LLM 的优缺点提供了为 LLM 选择最佳部署策略的关键考量。

让我们来理解上述代码片段在执行什么：

tokenizer.pad_token = tokenizer.eos_token：将填充 token 设置为与句子结束 token 相同。
model.resize_token_embeddings(len(tokenizer))：调整模型的 token 嵌入层大小以匹配分词器词表长度。
model = prepare_model_for_int8_training(model)：为 INT8 精度训练准备模型，通常会进行量化。
model = get_peft_model(model, lora_peft_config)：根据 PEFT 配置调整给定模型。
training_args = model_training_args：将预定义的训练参数赋给 training_args。
trainer = SFT_trainer：将 SFTTrainer 实例赋给变量 trainer。
trainer.train()：按照提供的设定触发模型的训练过程。

结语

本文清晰介绍了如何使用 PyTorch 训练大型语言模型。从数据集准备开始，依次讲解了前置条件、配置训练器，直至最终启动训练流程。

尽管我们使用了特定的数据集与预训练模型，但对其他兼容选项而言流程基本相同。现在您已经了解如何训练一款 LLM，便可将这些知识用于训练其他面向各类 NLP 任务的复杂模型。

欢迎阅读我们的指南如何使用 LangChain 构建 LLM 应用，进一步探索大型语言模型的强大功能。或者，如果您仍需系统学习大型语言模型的核心概念，也可查看我们的课程以深入进阶。

作者

Zoumana Keita

LinkedIn Twitter

一位多才多艺、乐于分享知识并回馈他人的数据科学家。Zoumana 是 YouTube 内容创作者，也是 Medium 上的顶尖科技写作者。他热爱演讲、编程与教学。Zoumana 拥有两个硕士学位：其一是在法国巴黎获得的计算机科学硕士，专注于机器学习；其二是在美国德州理工大学获得的数据科学硕士。他的职业生涯始于法国 Groupe OPEN 的软件开发工程师，随后加入 IBM 担任机器学习顾问，为保险公司构建端到端的 AI 解决方案。之后，Zoumana 加入了 Axionable——一家总部位于巴黎与蒙特利尔的首家“可持续 AI”初创公司。在那里，他担任数据科学家，为来自法国、蒙特利尔、新加坡和瑞士的客户交付 AI 产品，主要聚焦 NLP 场景。此外，他将 5% 的时间投入到研发工作。目前，他在 IFC（世界银行集团）担任高级数据科学家。

主题

人工智能