huggingface — Model Optimizer 0.27.1 (original) (raw)

ModelOpt plugin for enabling automatic save/restore of ModelOpt state for HuggingFace models.

Functions

enable_huggingface_checkpointing()

Enables automatic save/restore of ModelOpt state with HuggingFace checkpointing APIs.

ModelOpt automatically saves modelopt_state to save_directory/modelopt_state.pth when a Huggingface model is saved usingmodel.save_pretrained(save_directory).

Conversely, ModelOpt restores the saved state from pretrained_model_name_or_path/modelopt_state.pth if it exists when a Huggingface model is loaded usingcls.from_pretrained(pretrained_model_name_or_path).

This function should be called once in the program before loading/saving any HuggingFace models.

Here is an example usage:

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()