序列化

序列化和反序列化是一個人們特別關心的問題，尤其是在我們將 torchao 與其他庫整合時。在這裡，我們想描述 torchao 最佳化（量化或稀疏化）模型的序列化和反序列化是如何工作的。

序列化和反序列化流程

這是序列化和反序列化流程

import copy
import tempfile
import torch
from torchao.utils import get_model_size_in_bytes
from torchao.quantization.quant_api import (
    quantize_,
    Int4WeightOnlyConfig,
)

class ToyLinearModel(torch.nn.Module):
    def __init__(self, m=64, n=32, k=64):
        super().__init__()
        self.linear1 = torch.nn.Linear(m, n, bias=False)
        self.linear2 = torch.nn.Linear(n, k, bias=False)

    def example_inputs(self, batch_size=1, dtype=torch.float32, device="cpu"):
        return (torch.randn(batch_size, self.linear1.in_features, dtype=dtype, device=device),)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x)
        return x

dtype = torch.bfloat16
m = ToyLinearModel(1024, 1024, 1024).eval().to(dtype).to("cuda")
print(f"original model size: {get_model_size_in_bytes(m) / 1024 / 1024} MB")

example_inputs = m.example_inputs(dtype=dtype, device="cuda")
quantize_(m, Int4WeightOnlyConfig())
print(f"quantized model size: {get_model_size_in_bytes(m) / 1024 / 1024} MB")

ref = m(*example_inputs)
with tempfile.NamedTemporaryFile() as f:
    torch.save(m.state_dict(), f)
    f.seek(0)
    state_dict = torch.load(f)

with torch.device("meta"):
    m_loaded = ToyLinearModel(1024, 1024, 1024).eval().to(dtype)

# `linear.weight` is nn.Parameter, so we check the type of `linear.weight.data`
print(f"type of weight before loading: {type(m_loaded.linear1.weight.data), type(m_loaded.linear2.weight.data)}")
m_loaded.load_state_dict(state_dict, assign=True)
print(f"type of weight after loading: {type(m_loaded.linear1.weight), type(m_loaded.linear2.weight)}")

res = m_loaded(*example_inputs)
assert torch.equal(res, ref)

序列化最佳化模型時會發生什麼？

要序列化最佳化模型，我們只需要呼叫 torch.save(m.state_dict(), f)，因為在 torchao 中，我們使用張量子類來表示不同的資料型別或支援量化和稀疏性等不同的最佳化技術。因此，最佳化後，唯一改變的是權重張量被更改為最佳化後的權重張量，而模型結構完全沒有改變。例如

原始浮點模型 state_dict

{"linear1.weight": float_weight1, "linear2.weight": float_weight2}

量化模型 state_dict

{"linear1.weight": quantized_weight1, "linear2.weight": quantized_weight2, ...}

量化模型的大小通常會比原始浮點模型小，但這取決於您使用的具體技術和實現。您可以使用 torchao.utils.get_model_size_in_bytes 工具函式列印模型大小。特別是對於上面使用 Int4WeightOnlyConfig 量化的示例，我們可以看到尺寸減小了約 4 倍。

original model size: 4.0 MB
quantized model size: 1.0625 MB

反序列化最佳化模型時會發生什麼？

要反序列化最佳化模型，我們可以在 meta 裝置上初始化浮點模型，然後使用 assign=True 透過 model.load_state_dict 載入最佳化後的 state_dict。

with torch.device("meta"):
    m_loaded = ToyLinearModel(1024, 1024, 1024).eval().to(dtype)

print(f"type of weight before loading: {type(m_loaded.linear1.weight), type(m_loaded.linear2.weight)}")
m_loaded.load_state_dict(state_dict, assign=True)
print(f"type of weight after loading: {type(m_loaded.linear1.weight), type(m_loaded.linear2.weight)}")

我們在 meta 裝置上初始化模型的原因是避免初始化原始浮點模型，因為原始浮點模型可能不適合我們要用於推理的裝置。

在 m_loaded.load_state_dict(state_dict, assign=True) 中發生的事情是，相應的權重（例如 m_loaded.linear1.weight）會使用 state_dict 中的張量進行更新，這些張量是最佳化後的張量子類例項（例如 int4 AffineQuantizedTensor）。要使其工作，不需要依賴 torchao。

我們還可以透過檢查權重張量的型別來驗證權重是否已正確載入

type of weight before loading: (<class 'torch.Tensor'>, <class 'torch.Tensor'>)
type of weight after loading: (<class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>, <class 'torchao.dtypes.affine_quantized_tensor.AffineQuantizedTensor'>)