huggingface

ModelOpt插件用于启用HuggingFace模型的自动保存/恢复ModelOpt状态。

函数

enable_huggingface_checkpointing

启用与HuggingFace检查点API的自动保存/恢复ModelOpt状态。

enable_huggingface_checkpointing()

启用与HuggingFace检查点API的自动保存/恢复ModelOpt状态。

当使用model.save_pretrained(save_directory)保存Huggingface模型时,ModelOpt会自动将modelopt_state保存到save_directory/modelopt_state.pth

相反,如果存在,ModelOpt 会从 pretrained_model_name_or_path/modelopt_state.pth 恢复保存的状态 当使用 cls.from_pretrained(pretrained_model_name_or_path) 加载 Huggingface 模型时。

此函数应在程序中加载/保存任何HuggingFace模型之前调用一次。

以下是一个使用示例:

from transformers import AutoModelForCausalLM
import modelopt.torch.opt as mto

# Enable ModelOpt save/restore for HuggingFace models
# This only needs to be called once in the program.
mto.enable_huggingface_checkpointing()

# Instantiate a HuggingFace model, modelopt_state will be automatically loaded if it exists.
model = AutoModelForCausalLM.from_pretrained(model_path).cuda()