保存与恢复

ModelOpt优化方法，如剪枝、量化、稀疏化、蒸馏等，会修改原始模型的架构和权重。因此，保存和恢复模型架构以及新的权重非常重要，以便在下游任务中使用修改后的模型。

本指南描述了正确保存和恢复ModelOpt修改模型的多种选项。

保存ModelOpt模型

对模型所做的修改被捕获在模型的ModelOpt状态中，如modelopt_state所示。ModelOpt支持将架构修改与模型权重一起保存或分开保存。

保存ModelOpt状态和模型权重

mto.save 保存 ModelOpt 状态以及新的模型权重（即 Pytorch state_dict），这些可以在以后用于正确恢复模型。

以下是如何保存一个经过ModelOpt修改的模型的示例：

import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq

# Initialize the original model and set the original weights
model = ...

# Apply ModelOpt modifications to the model. For example, quantization
model = mtq.quantize(model, config, forward_loop)

# Save the model weights and the `modelopt_state` to 'modelopt_model.pth'
mto.save(model, "modelopt_model.pth")

分别保存ModelOpt状态和模型权重

如果你想使用自己的自定义方法保存模型权重，而不是使用ModelOpt的mto.save()保存，你应该单独保存modelopt_state。这样你可以稍后从保存的modelopt_state中正确恢复模型架构，如分别恢复ModelOpt状态和模型权重所示。

以下是如何单独保存 modelopt_state 的示例：

import torch
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq

# Initialize the original model and set the original weights
model = ...

# Apply ModelOpt modifications to the model. For example, quantization
model = mtq.quantize(model, config, forward_loop)

# Save the `modelopt_state` to 'modelopt_state.pth'
torch.save(mto.modelopt_state(model), "modelopt_state.pth")

# Save the model weights separately with your method
custom_method_to_save_model_weights(model)

注意

尽管上述示例展示了使用量化的ModelOpt修改，相同的工作流程适用于其他ModelOpt修改，如剪枝、稀疏性、蒸馏等，以及它们的组合。

恢复ModelOpt模型

恢复ModelOpt状态和模型权重

mto.restore 恢复模型的 modelopt_state() 和使用 mto.save() 保存的模型权重，如保存一个经过ModelOpt修改的模型所示。恢复后的模型可以用于推理或进一步的训练和优化。

以下是一个恢复经过ModelOpt修改的模型的示例：

import modelopt.torch.opt as mto

# Initialize the original model
model = ...

# Restore the model architecture and weights after applying ModelOpt modifications
mto.restore(model, "modelopt_model.pth")

# Use the restored model for inference or further training / optimization

分别恢复ModelOpt状态和模型权重

或者，如果您像单独保存模型优化状态中所示的那样单独保存了modelopt_state，您可以使用保存的modelopt_state恢复经过ModelOpt修改的模型架构。在此步骤之后，应单独加载经过ModelOpt修改后的模型权重。

以下是使用保存的modelopt_state恢复ModelOpt修改的模型架构的示例工作流程：

import torch
import modelopt.torch.opt as mto

# Initialize the original model
model = ...

# Restore the model architecture using the saved `modelopt_state`
modelopt_state = torch.load("modelopt_state.pth")
model = mto.restore_from_modelopt_state(model, modelopt_state)

# Load the model weights separately after restoring the model architecture
custom_method_to_load_model_weights(model)

使用Huggingface检查点API进行ModelOpt保存/恢复

ModelOpt 支持在使用 Huggingface 库（如 transformers 和 diffusers）中的 save_pretrained 和 from_pretrained API 时自动保存和恢复修改后的模型。

要启用此功能，您需要在加载/保存任何HuggingFace模型之前在程序中调用一次 mto.enable_huggingface_checkpointing()。

以下是如何使用 Huggingface API 启用 ModelOpt 保存/恢复的示例：

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

...

# Enable automatic ModelOpt save/restore with
# Huggingface checkpointing APIs `save_pretrained` and `from_pretrained`
mto.enable_huggingface_checkpointing()

# Load the original Huggingface model
model = AutoModelForCausalLM.from_pretrained(model_path)

# Apply ModelOpt modifications to the model. For example, quantization
model = mtq.quantize(model, config, forward_loop)

# Save the ModelOpt-modified model architecture and weights using Huggingface APIs
model.save_pretrained(f"ModelOpt_{model_path}")

如上所述保存的模型可以使用Huggingface的from_pretrained API进行恢复。在加载模型之前，不要忘记调用mto.enable_huggingface_checkpointing()。这只需要在程序中执行一次。

请参见以下示例：

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

...

# Enable automatic ModelOpt save/restore with huggingface checkpointing APIs
# This needs to be done only once in the program
mto.enable_huggingface_checkpointing()

# Load the ModelOpt-modified model architecture and weights using Huggingface APIs
model = AutoModelForCausalLM.from_pretrained(f"ModelOpt_{model_path}")