PyTorch 中的 Transformer 源码 [LLM]

1
_shape_t = int | list[int] | Size
2

3
class LayerNorm(Module):
4

5
  def __init__(
6
        self,
7
    # 定义要归一化的维度的大小
8
        normalized_shape: _shape_t,
9
    # 公式中的 ε
10
        eps: float = 1e-5,
11
    # 使用 γ 和 β
12
        elementwise_affine: bool = True,
13
        bias: bool = True,
14
        device=None,
15
        dtype=None,
16
    ) -> None:
17
        factory_kwargs = {"device": device, "dtype": dtype}
18
        super().__init__()
19
    # 给成员赋值
20
        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
21
        self.eps = eps
22
        self.elementwise_affine = elementwise_affine
23
        if self.elementwise_affine:
24
            self.weight = Parameter(
25
                torch.empty(self.normalized_shape, **factory_kwargs)
26
            )
27
            if bias:
28
                self.bias = Parameter(
29
                    torch.empty(self.normalized_shape, **factory_kwargs)
30
                )
31
            else:
32
                self.register_parameter("bias", None)
33
        else:
34
            self.register_parameter("weight", None)
35
            self.register_parameter("bias", None)
36

37
        self.reset_parameters()
38

39
    def reset_parameters(self) -> None:
40
        if self.elementwise_affine:
41
      # γ 初始值是 1
42
            init.ones_(self.weight)
43
            if self.bias is not None:
44
        # β 初始值是 0
45
                init.zeros_(self.bias)
46

47
    def forward(self, input: Tensor) -> Tensor:
48
    # 将参数传给 functional.layer_norm()
49
        return F.layer_norm(
50
            input, self.normalized_shape, self.weight, self.bias, self.eps
51
        )

1
def layer_norm(
2
    input: Tensor,
3
    normalized_shape: list[int],
4
    weight: Tensor | None = None,
5
    bias: Tensor | None = None,
6
    eps: float = 1e-5,
7
) -> Tensor:
8
    r"""Apply Layer Normalization for last certain number of dimensions.
9

10
    See :class:`~torch.nn.LayerNorm` for details.
11
    """
12
  # 如果自定义了张量 (variadic: 可变参数)，将调用交给自定义类的 __torch_function__ 方法
13
    if has_torch_function_variadic(input, weight, bias):
14
        return handle_torch_function(
15
            layer_norm,
16
            (input, weight, bias),
17
            input,
18
            normalized_shape,
19
            weight=weight,
20
            bias=bias,
21
            eps=eps,
22
        )
23
  # 否则调用 C++ 实现的高性能版本
24
    return torch.layer_norm(
25
        input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled
26
    )

1
std::tuple<Tensor, Tensor, Tensor> layer_norm_cuda(
2
    const Tensor& input,
3
    IntArrayRef normalized_shape,
4
    const std::optional<Tensor>& weight_opt /* optional */,
5
    const std::optional<Tensor>& bias_opt /* optional */,
6
    double eps) {
7

8
  c10::MaybeOwned<Tensor> weight_maybe_owned =
9
      at::borrow_from_optional_tensor(weight_opt);
10
 // 解引用获取 Tensor 引用
11
  const Tensor& weight = *weight_maybe_owned;
12

13
  c10::MaybeOwned<Tensor> bias_maybe_owned =
14
      at::borrow_from_optional_tensor(bias_opt);
15
  const Tensor& bias = *bias_maybe_owned;
16

17
  // auto 让编译器自动推断变量的类型
18
  auto M_N = _check_layer_norm_inputs(input, normalized_shape, weight, bias);
19
  // M_N 的 第一个和第二个元素
20
  auto M = M_N.first;
21
  auto N = M_N.second;
22
  auto X = input.expect_contiguous();
23
  auto gamma = weight.expect_contiguous();
24
  auto beta = bias.expect_contiguous();
25

26
  // 输出张量 Y
27
  Tensor Y = at::native::empty_like(...);
28
  // 均值张量和 std 张量
29
  Tensor mean = at::empty(...);
30
  Tensor rstd = at::empty(...);
31
  // M > 0 时调用 kernel
32
  if (M > 0) {
33
    LayerNormKernelImpl(*X, *gamma, *beta, M, N, eps, &Y, &mean, &rstd);
34
  }
35
  ...
36
  // 计算得到的均值和 std 通过 view 改变维度与输入对齐进行广播
37
  mean = mean.view(stat_shape);
38
  rstd = rstd.view(stat_shape);
39
  // 返回 mean, std 用于 bp
40
  return std::make_tuple(std::move(Y), std::move(mean), std::move(rstd));
41
}
42

43
void LayerNormKernelImpl(...) {
44
  // 根据输入张量的数据类型，自动生成并调用针对该类型的模板代码
45
  AT_DISPATCH_FLOATING_TYPES_AND2(
46
      at::ScalarType::Half,
47
      at::ScalarType::BFloat16,
48
      X.scalar_type(),
49
      "LayerNormKernelImpl",
50
      [&]() {
51
        using acc_t = acc_type<scalar_t, true>;
52
    // 执行计算
53
        LayerNormKernelImplInternal<scalar_t, acc_t>(
54
            X, gamma, beta, M, N, static_cast<acc_t>(eps), Y, mean, rstd);
55
      });
56
}
57

58
// 最后在 LayerNormKernelImplInternal 中调用 RowwiseMomentsCUDAKernel 和 LayerNormForwardCUDAKernel 计算

1
class TransformerEncoderLayer(Module):
2

3
  def __init__(
4
        self,
5
        d_model: int,
6
        nhead: int,
7
    # FFN 的维度
8
        dim_feedforward: int = 2048,
9
        dropout: float = 0.1,
10
    # 激活函数
11
        activation: str | Callable[[Tensor], Tensor] = F.relu,
12
        layer_norm_eps: float = 1e-5,
13
        batch_first: bool = False,
14
    # Pre-LN 还是 Post-LN
15
        norm_first: bool = False,
16
        bias: bool = True,
17
        device=None,
18
        dtype=None,
19
    ) -> None:
20
    factory_kwargs = {"device": device, "dtype": dtype}
21
        super().__init__()
22
    # 自注意力
23
        self.self_attn = MultiheadAttention(...)
24

25
    # FFN
26
        self.linear1 = Linear(d_model, dim_feedforward, bias=bias, **factory_kwargs)
27
        self.dropout = Dropout(dropout)
28
        self.linear2 = Linear(dim_feedforward, d_model, bias=bias, **factory_kwargs)
29

30
    # LN
31
        self.norm_first = norm_first
32
    # norm1 用于自注意力，norm2 用于 FFN
33
        self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
34
        self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, bias=bias, **factory_kwargs)
35
        self.dropout1 = Dropout(dropout)
36
        self.dropout2 = Dropout(dropout)
37
    # ... (略去激活函数处理细节)
38
    # 定义激活函数
39
    self.activation = activation
40

41
  def forward(
42
        self,
43
        src: Tensor,
44
    # 注意力掩码
45
        src_mask: Tensor | None = None,
46
    # pad 掩码
47
        src_key_padding_mask: Tensor | None = None,
48
        is_causal: bool = False,
49
    ) -> Tensor:
50
    # 将传入的 mask 格式规范化
51
    src_key_padding_mask = F._canonical_mask(...)
52
    src_mask = F._canonical_mask(...)
53
    # 计算逻辑
54
    x = src
55
    # Pre-LN
56
        if self.norm_first:
57
      # 注意力计算
58
            x = x + self._sa_block(
59
                self.norm1(x), src_mask, src_key_padding_mask, is_causal=is_causal
60
            )
61
            x = x + self._ff_block(self.norm2(x))
62
        else:
63
      # Post-LN
64
            x = self.norm1(
65
                x
66
                + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal)
67
            )
68
      # FFN
69
            x = self.norm2(x + self._ff_block(x))
70

71
        return x
72

73
  def _sa_block(...) -> Tensor:
74
    # 上面定义了 self.self_attn = MultiheadAttention(...)
75
        x = self.self_attn(
76
      # 自注意力 qkv 都是 x
77
            x,
78
            x,
79
            x,
80
            attn_mask=attn_mask,
81
            key_padding_mask=key_padding_mask,
82
            need_weights=False,
83
            is_causal=is_causal,
84
        )[0]
85
        return self.dropout1(x)
86

87
  def _ff_block(self, x: Tensor) -> Tensor:
88
        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
89
        return self.dropout2(x)

1
class TransformerEncoder(Module):
2
  def __init__(
3
        self,
4
        encoder_layer: "TransformerEncoderLayer",
5
        num_layers: int,
6
        norm: Module | None = None,
7
        enable_nested_tensor: bool = True,
8
        mask_check: bool = True,
9
    ) -> None:
10
    super().__init__()
11
    # 创建 num_layers 个编码器层的深拷贝
12
        self.layers = _get_clones(encoder_layer, num_layers)
13
        self.num_layers = num_layers
14
        self.norm = norm
15
        self.mask_check = mask_check
16

17
  def forward(
18
        self,
19
        src: Tensor,
20
        mask: Tensor | None = None,
21
        src_key_padding_mask: Tensor | None = None,
22
        is_causal: bool | None = None,
23
    ) -> Tensor:
24

25
  seq_len = _get_seq_len(src, batch_first)
26

27
  for mod in self.layers:
28
      # 循环调用 model 的 forward 方法
29
            output = mod(
30
                output,
31
                src_mask=mask,
32
                is_causal=is_causal,
33
                src_key_padding_mask=src_key_padding_mask_for_layers,
34
            )
35

36
  return output