当前位置：首页 > news >正文

LLMs-from-scratch ：PyTorch 缓冲区（Buffers）

news 2025/10/16 13:43:56

源代码链接：https://github.com/rasbt/LLMs-from-scratch/tree/main/ch03/03_understanding-buffers

《从零开始构建大型语言模型》一书的补充代码，作者：Sebastian Raschka

代码仓库：https://github.com/rasbt/LLMs-from-scratch

理解 PyTorch 缓冲区（Buffers）

本质上，PyTorch 缓冲区是与 PyTorch 模块或模型关联的张量属性，类似于参数，但与参数不同的是，缓冲区在训练过程中不会被更新。

PyTorch 中的缓冲区在处理 GPU 计算时特别有用，因为它们需要与模型参数一起在设备之间传输（比如从 CPU 到 GPU）。与参数不同，缓冲区不需要梯度计算，但它们仍然需要在正确的设备上，以确保所有计算都能正确执行。

在第3章中，我们通过 self.register_buffer 使用 PyTorch 缓冲区，这在书中只是简单解释了一下。由于这个概念和目的不是立即清楚的，这个代码笔记本提供了更长的解释和实际示例。

不使用缓冲区的示例

假设我们有以下代码，它基于第3章的代码。这个版本已经修改为不包含缓冲区。它实现了 LLM 中使用的因果自注意力机制：

import torch
import torch.nn as nnclass CausalAttentionWithoutBuffers(nn.Module):def __init__(self, d_in, d_out, context_length,dropout, qkv_bias=False):super().__init__()self.d_out = d_outself.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)self.dropout = nn.Dropout(dropout)self.mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)def forward(self, x):b, num_tokens, d_in = x.shapekeys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.transpose(1, 2)attn_scores.masked_fill_(self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)attn_weights = self.dropout(attn_weights)context_vec = attn_weights @ valuesreturn context_vec

我们可以在一些示例数据上初始化并运行该模块，如下所示：

torch.manual_seed(123)inputs = torch.tensor([[0.43, 0.15, 0.89], # Your     (x^1)[0.55, 0.87, 0.66], # journey  (x^2)[0.57, 0.85, 0.64], # starts   (x^3)[0.22, 0.58, 0.33], # with     (x^4)[0.77, 0.25, 0.10], # one      (x^5)[0.05, 0.80, 0.55]] # step     (x^6)
)batch = torch.stack((inputs, inputs), dim=0)
context_length = batch.shape[1]
d_in = inputs.shape[1]
d_out = 2ca_without_buffer = CausalAttentionWithoutBuffers(d_in, d_out, context_length, 0.0)with torch.no_grad():context_vecs = ca_without_buffer(batch)print(context_vecs)

tensor([[[-0.4519,  0.2216],[-0.5874,  0.0058],[-0.6300, -0.0632],[-0.5675, -0.0843],[-0.5526, -0.0981],[-0.5299, -0.1081]],[[-0.4519,  0.2216],[-0.5874,  0.0058],[-0.6300, -0.0632],[-0.5675, -0.0843],[-0.5526, -0.0981],[-0.5299, -0.1081]]])

到目前为止，一切都运行良好。

然而，在训练 LLM 时，我们通常使用 GPU 来加速这个过程。因此，让我们将 CausalAttentionWithoutBuffers 模块转移到 GPU 设备上。

请注意，此操作需要在配备 GPU 的环境中运行代码。

has_cuda = torch.cuda.is_available()
has_mps = torch.backends.mps.is_available()print("机器有 GPU:", has_cuda or has_mps)if has_mps:device = torch.device("mps")   # Apple Silicon GPU (Metal)
elif has_cuda:device = torch.device("cuda")  # NVIDIA GPU
else:device = torch.device("cpu")   # CPU 后备print(f"使用设备: {device}")batch = batch.to(device)
ca_without_buffer = ca_without_buffer.to(device)

机器有 GPU: True
使用设备: cuda

现在，让我们再次运行代码：

with torch.no_grad():context_vecs = ca_without_buffer(batch)print(context_vecs)

---------------------------------------------------------------------------RuntimeError                              Traceback (most recent call last)Cell In[4], line 21 with torch.no_grad():
----> 2     context_vecs = ca_without_buffer(batch)4 print(context_vecs)File d:\agent-llm2\LLMs-from-scratch\.venv\Lib\site-packages\torch\nn\modules\module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs)1734     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]1735 else:
-> 1736     return self._call_impl(*args, **kwargs)File d:\agent-llm2\LLMs-from-scratch\.venv\Lib\site-packages\torch\nn\modules\module.py:1747, in Module._call_impl(self, *args, **kwargs)1742 # If we don't have any hooks, we want to skip the rest of the logic in1743 # this function, and just call forward.1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks1745         or _global_backward_pre_hooks or _global_backward_hooks1746         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747     return forward_call(*args, **kwargs)1749 result = None1750 called_always_called_hooks = set()Cell In[1], line 23, in CausalAttentionWithoutBuffers.forward(self, x)20 values = self.W_value(x)22 attn_scores = queries @ keys.transpose(1, 2)
---> 23 attn_scores.masked_fill_(24     self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)25 attn_weights = torch.softmax(26     attn_scores / keys.shape[-1]**0.5, dim=-127 )28 attn_weights = self.dropout(attn_weights)RuntimeError: expected self and mask to be on the same device, but got mask on cpu and self on cuda:0

运行代码导致了错误。发生了什么？看起来我们试图在 GPU 上的张量和 CPU 上的张量之间进行矩阵乘法。但我们已经将模块移动到 GPU 了！？

让我们仔细检查一些张量的设备位置：

print("W_query.device:", ca_without_buffer.W_query.weight.device)
print("mask.device:", ca_without_buffer.mask.device)

W_query.device: cuda:0
mask.device: cpu

type(ca_without_buffer.mask)

torch.Tensor

如我们所见，mask 没有被移动到 GPU 上。这是因为它不是像权重（例如，W_query.weight）那样的 PyTorch 参数。

这意味着我们必须通过 .to("cuda") 手动将其移动到 GPU：

ca_without_buffer.mask = ca_without_buffer.mask.to(device)
print("mask.device:", ca_without_buffer.mask.device)

mask.device: cuda:0

让我们再次尝试我们的代码：

with torch.no_grad():context_vecs = ca_without_buffer(batch)print(context_vecs)

tensor([[[-0.4519,  0.2216],[-0.5874,  0.0058],[-0.6300, -0.0632],[-0.5675, -0.0843],[-0.5526, -0.0981],[-0.5299, -0.1081]],[[-0.4519,  0.2216],[-0.5874,  0.0058],[-0.6300, -0.0632],[-0.5675, -0.0843],[-0.5526, -0.0981],[-0.5299, -0.1081]]], device='cuda:0')

这次成功了！

然而，记住将单个张量移动到 GPU 可能很繁琐。正如我们将在下一节中看到的，使用 register_buffer 将 mask 注册为缓冲区更容易。

使用缓冲区的示例

现在让我们修改因果注意力类，将因果 mask 注册为缓冲区：

import torch
import torch.nn as nnclass CausalAttentionWithBuffer(nn.Module):def __init__(self, d_in, d_out, context_length,dropout, qkv_bias=False):super().__init__()self.d_out = d_outself.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)self.dropout = nn.Dropout(dropout)# 旧方式:# self.mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)# 新方式:self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1))def forward(self, x):b, num_tokens, d_in = x.shapekeys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.transpose(1, 2)attn_scores.masked_fill_(self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)attn_weights = self.dropout(attn_weights)context_vec = attn_weights @ valuesreturn context_vec

现在，方便的是，如果我们将模块移动到 GPU，mask 也会位于 GPU 上：

ca_with_buffer = CausalAttentionWithBuffer(d_in, d_out, context_length, 0.0)
ca_with_buffer.to(device)print("W_query.device:", ca_with_buffer.W_query.weight.device)
print("mask.device:", ca_with_buffer.mask.device)

W_query.device: cuda:0
mask.device: cuda:0

with torch.no_grad():context_vecs = ca_with_buffer(batch)print(context_vecs)

tensor([[[0.4772, 0.1063],[0.5891, 0.3257],[0.6202, 0.3860],[0.5478, 0.3589],[0.5321, 0.3428],[0.5077, 0.3493]],[[0.4772, 0.1063],[0.5891, 0.3257],[0.6202, 0.3860],[0.5478, 0.3589],[0.5321, 0.3428],[0.5077, 0.3493]]], device='cuda:0')

如上所示，将张量注册为缓冲区可以让我们的生活变得更轻松：我们不必记住手动将张量移动到像 GPU 这样的目标设备。

缓冲区和 `state_dict`

PyTorch 缓冲区相对于常规张量的另一个优势是，它们会被包含在模型的 state_dict 中
例如，考虑不使用缓冲区的因果注意力对象的 state_dict

ca_without_buffer.state_dict()

OrderedDict([('W_query.weight',tensor([[-0.2354,  0.0191, -0.2867],[ 0.2177, -0.4919,  0.4232]], device='cuda:0')),('W_key.weight',tensor([[-0.4196, -0.4590, -0.3648],[ 0.2615, -0.2133,  0.2161]], device='cuda:0')),('W_value.weight',tensor([[-0.4900, -0.3503, -0.2120],[-0.1135, -0.4404,  0.3780]], device='cuda:0'))])

mask 没有包含在上面的 state_dict 中
然而，由于将其注册为缓冲区，mask 被包含在下面的 state_dict 中

ca_with_buffer.state_dict()

OrderedDict([('mask',tensor([[0., 1., 1., 1., 1., 1.],[0., 0., 1., 1., 1., 1.],[0., 0., 0., 1., 1., 1.],[0., 0., 0., 0., 1., 1.],[0., 0., 0., 0., 0., 1.],[0., 0., 0., 0., 0., 0.]], device='cuda:0')),('W_query.weight',tensor([[-0.1362,  0.1853,  0.4083],[ 0.1076,  0.1579,  0.5573]], device='cuda:0')),('W_key.weight',tensor([[-0.2604,  0.1829, -0.2569],[ 0.4126,  0.4611, -0.5323]], device='cuda:0')),('W_value.weight',tensor([[ 0.4929,  0.2757,  0.2516],[ 0.2377,  0.4800, -0.0762]], device='cuda:0'))])

state_dict 在保存和加载训练好的 PyTorch 模型时很有用
在这个特定情况下，保存和加载 mask 可能不是特别有用，因为它在训练过程中保持不变；所以，为了演示目的，让我们假设它被修改了，所有的 1 都被改为 2：

ca_with_buffer.mask[ca_with_buffer.mask == 1.] = 2.
ca_with_buffer.mask

tensor([[0., 2., 2., 2., 2., 2.],[0., 0., 2., 2., 2., 2.],[0., 0., 0., 2., 2., 2.],[0., 0., 0., 0., 2., 2.],[0., 0., 0., 0., 0., 2.],[0., 0., 0., 0., 0., 0.]], device='cuda:0')

然后，如果我们保存并加载模型，我们可以看到 mask 以修改后的值被恢复

torch.save(ca_with_buffer.state_dict(), "model.pth")new_ca_with_buffer = CausalAttentionWithBuffer(d_in, d_out, context_length, 0.0)
new_ca_with_buffer.load_state_dict(torch.load("model.pth"))new_ca_with_buffer.mask

tensor([[0., 2., 2., 2., 2., 2.],[0., 0., 2., 2., 2., 2.],[0., 0., 0., 2., 2., 2.],[0., 0., 0., 0., 2., 2.],[0., 0., 0., 0., 0., 2.],[0., 0., 0., 0., 0., 0.]])

如果我们不使用缓冲区，这就不成立：

ca_without_buffer.mask[ca_without_buffer.mask == 1.] = 2.torch.save(ca_without_buffer.state_dict(), "model.pth")new_ca_without_buffer = CausalAttentionWithoutBuffers(d_in, d_out, context_length, 0.0)
new_ca_without_buffer.load_state_dict(torch.load("model.pth"))new_ca_without_buffer.mask

tensor([[0., 1., 1., 1., 1., 1.],[0., 0., 1., 1., 1., 1.],[0., 0., 0., 1., 1., 1.],[0., 0., 0., 0., 1., 1.],[0., 0., 0., 0., 0., 1.],[0., 0., 0., 0., 0., 0.]])

查看全文

http://www.dtcms.com/a/487979.html