huggingface
ModelOpt插件用于启用HuggingFace模型的自动保存/恢复ModelOpt状态。
函数
启用与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()