FullModelTorchTuneCheckpointer

class torchtune.training.FullModelTorchTuneCheckpointer(checkpoint_dir: str, checkpoint_files: List[str], model_type: str, output_dir: str, adapter_checkpoint: Optional[str] = None, recipe_checkpoint: Optional[str] = None, resume_from_checkpoint: bool = False)

检查点器，用于读取和写入与 torchtune 兼容的格式的检查点。无需转换权重。

目前，这仅支持读取单个检查点文件。 随着我们增加对更大模型的支持，这种情况可能会改变。

参数:

• checkpoint_dir (str) – 包含检查点文件的目录

• checkpoint_files (List[str]) – 要加载的检查点文件列表。 由于检查点器负责按文件 ID 排序，因此此列表中的顺序无关紧要

• model_type (str) – 要加载检查点器的模型的模型类型，例如 LLAMA3。

• output_dir (str) – 用于保存检查点文件的目录

• adapter_checkpoint (Optional[str]) – 适配器权重的路径。 如果为 None 且 resume_from_checkpoint=True，则在 output_dir/epoch_{largest_epoch} 中查找 adapter_model.pt。 默认为 None。

• recipe_checkpoint (Optional[str]) – 配方状态检查点文件的路径。 如果为 None 且 resume_from_checkpoint=True，则在 output_dir/RECIPE_STATE_DIRNAME 中查找 recipe_state.pt。 默认为 None。

• resume_from_checkpoint (bool) – 如果为 True，检查点器将加载额外的检查点文件，以从之前的运行恢复训练。 默认为 False

Raises:

ValueError – 如果提供了多个检查点文件

load_checkpoint(weights_only: bool = True) → Dict[str, Any]

从文件加载 torchtune 检查点。 目前仅支持从单个文件加载。

输出 state_dict 具有以下格式，只有当 resume_from_checkpoint 为 True 时，才会存在 “model” 以外的键

>>>     {
>>>         "model": {
>>>             "key_1": weight
>>>             ...
>>>         },
>>>         "optimizer": {...},
>>>         ...
>>>     }

参数:

weights_only (bool) – 传递给 torch.load 的标志。 我们公开此标志，因为量化模型无法通过 weights_only=True 加载。

返回:

来自输入检查点的 state_dict

返回类型:

Dict[str, Any]

save_checkpoint(state_dict: Dict[str, Any], epoch: int, intermediate_checkpoint: bool = False, adapter_only: bool = False) → None

将 torchtune 检查点保存到文件。 如果 intermediate_checkpoint 为 True，则会在 _output_dir/RECIPE_STATE_DIRNAME 中创建一个包含配方状态的额外检查点文件 recipe_state.pt。 输出 state_dict 具有以下格式

>>> # Model
>>> {
>>>     "key_1": weight
>>>     ...
>>> }
>>>
>>> # Recipe state
>>> {
>>>     "optimizer": ...,
>>>     "epoch": ...,
>>>     ...
>>> }

参数:

• state_dict (Dict[str, Any]) – 包含模型和（可选）配方状态的 State dict

• epoch (int) – 当前 epoch 编号。 这已添加到检查点文件名中，以确保我们不会覆盖中间检查点文件

• intermediate_checkpoint (bool) – 如果为 True，则保存一个包含配方状态的额外检查点文件

• adapter_only (bool) – 如果为 True，则仅保存适配器权重。 默认为 False

Raises:

ValueError – 如果 adapter_only 为 True 且在 state_dict 中未找到适配器检查点。