精细调优大语言模型：揭开HuggingFace的神秘面纱！如何克服GPU内存束缚？

每次大型语言模型 (LLMs) 的新公告往往将性能推向新的高度，常常超越之前的基准（例如，巨量多任务语言理解或 MMLU）。这一进展激发了许多应用程序的出现，利用最大的和最优秀的模型。在我们之前的帖子中，我们讨论了 LLMs 的规模法则，并解释了为什么更大的模型在预测下一个标记时表现更好。

然而，从原型 LLM 演示到功能性生产系统的旅程并非没有挑战。用户隐私和信任非常重要，尤其是在我们运营会计和金融领域时。确保 LLMs 增强我们的应用程序并为客户提供价值仍然是我们的首要任务。

我们观察到，当使用专有模型（如 GPT-4 及其变体）时，模型响应会随着时间的推移而漂移。这是可以理解的，因为模型的权重会更新。然而，这会造成明显的推理差异，从而使下游应用程序的性能不稳定。此外，我们还观察到，即使为随机种子设置了一切可能的种子并选择贪婪解码策略，也很难获得确定性的结果。此外，我们意识到 GPT-4 在回答某些特定于我们领域的问题时表现不佳，可能会生成虚假的结果。

为了缓解这些问题，我们决定训练自己的会计领域特定模型，并与 Amazon AWS 团队合作，利用他们的基础设施和知识提供最先进的训练和推理计算资源。然而，微调 LLMs 的一个挑战是高 GPU 内存消耗。在接下来的部分中，我们将探讨模型并行和数据并行策略，这些策略通常用于解决这一问题。

模型并行与数据并行

在训练 LLM 时，选择 模型并行 (MP) 和数据并行 (DP) 取决于模型是否能够适应单个 GPU。 当模型的权重过大而无法放入单个 GPU 时，使用 MP。MP 将模型权重分片到多个 GPU 上。另一方面，DP 将数据分片到多个 GPU 上，每个 GPU 持有模型权重的完整副本。

在 DP 的情况下，PyTorch 提供了 分布式数据并行 (DDP) 功能，它包装了模型类对象，允许通过其 torchrun 启动工具启动 DDP。

在 MP 的情况下，PyTorch 原生支持 完全分片数据并行 (FSDP) 以实现模型并行。此外，DeepSpeed 是由微软团队开发的另一种流行的 MP 技术，我们在之前的文章中使用过。

正如您可能注意到的，使用这些并行技术需要对代码进行修改。为了简化这个过程，HuggingFace 提供了 accelerate 包，它处理复杂的包装函数和 GPU 设备放置。这使得使用这些并行技术变得更加简单，而无需编写特定于并行化的样板代码。

准备用于指令微调的数据集

首先，我们将数据集整理成问题和答案的对。由于 LLM 以 token-in 和 token-out 的方式运作，因此对中的文本使用 tokenizer 对象转换为 tokens。对于指令微调模型，这些对需要符合每个模型独特的指令格式。我们可以使用 tokenizer.apply_chat_template() 函数来实现这一点。

def apply_chat_template(
    datapoint,
    tokenizer,
    prefix="",
):
    messages = [
        {"role": "user", "content": f"""{prefix}: [{datapoint["input"]}]"""},
        {"role": "assistant", "content": datapoint["output"]},
    ]
    datapoint["text"] = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    return datapoint


def apply_chat_template_test(
    datapoint,
    tokenizer,
    prefix="",
):
    print("using chat template prompting.")
    messages = [
        {"role": "user", "content": f"""{prefix}: [{datapoint["input"]}]"""},
    ]
    datapoint["text"] = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    return datapoint
```python
def apply_chat_template(
    datapoint,
    tokenizer,
    prefix="",
):
    messages = [
        {"role": "user", "content": f"""{prefix}: [{datapoint["input"]}]"""},
        {"role": "assistant", "content": datapoint["output"]},
    ]
    datapoint["text"] = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    return datapoint


def apply_chat_template_test(
    datapoint,
    tokenizer,
    prefix="",
):
    print("using chat template prompting.")
    messages = [
        {"role": "user", "content": f"""{prefix}: [{datapoint["input"]}]"""},
    ]
    datapoint["text"] = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    return datapoint

这是完整的代码片段。

from sklearn.model_selection import train_test_split

import pandas as pd
df = pd.read_csv(data_path)


## 将数据集拆分为训练集和测试集
X_train, X_test = train_test_split(
  df, test_size=0.00001, random_state=42
)


## 重新格式化输入和输出以符合指令格式
X_train = pd.DataFrame(
      X_train.apply(
          lambda row: apply_chat_template(
              row, tokenizer, prefix
          ),
        axis=1,
    ),
    columns=["text"],
)
X_test = pd.DataFrame(
    X_test.apply(
        lambda row: apply_chat_template_test(
            row, tokenizer, prefix
        ),
        axis=1,
    ),
    columns=["text"],
)

## 将数据加载到 Dataset 对象中
from datasets import Dataset

train_data = Dataset.from_pandas(X_train)
test_data = Dataset.from_pandas(X_test)
```python
from sklearn.model_selection import train_test_split

import pandas as pd
df = pd.read_csv(data_path)


## 将数据集拆分为训练集和测试集
X_train, X_test = train_test_split(
  df, test_size=0.00001, random_state=42
)


## 重新格式化输入和输出以符合指令格式
X_train = pd.DataFrame(
      X_train.apply(
          lambda row: apply_chat_template(
              row, tokenizer, prefix
          ),
        axis=1,
    ),
    columns=["text"],
)
X_test = pd.DataFrame(
    X_test.apply(
        lambda row: apply_chat_template_test(
            row, tokenizer, prefix
        ),
        axis=1,
    ),
    columns=["text"],
)

## 将数据加载到 Dataset 对象中
from datasets import Dataset

train_data = Dataset.from_pandas(X_train)
test_data = Dataset.from_pandas(X_test)

模型训练

LLM 以监督的方式进行训练，因为我们为问题提供了标记答案。为了实现这个训练过程，HuggingFace 实现了 SFTTrainer() 对象。

加载预训练模型和分词器

首先，我们使用 HuggingFace 的 AutoModelForCausalLM 和 AutoTokenizer 对象加载预训练模型和分词器。一些模型（例如 Phi-3 系列）需要 trust_remote_code=True，因此我们默认将此参数设置为 true。

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)
import torch


## Load the model
def load_model(
    base_model,
):
    return AutoModelForCausalLM.from_pretrained(
        base_model,
        trust_remote_code=True,
    )


## Load tokenizer
def load_tokenizer(base_model,):
    return AutoTokenizer.from_pretrained(
        base_model, trust_remote_code=True,
    )
```python
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)
import torch


## Load the model
def load_model(
    base_model,
):
    return AutoModelForCausalLM.from_pretrained(
        base_model,
        trust_remote_code=True,
    )


## Load tokenizer
def load_tokenizer(base_model,):
    return AutoTokenizer.from_pretrained(
        base_model, trust_remote_code=True,
    )

在开始训练之前，我们为该过程设置一个种子，以确保可重复性。

from transformers import set_seed


## Set seed for reproducibility
set_seed(123)
```python
from transformers import set_seed


## Set seed for reproducibility
set_seed(123)

LoRA 微调

从零开始训练大型语言模型（LLMs）需要大量的计算资源和庞大的数据集（例如，15万亿个标记用于 Llama 3.1 405B，16,000 个 H100 GPU）。幸运的是，如果仅对线性层进行微调，微调 LLMs 所需的计算资源要少得多。

LoRA 利用低秩矩阵近似的优势，这是一种将高维矩阵近似为两个较小矩阵（例如，秩为 1）的乘积的概念，从而将矩阵更新维度从 25 降低到 10，如图 1 所示。这种方法源于奇异值分解（SVD）的数学技术。通过将其应用于变换器模型的注意力层中的权重矩阵，LoRA 有效地减少了需要训练的参数数量。

peft 包提供了 prepare_model_for_kbit_training 和 LoraConfig 对象，以便您可以向目标线性层添加可调权重，同时冻结其他层的权重。

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training


def add_lora_layers(
    model,
    lora_alpha=<an interger>,
    r=<an interger>,
    target_modules="all-linear",
):
    model = prepare_model_for_kbit_training(
        model,
        use_gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )
    peft_config = LoraConfig(
        lora_alpha=lora_alpha,
        lora_dropout=<a floating number>,
        r=r,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=target_modules,
    )
    return get_peft_model(model, peft_config), peft_config
```python
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training


def add_lora_layers(
    model,
    lora_alpha=<an interger>,
    r=<an interger>,
    target_modules="all-linear",
):
    model = prepare_model_for_kbit_training(
        model,
        use_gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )
    peft_config = LoraConfig(
        lora_alpha=lora_alpha,
        lora_dropout=<a floating number>,
        r=r,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=target_modules,
    )
    return get_peft_model(model, peft_config), peft_config

定义模型超参数

我们需要为训练模型定义超参数。我们选择余弦学习率调度，这提供了平滑的衰减。您应该仔细选择每个设备的批大小和梯度累积步数，以避免耗尽您的GPU内存。有效批大小是通过将每个设备的批大小乘以梯度累积步数和集群（节点）中的GPU数量来计算的。例如，Nvidia为云服务提供商（CSP）打包并运输8个A100 GPU在单个节点中。此外，我们保存达到最低交叉熵损失的前三个模型。

from transformers import EarlyStoppingCallback, TrainingArguments, trainer_utils
import torch


## If your gpu is above 8, you can accelerate training with bf16
major, _ = torch.cuda.get_device_capability()

## Set up training hyperparameters
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=<an interger>,
    per_device_train_batch_size=<an interger>,
    per_device_eval_batch_size=<an interger>,
    gradient_accumulation_steps=<an interger>,
    optim="paged_adamw_32bit",
    save_steps=<an integer>,
    logging_steps=<an integer>,
    learning_rate=<a floating number>,
    weight_decay=<a floating number>,
    fp16=False if major >= 8 else True,
    bf16=True if major >= 8 else False,
    max_grad_norm=<a floating number>,
    warmup_ratio=<a floating number>,
    group_by_length=True,
    lr_scheduler_type="cosine",
    gradient_checkpointing_kwargs={"use_reentrant": False},
    disable_tqdm=False,
    resume_from_checkpoint=True,
    seed=123,
    log_level="info",
    remove_unused_columns=True,
    ...
)
```python
from transformers import EarlyStoppingCallback, TrainingArguments, trainer_utils
import torch


## If your gpu is above 8, you can accelerate training with bf16
major, _ = torch.cuda.get_device_capability()

## Set up training hyperparameters
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=<an interger>,
    per_device_train_batch_size=<an interger>,
    per_device_eval_batch_size=<an interger>,
    gradient_accumulation_steps=<an interger>,
    optim="paged_adamw_32bit",
    save_steps=<an integer>,
    logging_steps=<an integer>,
    learning_rate=<a floating number>,
    weight_decay=<a floating number>,
    fp16=False if major >= 8 else True,
    bf16=True if major >= 8 else False,
    max_grad_norm=<a floating number>,
    warmup_ratio=<a floating number>,
    group_by_length=True,
    lr_scheduler_type="cosine",
    gradient_checkpointing_kwargs={"use_reentrant": False},
    disable_tqdm=False,
    resume_from_checkpoint=True,
    seed=123,
    log_level="info",
    remove_unused_columns=True,
    ...
)

应用 SFTTrainer 对象

接下来，我们通过调用 SFTTrainer() 对象来启动训练，传递训练数据集和超参数。我们还加入了 callbacks，以便在损失超过我们定义的限制后未能减小时提前结束训练。

from trl import SFTTrainer


## Set up the trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=(
        test_data
    ),
    peft_config=peft_config,
    max_seq_length=max_seq_length,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_arguments,
    packing=False,
    callbacks=[
        EarlyStoppingCallback(
            early_stopping_patience=<an interger>,
            early_stopping_threshold=<a floating number>,
        ),
    ],
)
```python
from trl import SFTTrainer


## Set up the trainer
trainer = SFTTrainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=(
        test_data
    ),
    peft_config=peft_config,
    max_seq_length=max_seq_length,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_arguments,
    packing=False,
    callbacks=[
        EarlyStoppingCallback(
            early_stopping_patience=<an interger>,
            early_stopping_threshold=<a floating number>,
        ),
    ],
)

日志模型指标

方便的是，我们可以使用 logging 函数记录超参数和损失指标。

import datasets
import logging
import transformers


logger = logging.getLogger(__name__)


## setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
log_level = training_arguments.get_process_log_level()
logger.setLevel(log_level)
datasets.utils.logging.set_verbosity(log_level)
transformers.utils.logging.set_verbosity(log_level)
transformers.utils.logging.enable_default_handler()
transformers.utils.logging.enable_explicit_format()
```python
import datasets
import logging
import transformers


logger = logging.getLogger(__name__)


## setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
log_level = training_arguments.get_process_log_level()
logger.setLevel(log_level)
datasets.utils.logging.set_verbosity(log_level)
transformers.utils.logging.set_verbosity(log_level)
transformers.utils.logging.enable_default_handler()
transformers.utils.logging.enable_explicit_format()

最后，我们记录训练和测试数据集的指标，包括损失和训练过程中使用的超参数。

resume_from_checkpoint = True

if resume_from_checkpoint:
    print("Continue training from the last checkpoint.")
    train_result = trainer.train(
        resume_from_checkpoint=resume_from_checkpoint
    )
else:
    train_result = trainer.train()

metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()

## Evaluation
tokenizer.padding_side = "left"
metrics = trainer.evaluate()
metrics["eval_samples"] = (
    len(test_data)
)
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)
```python
resume_from_checkpoint = True

if resume_from_checkpoint:
    print("Continue training from the last checkpoint.")
    train_result = trainer.train(
        resume_from_checkpoint=resume_from_checkpoint
    )
else:
    train_result = trainer.train()

metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()

## Evaluation
tokenizer.padding_side = "left"
metrics = trainer.evaluate()
metrics["eval_samples"] = (
    len(test_data)
)
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)

使用 Accelerate 和 DeepSpeed 配置启动

我们可以通过命令行启动训练任务，并将到目前为止的所有代码放入一个 Python 脚本中。DeepSpeed 配置在标志 --config_file 后传递。请确保 gradient_accumulation_steps 和 gradient_clipping 与 TrainingArguments 中定义的值匹配。num_machines 定义使用的节点数量。num_processes 定义节点中的 GPU 数量。

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  gradient_accumulation_steps: <an interger>
  gradient_clipping: <a floating number>
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```python
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  gradient_accumulation_steps: <an interger>
  gradient_clipping: <a floating number>
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

CUDA_VISIBLE_DEVICE=0,1,2,3,4,5,6,7 accelerate launch --config_file "deepspeed_stage3.yaml" <put all of the code up to this point into a python script>
```python
CUDA_VISIBLE_DEVICE=0,1,2,3,4,5,6,7 accelerate launch --config_file "deepspeed_stage3.yaml" <put all of the code up to this point into a python script>

权衡：精度与速度

我们希望通过强调其他重要的权衡来结束这篇文章：精度与速度。图2显示了不同数据类型与计算速度之间的关系，以TFLOPS为单位，详细信息见Nvidia的A100白皮书。要深入了解计算速度，您可以参考我们早期博客文章中的附录。

不同的数据类型提供了不同的指数范围和精度。此外，不同的GPU（如V100和A100）支持不同的数据类型。检查您的GPU支持哪些数据类型非常重要。例如，根据上面的图2，FP16在操作/秒（计算速度）中提供了FP32的8倍吞吐量，尽管其指数范围和精度有所降低。我们建议仔细查看PyTorch CUDA设置以及您的GPU白皮书，并对您的数据进行实验，以找到精度与速度之间的正确平衡。例如，以下代码片段启用了TF32用于矩阵乘法。

## https://pytorch.org/docs/stable/notes/cuda.html#tf32-on-ampere

## The flag below controls whether to allow TF32 on matmul. This flag defaults to False
## in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True

## The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True
```python
## https://pytorch.org/docs/stable/notes/cuda.html#tf32-on-ampere

## The flag below controls whether to allow TF32 on matmul. This flag defaults to False
## in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True

## The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True

结论

在这篇文章中，我们展示了如何使用 HuggingFace 提供的工具通过 MP 技术对 LLM 进行微调。通过分享这些经验，我们希望您能训练自己的 LLM 并拥有权重。请记住，即使是经过微调的 LLM 也可能产生幻觉响应。虽然微调可以减少幻觉的可能性，但并不能消除它。开发人员和科学家在设计 AI 产品时应理解，LLM 仍然是通过根据前面的上下文生成最可能的标记来操作的，并且可能生成无根据或不真实的响应。尽管存在这些限制，小型特定领域语言模型更受欢迎，因为它们可以被优化以理解特定领域的术语，并且与大型封闭源语言模型相比，推理成本更低。