nn.Module 用于 load_state_dict 和张量子类"

recipes/recipes/swap_tensors

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

nn.Module 中用于 load_state_dict 和张量子类的扩展点

创建日期：2024年4月19日 | 最后更新：2024年4月19日 | 最后验证：未验证

作者： Mikayla Gawarecki

本教程介绍了一个新的实用函数 torch.utils.swap_tensors，以及它在 nn.Module 中集成的两个新扩展点：

• nn.Module.to() 及相关方法

• nn.Module.load_state_dict()

注意

本教程需要 PyTorch 2.3.0 或更高版本。

torch.utils.swap_tensors

torch.utils.swap_tensors（以下简称 swap_tensors）是一个实用函数，它接收两个 Python 张量并将它们交换。

import torch
import torch.nn as nn
t1 = torch.arange(2)
t2 = torch.arange(3)
print(f"Before swapping, t1: {t1}, t2: {t2}")
torch.utils.swap_tensors(t1, t2)
print(f"After swapping, t1: {t1}, t2: {t2}")

Before swapping, t1: tensor([0, 1]), t2: tensor([0, 1, 2])
After swapping, t1: tensor([0, 1, 2]), t2: tensor([0, 1])

更具体地说，swap_tensors 会交换两个张量的 Python __class__、__dict__ 和 __slots__，以及它们关联的 at::Tensor。

在 nn.Module 中的应用

当模块外部的 Python 对象持有该模块参数的引用时，此工具对于 nn.Module 尤为重要。如果 nn.Module 原位（out of place）修改了其任何参数，持有该参数引用的对象将无法感知到变化。一个典型的例子是优化器，它持有 nn.Module 参数的引用。这会导致一个静默的正确性问题，即 optimizer.step() 会在没有报错的情况下运行，但 nn.Module 的权重实际上并未更新。

mod = torch.nn.Linear(1, 2, bias=False)
optimizer = torch.optim.SGD(mod.parameters())
print(f"weight in mod: {mod.weight}")
print(f"weight in optimizer: {optimizer.param_groups[0]['params']}")
mod.weight = torch.nn.Parameter(2 * mod.weight)
print(f"weight in mod: {mod.weight}")
print(f"weight in optimizer: {optimizer.param_groups[0]['params']}")

weight in mod: Parameter containing:
tensor([[-0.4580],
        [-0.8483]], requires_grad=True)
weight in optimizer: [Parameter containing:
tensor([[-0.4580],
        [-0.8483]], requires_grad=True)]
weight in mod: Parameter containing:
tensor([[-0.9160],
        [-1.6965]], requires_grad=True)
weight in optimizer: [Parameter containing:
tensor([[-0.4580],
        [-0.8483]], requires_grad=True)]

nn.Module.to() 及相关方法

这包括改变模块设备的方法（如 nn.Module.cpu()）、改变模块 dtype 的方法（如 nn.Module.float()），以及允许模块实例化的方法（如 nn.Module.to_empty()）。

乍一看，这些方法能够原位修改模块参数可能显得不太直观。现有的方法是使用一种自 PyTorch 早期就存在的“蹩脚”技巧。

值得注意的是，现有的方法在以下情况下无法正常工作：

• 使用 __torch_dispatch__ 子类时

• 当 param 和 new_param 的 Python type() 不同时

• 对于具有特殊 C++ 表示的张量（如稀疏张量和 XLA 张量）

在本教程的后续部分，我们将定义一个用于演示的 __torch_dispatch__ 子类 MyQuantizedLinearWeight，它代表量化的线性层权重。此子类将在本教程的其余部分用于说明。为了简洁，我们省略了大部分 __torch_dispatch__ 的实现。

aten = torch.ops.aten

class MyQuantizedLinearWeight(torch.Tensor):
    @staticmethod
    def __new__(cls, elem, scale):
        return torch.Tensor._make_wrapper_subclass(
            cls,
            elem.shape,
            dtype=elem.dtype,
            layout=elem.layout,
            device=elem.device,
            strides=elem.stride(),
            storage_offset=elem.storage_offset())

    def __init__(self, elem: torch.Tensor, scale: float):
        self.elem = elem
        self.scale = scale

    def __repr__(self):
        return f"MyQuantizedLinearWeight({self.elem}, scale={self.scale})"

    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs):
        if func in (aten.detach.default, aten._to_copy.default):
            new_elem = func(args[0].elem, *args[1:], **kwargs)
            return cls(new_elem, args[0].scale)
        # Implementations for certain ops would be added to ``OP_TABLE``.
        # We omit this for brevity.
        OP_TABLE = dict()
        if func in OP_TABLE:
          return OP_TABLE[func](func, args, kwargs)
        raise NotImplementedError(f"Unsupported function {func}")

让我们创建一个 dtype 为 torch.float32 的 nn.Linear 层，其权重是一个 MyQuantizedLinearWeight，并尝试将其转换为 torch.bfloat16。观察到权重的 dtype 如预期般发生了变化。然而，子类负载（elem）的 dtype 并没有改变。

m = nn.Linear(3, 5, dtype=torch.float32)
m.weight = torch.nn.Parameter(MyQuantizedLinearWeight(m.weight, 0.5))
print(f"Before: id(m.weight)={id(m.weight)}, id(m.bias)={id(m.bias)}")
m.bfloat16()
print(f"After: id(m.weight)={id(m.weight)}, id(m.bias)={id(m.bias)}")
print(f"m.weight.dtype: {m.weight.dtype}")
print(f"m.weight.elem.dtype: {m.weight.elem.dtype}")
print(f"m.bias.dtype: {m.bias.dtype}")

Before: id(m.weight)=139631470400096, id(m.bias)=139631470397696
After: id(m.weight)=139631470400096, id(m.bias)=139631470397696
m.weight.dtype: torch.bfloat16
m.weight.elem.dtype: torch.float32
m.bias.dtype: torch.bfloat16

为此，我们引入了一个全局配置 torch.__future__.set_swap_module_params_on_conversion，它将使用 swap_tensors 来交换模块参数，同时保留引用，以替代原来的 .data 设置。当设置此配置后，转换过程中将使用 swap_tensors，从而确保负载的 dtype 被正确转换。

torch.__future__.set_swap_module_params_on_conversion(True)
m = nn.Linear(3, 5, dtype=torch.float32)
m.weight = torch.nn.Parameter(MyQuantizedLinearWeight(m.weight, 0.5))
print(f"Before: id(m.weight)={id(m.weight)}, id(m.bias)={id(m.bias)}")
m.bfloat16()
print(f"After: id(m.weight)={id(m.weight)}, id(m.bias)={id(m.bias)}")
print(f"m.weight.dtype: {m.weight.dtype}")
print(f"m.weight.elem.dtype: {m.weight.elem.dtype}")
print(f"m.bias.dtype: {m.bias.dtype}")
torch.__future__.set_swap_module_params_on_conversion(False)

Before: id(m.weight)=139631492812048, id(m.bias)=139631264839328
After: id(m.weight)=139631492812048, id(m.bias)=139631264839328
m.weight.dtype: torch.bfloat16
m.weight.elem.dtype: torch.bfloat16
m.bias.dtype: torch.bfloat16

nn.Module.load_state_dict()

根据传递给 load_state_dict() 的 assign 关键字参数的值，有两种加载 state_dict 的方式：

• assign=False：保留 module.param 的属性，仅获取 state_dict['param_name'] 中的值。

• assign=True：保留 state_dict['param_name'] 的属性和值。

以前，这些分别通过原位 copy_ 和 __setattr__ 实现。在现有实现中，每种方法都有其局限性——assign=False 强制要求 state_dict 中的参数类型必须与模块中的参数类型相同；而 assign=True 则要求所有持有模块参数引用的对象必须在 nn.Module.load_state_dict() 之后初始化。

现在，我们通过在 load_state_dict() 中添加 swap_tensors 路径，并引入一个新的扩展点 torch.Tensor.module_load(self, other, assign=False) 来解决这两个限制。当通过上述 __future__ 启用 swap_tensors 路径时，我们可以使用 module_load 的 __torch_function__ 处理器，对 state_dict 中的值应用自定义转换。该转换的结果将与模块中的参数进行交换。

在以下示例中，我们将使用上面定义的 MyQuantizedLinearWeight 子类，演示如何在加载 state_dict 时利用这些特性将自定义量化方案应用于线性层权重。

回想一下，如果 self 或 other（在本例中为 param 或 state_dict[param_key]）是 MyQuantizedLinearWeight 子类，那么 module_load 的 __torch_function__ 处理器就会被调用。

假设我们预期 state_dict 包含普通张量，而模块包含 MyQuantizedLinearWeight 参数，我们希望将 state_dict 中的张量转换为该子类。那么我们可以这样定义 torch.Tensor.module_load 的 __torch_function__ 处理器：

@classmethod
def custom_torch_function(cls, func, types, args=(), kwargs=None):
    kwargs = {} if kwargs is None else kwargs

    if func is torch.Tensor.module_load:
        dest, src = args[0], args[1]
        assert type(dest) == cls and type(src) == torch.Tensor
        return MyQuantizedLinearWeight(src, dest.scale)
    else:
        with torch._C.DisableTorchFunctionSubclass():
                return func(*args, **kwargs)

MyQuantizedLinearWeight.__torch_function__ = custom_torch_function

首先，让我们在 meta 设备上创建一个模型的骨架，以避免实例化存储。我们将模块中的所有权重转换为 MyQuantizedLinearWeight 子类，同时保持偏置（biases）不变。

def fn(m):
    if isinstance(m, nn.Linear):
        requires_grad = m.weight.requires_grad
        m.weight = torch.nn.Parameter(
                    MyQuantizedLinearWeight(m.weight, 0.5), requires_grad=requires_grad
                   )

with torch.device("meta"):
    m = nn.Linear(3, 5)
    m.apply(fn)

然后我们可以加载 state_dict。注意我们使用了 assign=True，因为对于偏置，我们希望保留 state_dict 中张量的属性（例如，我们不希望加载后偏置位于 meta 设备上）。

torch.__future__.set_swap_module_params_on_conversion(True)
print(f"Before: id(weight)={id(m.weight)}, id(bias)={id(m.bias)}")
print(f"m.state_dict() before load_state_dict():\n {m.state_dict()}")
state_dict = nn.Linear(3, 5).state_dict()
print(f"state_dict:\n {state_dict}")
m.load_state_dict(state_dict, assign=True)
print(f"After: id(weight)={id(m.weight)}, id(bias)={id(m.bias)}")
print(f"m.state_dict() after load_state_dict():\n {m.state_dict()}")

Before: id(weight)=139631492811168, id(bias)=139631224831520
m.state_dict() before load_state_dict():
 OrderedDict([('weight', MyQuantizedLinearWeight(tensor(..., device='meta', size=(5, 3)), scale=0.5)), ('bias', tensor(..., device='meta', size=(5,)))])
state_dict:
 OrderedDict([('weight', tensor([[ 0.4500,  0.1149, -0.2731],
        [-0.1652, -0.5626, -0.5771],
        [ 0.5429,  0.2590,  0.5639],
        [ 0.3249,  0.5627, -0.0740],
        [-0.1397, -0.1811,  0.2544]])), ('bias', tensor([-0.1232,  0.0864, -0.5610, -0.3004,  0.1822]))])
After: id(weight)=139631492811168, id(bias)=139631224831520
m.state_dict() after load_state_dict():
 OrderedDict([('weight', MyQuantizedLinearWeight(tensor([[ 0.4500,  0.1149, -0.2731],
        [-0.1652, -0.5626, -0.5771],
        [ 0.5429,  0.2590,  0.5639],
        [ 0.3249,  0.5627, -0.0740],
        [-0.1397, -0.1811,  0.2544]]), scale=0.5)), ('bias', tensor([-0.1232,  0.0864, -0.5610, -0.3004,  0.1822]))])

以上是一个关于如何使用 nn.Module.load_state_dict() 中新扩展点的简单示例。人们还可以设想其他场景，例如 state_dict 中有张量子类，而模块中有普通 nn.Parameters/张量，或者两者均为张量子类。根据具体用例，我们可以定义 module_load 的 __torch_function__ 处理器以按需应用转换。

结论

在本教程中，我们了解了 swap_tensors，了解了在 nn.Module 中保留参数引用的重要性，以及如何使用受 torch.__future__.set_swap_module_params_on_conversion 控制的两个新扩展点。