检查点机制在 torchtune 中的应用¶
本次深入探讨将带您了解检查点机制及其相关实用程序的设计和行为。
torchtune 的检查点机制设计
检查点格式以及我们如何处理它们
检查点机制应用场景：中间检查点与最终检查点，以及 LoRA 与完全微调
概述¶
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 个 token，每个 token 的嵌入维度为 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 │ ├── 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 检查点机制支持两种检查点应用场景
训练结束检查点
完成训练运行结束时的模型权重被写入文件。检查点机制确保输出检查点文件具有与用于开始训练的输入检查点文件相同的键。检查点机制还确保键被分区到与原始检查点相同数量的文件中。输出 state dict 具有以下标准格式
{ "key_1": weight_1, "key_2": weight_2, ... }
训练中期检查点.
如果在训练中期进行检查点，则输出检查点需要存储额外的信息，以确保后续的训练运行可以正确地重新启动。除了模型检查点文件外，我们还为中间检查点输出一个 recipe_state.pt 文件。这些文件目前在每个 epoch 结束时输出，并包含诸如优化器状态、已完成 epoch 数等信息。
为了防止我们用检查点文件淹没 output_dir，配方状态在每个 epoch 结束时被覆盖。
输出 state dict 具有以下格式
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 恢复时，我们将从 checkpoint dir 中获取原始未训练的权重，并从 output_dir 中获取训练好的适配器。有关更多详细信息，请查看我们的 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 支持的任何模型来执行此操作。您可以在此处找到模型和模型构建器的完整列表。
我们希望本次深入探讨能让您更深入地了解 torchtune 中的检查点机制和相关实用程序。祝您调优愉快！
