序列化

序列化與反序列化是人們關心的重要問題，特別是在我們將 torchao 與其他函式庫整合時。在此，我們想要描述序列化和反序列化如何用於 torchao 最佳化的（量化或稀疏化）模型。

序列化與反序列化流程

以下是序列化與反序列化流程

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

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, int4_weight_only())
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 實用函數列印模型大小，特別是對於上述使用 int4_weight_only 量化的範例，我們可以發現大小減少了約 4 倍

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

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

為了反序列化最佳化模型，我們可以在 meta 裝置中初始化浮點模型，然後使用 model.load_state_dict 並使用 assign=True 來載入最佳化的 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'>)