huggingface

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.

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 using model.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 using cls.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()