FullModelHFCheckpointer¶
- class torchtune.training.FullModelHFCheckpointer(checkpoint_dir: str, checkpoint_files: Union[List[str], Dict[str, str]], model_type: ModelType, output_dir: str, adapter_checkpoint: Optional[str] = None, recipe_checkpoint: Optional[str] = None, resume_from_checkpoint: bool = False, safe_serialization: bool = False)[源代码]¶
检查点管理器,以 HF 格式读取和写入检查点。对于 LoRA 模型,这包括以可通过例如
from_pretrained
加载到 PEFT 的格式保存检查点。示例包括来自 meta-llama 存储库的 Llama-2-7b-hf 模型 (https://hugging-face.cn/meta-llama/Llama-2-7b-hf)。注意
HF 检查点名称通常按 ID 排序(例如:0001_of_0003、0002_of_0003 等)。为了确保我们按正确的顺序读取文件,我们在读取之前对检查点文件名进行排序。
注意
检查点到 HF 格式的转换以及从 HF 格式的转换需要访问模型参数,这些参数直接从
config.json
文件中读取。这有助于确保我们正确加载权重,或者在 HF 检查点文件与 torchtune 的模型实现之间存在差异时出错。- 参数:
checkpoint_dir (str) – 包含检查点文件的目录
checkpoint_files (Union[List[str], Dict[str, str]]) – 要加载的检查点文件的列表。由于检查点管理器负责按文件 ID 排序,因此此列表中的顺序无关紧要。待办事项:更新此内容
model_type (ModelType) – 加载检查点管理器的模型的模型类型
output_dir (str) – 保存检查点文件的目录
adapter_checkpoint (Optional[str]) – 自适应器权重的路径。默认为 None
recipe_checkpoint (Optional[str]) – 方案状态检查点文件的路径。默认为 None
resume_from_checkpoint (bool) – 如果为 True,则检查点管理器将加载其他检查点文件以从之前的运行恢复训练。默认为 False
safe_serialization (bool) – 如果为 True,则检查点管理器将使用 safetensors 保存检查点文件
- 引发:
ValueError – 如果
resume_from_checkpoint
为 True 但recipe_checkpoint
为 None
- load_checkpoint() Dict[str, Any] [source]¶
从文件中加载 HF 检查点。
来自所有检查点文件中的键和权重合并到单个 state_dict 中。我们保留“state_dict 键” <-> “检查点文件”映射在 weight_map 中,以便我们可以在
save_checkpoint
中正确写入 state_dict。在返回之前,模型 state_dict 使用适当的 convert_weights 函数(取决于
self._model_type
)转换为与 torchtune 兼容的格式。- 返回值:
torchtune 检查点状态字典
- 返回类型:
state_dict (Dict[str, Any])
- 引发:
ValueError – 如果输入 state_dict 中的值不是张量
- save_checkpoint(state_dict: Dict[str, Any], epoch: int, intermediate_checkpoint: bool = False, adapter_only: bool = False) None [source]¶
将 HF 检查点保存到文件。如果
intermediate_checkpoint
为 True,则在_output_dir
中创建额外的检查点文件recipe_state.pt
,其中包含配方状态。state_dict 首先转换回 HF 格式,然后根据
_weight_map
分割成单独的检查点文件。- 参数:
- 引发:
ValueError – 如果
adapter_only
为 True 且在 state_dict 中找不到适配器检查点。