torchtune 中的 Checkpointing¶

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

本次深入探讨将涵盖的内容

torchtune 的检查点工具设计
检查点格式以及我们如何处理它们
Checkpointing 场景：中间检查点 vs 最终检查点，以及 LoRA vs 全模型微调

概述¶

torchtune 检查点工具被设计为可组合组件，可以插入到任何代码示例中——无论是训练、评估还是生成。每个检查点工具都支持一组模型和场景，使它们易于理解、调试和扩展。

在我们深入了解 torchtune 中的检查点工具之前，先定义一些概念。

检查点格式¶

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

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

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

Meta 格式

这是官方 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 Model Hub 中最流行的格式，也是每个 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 检查点工具通过为您管理 state dict 来减少出错的可能性。torchtune 被设计为“state-dict 不变”。

加载时，torchtune 接受来自多个来源、采用多种格式的检查点。您无需在每次运行代码示例时都显式转换检查点。
保存时，torchtune 生成与源格式相同的检查点。这包括将 state_dict 转换回原始形式，并将键和权重分割到相同数量的文件中。

成为“state-dict 不变”的一个重要优势是，您应该能够在不更改任何代码或转换脚本的情况下，将 torchtune 微调的检查点用于任何训练后工具（量化、评估、推理），只要该工具支持源格式。这是 torchtune 与周边生态系统互操作的方式之一。

注意

为了以这种方式实现 state-dict“不变”，每个检查点工具的 load_checkpoint 和 save_checkpoint 方法都使用了权重转换器，这些转换器可以正确地在检查点格式之间映射权重。例如，从 Hugging Face 加载权重时，我们在加载和保存时会对某些权重进行置换，以确保检查点行为完全相同。为了进一步说明这一点，Llama 系列模型使用了一个通用权重转换函数，而 Phi3 等其他一些模型则有自己的转换函数，可以在其模型文件夹中找到。

处理不同的检查点格式¶

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

`HFCheckpointer`¶

此检查点工具以与 Hugging Face 的 transformers 框架兼容的格式读写检查点。如上所述，这是 Hugging Face Model Hub 中最流行的格式，也是每个 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 格式之间的相互转换需要访问直接从 config.json 文件读取的模型参数。这有助于确保我们正确加载权重，或者在 HF 检查点文件和 torchtune 模型实现之间存在差异时报错。此 json 文件与模型检查点一起从 hub 下载。

`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 文件，其中包含从上一个中间 epoch 重新启动训练运行所需的信息。稍后会详细介绍；
logs: 您的 metric_logger 的输出（如果有）；
epoch_{}: 包含您训练好的模型权重和模型元数据。如果运行推理或推送到模型 hub，应直接使用此文件夹；

注意

对于每个 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
│   ├── model-00001-of-00002.safetensors
│   ├── 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
│   ├── model-00001-of-00002.safetensors
│   ├── 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

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

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

训练结束时检查点

训练运行完成后，模型权重会被写入文件。检查点工具确保输出检查点文件与用于开始训练的输入检查点文件具有相同的键。检查点工具还确保键被分割到与原始检查点相同数量的文件中。输出 state dict 具有以下标准格式：

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

.

训练中途检查点

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

为了防止我们的 output_dir 被检查点文件淹没，recipe state 在每个 epoch 结束时会被覆盖。

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

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

输出 state dict 具有以下格式：

从检查点恢复 - 全模型微调¶

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

resume_from_checkpoint: 将其设置为 True；

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

checkpointer:
    # [... rest of the config...]

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

# set to True if restarting training
resume_from_checkpoint: True

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

从检查点恢复 - LoRA 微调¶

checkpointer:
    # [... rest of the config...]

    # 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.pt

# 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

注意

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

注意

在 torchtune 中，我们为 LoRA 输出适配器权重和全模型合并权重。合并后的检查点是一种便利，因为它可以在无需特殊工具处理适配器的情况下使用。然而，它们不应用于恢复训练，因为加载合并权重 + 适配器将会出错。因此，在恢复 LoRA 训练时，我们将从 checkpoint dir 中获取原始未训练的权重，并从 output_dir 中获取训练好的适配器。更多详细信息，请查看我们的LoRA 微调教程。

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

综合所有内容¶

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

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')

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

您可以使用 torchtune 支持的任何模型执行此操作。您可以在此处找到完整的模型和模型构建器列表。

torchtune 中的 Checkpointing¶

概述¶

检查点格式¶

处理不同的检查点格式¶

`HFCheckpointer`¶

`MetaCheckpointer`¶

`TorchTuneCheckpointer`¶

检查点输出¶

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

输出 state dict 具有以下格式：

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

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

文档

教程

查看教程

torchtune 中的 Checkpointing¶

概述¶

检查点格式¶

处理不同的检查点格式¶

HFCheckpointer¶

MetaCheckpointer¶

TorchTuneCheckpointer¶

检查点输出¶

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

输出 state dict 具有以下格式：

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

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

文档

教程

查看教程

`HFCheckpointer`¶

`MetaCheckpointer`¶

`TorchTuneCheckpointer`¶

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