torchtune中的检查点¶

本次深入探讨将引导您了解检查点及其相关实用程序的设计和行为。

What this deep-dive will cover:

torchtune的检查点设计
检查点格式以及我们如何处理它们
检查点场景：中间与最终以及LoRA与完全微调

概述¶

torchtune检查点设计为可组合的组件，可以插入任何配方中 - 训练、评估或生成。每个检查点支持一组模型和场景，使得这些易于理解、调试和扩展。

在我们深入探讨torchtune中的检查点之前，让我们先定义一些概念。

检查点格式¶

在这次深入探讨中，我们将讨论不同的检查点格式以及torchtune如何处理它们。让我们仔细看看这些不同的格式。

简单来说，检查点的格式由state_dict决定，以及它如何存储在磁盘上的文件中。每个权重都与一个字符串键相关联，该键在state dict中标识它。如果存储的检查点中的键的字符串标识符与模型定义中的不完全匹配，你可能会遇到明确的错误（加载state dict会引发异常）或更糟糕的情况——静默错误（加载会成功，但训练或推理不会按预期工作）。除了键要匹配外，你还需要权重的形状（state_dict中的值）与模型定义中预期的完全匹配。

让我们来看看Llama 3.2的两种流行格式。

元格式

这是官方Llama 3.2实现支持的格式。当你从meta-llama网站下载Llama 3.2 3B模型时，你将获得一个单一的.pth检查点文件。你可以使用torch.load轻松检查此检查点的内容。

>>> import torch
>>> state_dict = torch.load('consolidated.00.pth', mmap=True, weights_only=True, map_location='cpu')
>>> # inspect the keys and the shapes of the associated tensors
>>> for key, value in state_dict.items():
>>>    print(f'{key}: {value.shape}')

tok_embeddings.weight: torch.Size([128256, 3072])
...
...
>>> print(len(state_dict.keys()))
255

state_dict 包含 255 个键，其中包括一个名为 tok_embeddings 的输入嵌入表。该 state_dict 的模型定义期望一个嵌入层，其中包含 128256 个标记，每个标记的嵌入维度为 3072。

HF格式

这是Hugging Face模型中心中最流行的格式，也是每个torchtune配置中的默认格式。这也是您从Llama-3.2-3B-Instruct仓库下载llama3.2模型时获得的格式。

第一个主要区别是state_dict被拆分到两个.safetensors文件中。要正确加载检查点，你需要将这些文件拼接在一起。让我们检查其中一个文件。

>>> from safetensors import safe_open
>>> state_dict = {}
>>> with safe_open("model-00001-of-00002.safetensors", framework="pt", device="cpu") as f:
>>>     for k in f.keys():
>>>         state_dict[k] = f.get_tensor(k)

>>> # inspect the keys and the shapes of the associated tensors
>>> for key, value in state_dict.items():
>>>     print(f'{key}: {value.shape}')

model.embed_tokens.weight: torch.Size([128256, 3072])
...
...
>>> print(len(state_dict.keys()))
187

不仅state_dict包含的键更少（预期如此，因为这是两个文件之一），而且嵌入表被称为model.embed_tokens而不是tok_embeddings。这种名称不匹配会在尝试加载state_dict时导致异常。这一层的大小在两者之间是相同的，这是预期的。

正如你所见，如果不小心，你可能在检查点加载和保存过程中犯下许多错误。torchtune检查点管理器通过为你管理状态字典，减少了这种错误的可能性。torchtune被设计为“状态字典不变”。

加载时，torchtune 接受来自多个来源的多种格式的检查点。每次运行配方时，您不必担心显式转换检查点。
保存时，torchtune 会生成与源格式相同的检查点。这包括将 state_dict 转换回原始形式，并将键和权重拆分到相同数量的文件中。

“状态字典不变”的一个大优势是，您应该能够使用来自torchtune的微调检查点与任何支持源格式的训练后工具（量化、评估、推理），而无需任何代码更改或转换脚本。这是torchtune与周围生态系统互操作的方式之一。

注意

为了以这种方式保持状态字典的“不变性”，每个检查点的load_checkpoint和save_checkpoint方法都使用了权重转换器，这些转换器能够正确地在检查点格式之间映射权重。例如，当从Hugging Face加载权重时，我们会在加载和保存时对某些权重应用排列，以确保检查点的行为完全相同。为了进一步说明这一点，Llama系列模型使用了通用权重转换函数，而像Phi3这样的其他模型则有自己的转换函数，这些函数可以在它们的模型文件夹中找到。

处理不同的检查点格式¶

torchtune 支持三种不同的检查点，每种都支持不同的检查点格式。

`HFCheckpointer`¶

此检查点读取器以与Hugging Face的transformers框架兼容的格式读取和写入检查点。如上所述，这是Hugging Face模型中心中最流行的格式，也是每个torchtune配置中的默认格式。

为了使此检查点正常工作，我们假设checkpoint_dir包含必要的检查点和json文件。确保一切正常工作的最简单方法是使用以下流程：

使用tune download从HF仓库下载模型。这将忽略“pth”文件，因为我们将加载“safetensors”。

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

使用此处指定的output_dir作为检查点的checkpoint_dir参数。

以下代码片段解释了如何在torchtune配置文件中设置HFCheckpointer。

checkpointer:

    # checkpointer to use
    _component_: torchtune.training.FullModelHFCheckpointer

    # directory with the checkpoint files
    # this should match the folder you used when downloading the model
    checkpoint_dir: /tmp/Llama-3.2-3B-Instruct

    # checkpoint files. For the Llama-3.2-3B-Instruct model we have
    # 2 .safetensor files. The checkpointer takes care of sorting
    # by id and so the order here does not matter
    checkpoint_files: [
        model-00001-of-00002.safetensors,
        model-00002-of-00002.safetensors,
    ]

    # dir for saving the output checkpoints
    output_dir: <output_dir>

    # model_type which specifies how to convert the state_dict
    # into a format which torchtune understands
    model_type: LLAMA3_2

# set to True if restarting training. More on that later.
resume_from_checkpoint: False

注意

将检查点转换为HF格式或从HF格式转换需要访问模型参数，这些参数直接从config.json文件中读取。这有助于确保我们正确加载权重，或者在HF检查点文件与torchtune的模型实现之间存在差异时出错。此json文件与模型检查点一起从中心下载。

`MetaCheckpointer`¶

此检查点读取器以与原始meta-llama GitHub仓库兼容的格式读取和写入检查点。

为了使此检查点正常工作，我们假设checkpoint_dir包含必要的检查点和json文件。确保一切正常工作的最简单方法是使用以下流程：

使用tune download从HF仓库下载模型。默认情况下，这将忽略“safetensors”文件。

tune download meta-llama/Llama-3.2-3B-Instruct \
--output-dir /tmp/Llama-3.2-3B-Instruct \
--ignore-patterns "*.safetensors"

使用上面的output_dir作为检查点的checkpoint_dir。

以下代码片段解释了如何在torchtune配置文件中设置MetaCheckpointer。

checkpointer:

    # checkpointer to use
    _component_: torchtune.training.FullModelMetaCheckpointer

    # directory with the checkpoint files
    # this should match the folder you used when downloading the model
    checkpoint_dir: <checkpoint_dir>

    # checkpoint files. For the llama3.2 3B model we have
    # a single .pth file
    checkpoint_files: [consolidated.00.pth]

    # dir for saving the output checkpoints.
    output_dir: <checkpoint_dir>

    # model_type which specifies how to convert the state_dict
    # into a format which torchtune understands
    model_type: LLAMA3_2

# set to True if restarting training. More on that later.
resume_from_checkpoint: False

`TorchTuneCheckpointer`¶

此检查点读取器以与torchtune模型定义兼容的格式读取和写入检查点。它不执行任何state_dict转换，目前用于测试或加载量化模型以进行生成。

检查点输出¶

恭喜你走到这一步！假设你已经按照我们的使用torchtune的端到端工作流程并使用我们的LoRA配方之一训练了一个llama 3.2 3B模型。

现在让我们可视化输出。一个简单的方法是运行tree -a path/to/outputdir，这应该会显示类似于下面的树结构。有3种类型的文件夹：

recipe_state: 保存了recipe_state.pt，其中包含从最后一个中间时期重新开始训练运行所需的信息。更多内容稍后介绍；
logs: 你的metric_logger的输出，如果有的话；
epoch_{}: 包含您训练好的模型权重以及模型元数据。如果进行推理或推送到模型中心，您应该直接使用此文件夹；

注意

对于每个epoch，我们复制原始检查点文件夹的内容，不包括原始检查点和大文件。这些文件是轻量级的，主要是配置文件，使用户更容易在下游应用程序中直接使用epoch文件夹。

有关每个文件的更多详细信息，请查看上面提到的端到端教程。

>>> tree -a /tmp/torchtune/llama3_2_3B/lora_single_device
/tmp/torchtune/llama3_2_3B/lora_single_device
├── epoch_0
│   ├── adapter_config.json
│   ├── adapter_model.pt
│   ├── adapter_model.safetensors
│   ├── config.json
│   ├── ft-model-00001-of-00002.safetensors
│   ├── ft-model-00002-of-00002.safetensors
│   ├── generation_config.json
│   ├── LICENSE.txt
│   ├── model.safetensors.index.json
│   ├── original
│   │   ├── orig_params.json
│   │   ├── params.json
│   │   └── tokenizer.model
│   ├── original_repo_id.json
│   ├── README.md
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   ├── tokenizer.json
│   └── USE_POLICY.md
├── epoch_1
│   ├── adapter_config.json
│   ├── adapter_model.pt
│   ├── adapter_model.safetensors
│   ├── config.json
│   ├── ft-model-00001-of-00002.safetensors
│   ├── ft-model-00002-of-00002.safetensors
│   ├── generation_config.json
│   ├── LICENSE.txt
│   ├── model.safetensors.index.json
│   ├── original
│   │   ├── orig_params.json
│   │   ├── params.json
│   │   └── tokenizer.model
│   ├── original_repo_id.json
│   ├── README.md
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   ├── tokenizer.json
│   └── USE_POLICY.md
├── logs
│   └── log_1734652101.txt
└── recipe_state
    └── recipe_state.pt

中间检查点与最终检查点¶

torchtune 检查点支持两种检查点场景：

训练结束检查点

在训练运行结束时，模型权重会被写入文件。检查点确保输出检查点文件具有与用于开始训练的输入检查点文件相同的键。检查点还确保键被分区到与原始检查点相同数量的文件中。输出状态字典具有以下标准格式：

{
    "key_1": weight_1,
    "key_2": weight_2,
    ...
}

训练中期检查点.

如果在训练过程中进行检查点保存，输出检查点需要存储额外的信息，以确保后续的训练运行可以正确重新启动。除了模型检查点文件外，我们还会为中间检查点输出一个recipe_state.pt文件。这些文件目前在每个epoch结束时输出，包含诸如优化器状态、完成的epoch数量等信息。

为了防止我们用检查点文件淹没output_dir，配方状态在每个周期结束时被覆盖。

输出状态字典具有以下格式：

Model:
    {
        "key_1": weight_1,
        "key_2": weight_2,
        ...
    }

Recipe State:
    {
        "optimizer": ...,
        "epoch": ...,
        ...
    }

从检查点恢复 - 完全微调¶

有时我们的训练会因某些原因中断。要从之前的检查点文件重新开始训练，您需要在配置中更新以下字段：

resume_from_checkpoint: 设置为True;

checkpoint_files: 将路径更改为 epoch_{YOUR_EPOCH}/ft-model={}-of-{}.safetensors;

请注意，我们不更改我们的checkpoint_dir或output_dir。由于我们是从检查点恢复的，我们知道在output_dir中查找它。

checkpointer:
    # checkpoint files. Note that you will need to update this
    # section of the config with the intermediate checkpoint files
    checkpoint_files: [
        epoch_{YOUR_EPOCH}/ft-model-00001-of-00002.safetensors,
        epoch_{YOUR_EPOCH}/ft-model-00001-of-00002.safetensors,
    ]

# set to True if restarting training
resume_from_checkpoint: True

从检查点恢复 - LoRA 微调¶

与完全微调类似，我们也只需要修改两个字段：resume_from_checkpoint 和 adapter_checkpoint，它们将从 output_dir 加载。我们不需要修改 checkpoint_files，因为加载的基础模型仍然是相同的。

checkpointer:

    # adapter_checkpoint. Note that you will need to update this
    # section of the config with the intermediate checkpoint files
    adapter_checkpoint: epoch_{YOUR_EPOCH}/adapter_model.safetensors

# set to True if restarting training
resume_from_checkpoint: True

# set to True to save only the adapter weights
# it does not influence resuming_from_checkpointing
save_adapter_weights_only: False

注意

在torchtune中，我们输出了LoRA的适配器权重和完整模型合并权重。合并的检查点是为了方便，因为它可以在不需要特殊工具处理适配器的情况下使用。然而，在恢复训练时不应使用它们，因为加载合并的权重+适配器会导致错误。因此，在恢复LoRA训练时，我们将从检查点目录中获取原始的未训练权重，并从输出目录中获取训练好的适配器。更多详情，请查看我们的LoRA微调教程。

注意

此外，通过设置选项save_adapter_weights_only，您可以选择仅保存适配器权重。这减少了保存检查点所需的存储空间和时间，但对从检查点恢复没有影响。

将所有内容整合在一起¶

现在让我们把所有这些知识结合起来！我们将加载一些检查点，创建一些模型并运行一个简单的前向传播。

在本节中，我们将使用HF格式的Llama-3.2-3B-Instruct模型。

import torch
from torchtune.models.llama3_2 import llama3_2_3b
from torchtune.training import FullModelHFCheckpointer

# Set the right directory and files
checkpoint_dir = "/tmp/Llama-3.2-3B-Instruct/"
output_dir = "/tmp/torchtune/llama3_2_3B/full_single_device"

pytorch_files = [
    "model-00001-of-00002.safetensors",
    "model-00002-of-00002.safetensors",
]

# Set up the checkpointer and load state dict
checkpointer = FullModelHFCheckpointer(
    checkpoint_dir=checkpoint_dir,
    checkpoint_files=pytorch_files,
    output_dir=output_dir,
    model_type="LLAMA3_2",
)
torchtune_sd = checkpointer.load_checkpoint()

# Setup the model and the input
model = llama3_2_3b()

# Model weights are stored with the key="model"
model.load_state_dict(torchtune_sd["model"])
model.to("cuda")

# We have 128256 vocab tokens; lets generate an input with 24 tokens
x = torch.randint(0, 128256, (1, 24), dtype=torch.long, device="cuda")

tensor([[[ 1.4299,  1.1658,  4.2459,  ..., -2.3259, -2.3262, -2.3259],
        [ 6.5942,  7.2284,  2.4090,  ..., -6.0129, -6.0121, -6.0127],
        [ 5.6462,  4.8787,  4.0950,  ..., -4.6460, -4.6455, -4.6457],
        ...,
        [-0.4156, -0.0626, -0.0362,  ..., -3.6432, -3.6437, -3.6427],
        [-0.5679, -0.6902,  0.5267,  ..., -2.6137, -2.6138, -2.6127],
        [ 0.3688, -0.1350,  1.1764,  ..., -3.4563, -3.4565, -3.4564]]],
    device='cuda:0')

你可以使用torchtune支持的任何模型来完成这个操作。你可以在这里找到完整的模型和模型构建器列表here。

我们希望这次深入探讨能让你对torchtune中的检查点工具及其相关实用程序有更深入的了解。祝你调参愉快！