从零开始实现大语言模型（十一）：构建大语言模型GPTModel

1. 前言

OpenAI的GPT系列大语言模型GPTModel共包含两个Embedding层（Token Embedding和Positional Embedding），一个Dropout层，多个结构完全相同的Transformer Block，一个Layer Normalization以及一个Linear层。

对输入文本做tokenization，将输入文本转换成token ID列表，并输入Embedding层，可以得到维度为embedding_dim的Embedding向量序列。将经过Dropout的Embedding向量序列依次输入多个Transformer Block，前一个Transformer Block的输出作为后一个Transformer Block的输入，最后一个Transformer Block输出的维度为embedding_dim的向量序列输入Layer Normalization做变换处理，并将经过变换处理的向量输入Linear层。Linear层将各个token对应向量的维度变换拓展至vocabulary_size维，得到大语言模型GPTModel的输出。

本文将介绍由多头注意力模块、Layer Normalization、Dropout、残差连接、前馈神经网络与GELU激活函数构成的Transformer Block，实现OpenAI的GPT系列大语言模型GPTModel，并对大语言模型参数进行度量分析。

2. Transformer Block

Transformer Block由多头注意力模块、Layer Normalization、Dropout、残差连接、前馈神经网络与GELU激活函数构成。如下图所示，Transformer Block可以被划分为两个子模块，第一个子模块的核心是多头注意力模块MultiHeadAttention，其每个注意力头均为CausalAttention。第二个子模块的核心是前馈神经网络模块FeedForward，其共包含两个Linear层以及一个GELU激活函数。

Layer Normalization会对每个子模块的输入张量做变换，使shape为[batch_size, num_tokens, embedding_dim]的输入张量的embedding_dim维度数据的均值为0，方差为1。对第一个子模块中多头注意力模块的输出张量做Dropout，并使用残差连接将其输出与输入直接相加，即可得到第一个子模块的输出。将第一个子模块的输出作为第二个子模块的输入，对第二个子模块中前馈神经网络模块的输出做Dropout，并使用残差连接将其输出与输入直接相加，得到的shape为[batch_size, num_tokens, embedding_dim]的张量即为Transformer Block的输出。

Transformer Block中这种对多头注意力模块及前馈神经网络模块的输入张量做Layer Normalization的方法被称为Pre-LayerNorm。文章Attention is all you need中提出的Transformer模型中使用Layer Normalization对多头注意力模块及前馈神经网络模块的输出张量做变换的方法被称为Post-LayerNorm。深度学习实践已经证明使用Post-LayerNorm可以使模型训练过程更稳定，模型更容易收敛。OpenAI的GPT系列大语言模型以及Meta的LLaMA等现在流行的绝大部分大语言模型均使用了这种Post-LayerNorm方法。

构建Transformer Block可以实现一个继承自torch.nn.Module的TransformerBlock类，并在其__init__方法中初始化一个从零开始实现大语言模型（七）：多头注意力机制所述的MultiHeadAttention类对象，一个从零开始实现大语言模型（九）：前馈神经网络与GELU激活函数所述的FeedForward类对象，两个从零开始实现大语言模型（八）：Layer Normalization所述的LayerNorm类对象，以及一个对多头注意力模块及前馈神经网络模块输出张量做Dropout的torch.nn.Dropout对象。forward方法将x输入第一个子模块，并将第一个子模块的输出作为第二个子模块的输入，第二个子模块可以计算得到Transformer Block的输出。具体代码如下所示：

import torch
# from [从零开始实现大语言模型（七）：多头注意力机制] import MultiHeadAttention
# from [从零开始实现大语言模型（八）：Layer Normalization] import LayerNorm
# from [从零开始实现大语言模型（九）：前馈神经网络与GELU激活函数] import GELU, FeedForward

torch.manual_seed(123)


class TransformerBlock(torch.nn.Module):
    def __init__(self, embedding_dim, context_len, num_heads, dropout, qkv_bias=False):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=embedding_dim,
            d_out=embedding_dim,
            context_len=context_len,
            dropout=dropout,
            num_heads=num_heads, 
            qkv_bias=qkv_bias
        )
        self.ff = FeedForward(embedding_dim)
        self.norm1 = LayerNorm(embedding_dim)
        self.norm2 = LayerNorm(embedding_dim)
        self.dropout = torch.nn.Dropout(dropout)
 
    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.dropout(x)
        x = x + shortcut
 
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.dropout(x)
        x = x + shortcut
        return x

可以使用如下代码实例化上述TransformerBlock类对象，初始化一个shape为[2, 4, 768]的张量batch_example，并将其输入实例化的TransformerBlock类对象，打印输出张量及其shape：

batch_example = torch.rand(2, 4, 768)
trf_block = TransformerBlock(
    embedding_dim=768,
    context_len=1024,
    num_heads=12,
    dropout=0.1,
    qkv_bias=False
)
out = trf_block(batch_example)
print(out.shape)
print(out)

执行上面代码，打印结果如下：

torch.Size([2, 4, 768])
tensor([[[-0.0055,  0.0972, -0.1122,  ...,  1.2889,  0.2623,  0.6685],
         [ 0.0023, -0.2369,  0.1720,  ...,  0.5952,  0.2497,  0.7447],
         [ 0.4673,  0.4472,  0.1791,  ...,  1.2525,  0.3045,  0.7750],
         [ 0.0662,  0.7224,  0.9206,  ...,  0.4790,  0.7428,  0.7015]],

        [[ 0.3622,  1.2144,  0.5221,  ...,  0.1854,  0.0111, -0.5034],
         [-0.0225,  0.7789,  0.2770,  ...,  0.1734,  0.5419,  0.1143],
         [ 0.7425,  0.4013,  0.3211,  ...,  0.3268,  0.7523, -0.1642],
         [ 0.5745,  0.6241,  0.4410,  ...,  1.1963,  1.2650,  0.2243]]],
       grad_fn=<AddBackward0>)

3. 构建大语言模型GPTModel

OpenAI的GPT系列大语言模型由两个Embedding层（Token Embedding和Positional Embedding），一个Dropout层，多个结构完全相同的Transformer Block，一个Layer Normalization以及一个Linear层构成。如下图所示，对输入文本Every effort moves you做tokenization，将输入文本转换成包含4个token ID的列表，并输入GPT-2 small版本的Embedding层，可以得到维度为768的Embedding向量序列。对Embedding层的输出张量做Dropout，并依次输入12个Transformer Block，前一个Transformer Block的输出作为后一个Transformer Block的输入，最后一个Transformer Block输出的维度为768的向量序列输入Layer Normalization做变换处理，并将经过变换处理的向量输入Linear层。Linear层将各个token对应向量的维度变换拓展至50257维，得到大语言模型GPT-2 small版本的输出。

如下面的代码所示，构建OpenAI的GPT系列大语言模型可以实现一个继承自torch.nn.Module的GPTModel类，并在其__init__方法中使用torch.nn.Embedding分别初始化将token ID及其位置映射成Embedding向量的两个Embedding层，一个对Embedding层输出张量做Dropout的torch.nn.Dropout对象，一个包含num_layers个Transformer Block的torch.nn.Sequential对象，一个从零开始实现大语言模型（八）：Layer Normalization所述的LayerNorm类对象，以及一个将向量维度由embedding_dim维变换拓展至vocabulary_size维的不包含bias的torch.nn.Linear层。

forward方法将x输入Token Embedding层，可以将输入token ID映射成token Embedding向量。各个token对应的位置ID输入Positional Embedding层，可以将输入token对应位置映射成positional Embedding向量。最后将token Embedding张量与positional Embedding张量直接相加，得到Embedding层的输出。使用在__init__方法中初始化的torch.nn.Dropout对象对Embedding层输出张量做Dropout，并输入包含num_layers个Transformer Block的torch.nn.Sequential对象。使用Layer Normalization对最后一个Transformer Block的输出张量做变换，并将经过变换后的张量输入Linear层，得到大语言模型GPTModel的输出：

class GPTModel(torch.nn.Module):
    def __init__(self, embedding_dim, num_layers, num_heads, context_len, vocabulary_size, dropout, qkv_bias=False):
        super().__init__()
        self.tok_emb = torch.nn.Embedding(vocabulary_size, embedding_dim)
        self.pos_emb = torch.nn.Embedding(context_len, embedding_dim)
        self.drop_emb = torch.nn.Dropout(dropout)

        self.trf_blocks = torch.nn.Sequential(
            *[TransformerBlock(embedding_dim, context_len, num_heads, dropout, qkv_bias) for _ in range(num_layers)]
        )

        self.final_norm = LayerNorm(embedding_dim)
        self.out_linear = torch.nn.Linear(embedding_dim, vocabulary_size, bias=False)

    def forward(self, x):
        batch_size, sequence_len = x.shape

        tok_embeds = self.tok_emb(x)
        pos_embeds = self.pos_emb(torch.arange(sequence_len, device=x.device))
        x = tok_embeds + pos_embeds

        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_linear(x)
        return logits

可以使用如下代码创建上图所示的GPT-2 small版本大语言模型，初始化一个shape为[2, 4]的张量batch_example，并将其输入GPT-2 small版本大语言模型，打印输出张量及其shape：

embedding_dim = 768
num_layers = 12
num_heads = 12
context_len = 1024
vocabulary_size = 50257
dropout = 0.1
qkv_bias = False

batch_example = torch.randint(0, vocabulary_size, [2, 4])
gpt2_small = GPTModel(
    embedding_dim=embedding_dim,
    num_layers=num_layers,
    num_heads=num_heads,
    context_len=context_len,
    vocabulary_size=vocabulary_size,
    dropout=dropout,
    qkv_bias=qkv_bias
)
logits = gpt2_small(batch_example)
print(logits.shape)
print(logits)

执行上面代码，打印结果如下：

torch.Size([2, 4, 50257])
tensor([[[ 0.1942, -0.5203, -0.6836,  ..., -0.2983,  0.1723,  0.4619],
         [ 0.0668, -0.0233, -0.1506,  ...,  0.5952, -0.2082,  0.7683],
         [ 0.1064, -0.6752,  0.8593,  ...,  0.1733,  0.0344,  0.3323],
         [-0.5962,  0.7507,  0.1938,  ...,  0.6638,  0.5102, -0.1514]],

        [[-0.0182,  0.4398, -0.3168,  ...,  0.4502,  0.1311, -0.1457],
         [ 0.1600,  0.4703,  0.6581,  ...,  0.6061, -0.9757,  0.6010],
         [ 0.4367, -0.4918,  0.2199,  ..., -0.0993,  0.7368,  0.5137],
         [-0.0530, -0.0455,  0.4997,  ...,  0.8101,  0.4854,  0.0509]]],
       grad_fn=<UnsafeViewBackward0>)

输出张量logits的shape为[2, 4, 50257]。其中第一个维度2表明输入中包含两个样本数据，第二个维度4表示输入的每个样本数据包含4个tokens，第三个维度50257表示GPT-2 small版本大语言模型所使用的词汇表中共有50257个不同的token。

4. 大语言模型参数度量分析

如下表所示，OpenAI的GPT系列大语言模型中参数量为124M的GPT-2 small版本的Embedding向量维度embedding_dim=768，多头注意力机制中注意力头数num_heads=12，Transformer Block的数量num_layers=12，支持的最大上下文长度context_len=1024，使用的词汇表中共包含vocabulary_size=50257个不同的tokens。

如下面的代码所示，可以使用.numel(number of elements)方法计算3中所述大语言模型gpt2_small的参数数量：

total_params = sum(p.numel() for p in gpt2_small.parameters())
print(f"Total number of parameters: {total_params:,}")

执行上面代码，打印结果如下：

Total number of parameters: 163,009,536

上面代码打印出大语言模型GPT-2 small版本中参数数量为163M。在原始的GPT-2模型中使用了一种被称为Weight Tying的架构设计方法，这种架构设计方法会让Token Embedding层与Linear输出层共享权重参数矩阵。可以使用下面的代码分别打印输出Token Embedding层及Linear输出层权重参数矩阵的shape：

print("Token embedding layer shape:", gpt2_small.tok_emb.weight.shape)
print("Output linear layer shape:", gpt2_small.out_linear.weight.shape)

执行上面代码，打印结果如下：

Token embedding layer shape: torch.Size([50257, 768])
Output linear layer shape: torch.Size([50257, 768])

可以使用如下代码剔除Linear输出层参数，并重新统计大语言模型gpt2_small的参数数量：

total_params_2 = total_params - sum(p.numel() for p in gpt2_small.out_linear.parameters())
print(f"Number of trainable parameters considering weight tying: {total_params_2:,}")

执行上面代码，打印结果如下：

Number of trainable parameters considering weight tying: 124,412,160

Weight Tying可以降低训练神经网络模型所需的GPT内存资源，但是深度学习实践已经证明Token Embedding层与Linear输出层不共享权重参数矩阵可以使大语言模型的效果更好。现在流行的大语言模型一般不会采用这种Weight Tying架构设计方法。

假设3中所述大语言模型gpt2_small的每个参数类型均为torch.float32，即每个参数均需要4 bytes的存储空间。可以使用如下代码计算存储该参数量为163M的大语言模型所需内存大小：

total_size_bytes = total_params * 4
total_size_mb = total_size_bytes / (1024 * 1024)
print(f"Total size of the gpt2_small: {total_size_mb:.2f} MB")

执行上面代码，打印结果如下：

Total size of the gpt2_small: 621.83 MB

上述大语言模型GPT-2 small版本的参数量仅为163M，但是存储该模型共需要621.83MB存储空间。训练大语言模型时，需要在内存或显存中存储的数据主要包括以下5个部分：

• 大语言模型参数：如上述大语言模型GPT-2 small版本共需要621.83MB存储空间
• 大语言模型参数梯度：训练大语言模型时，每个可训练参数均会由一个对应梯度。梯度通常与参数的数据类型相同，因此训练上述大语言模型GPT-2 small版本同样需要621.83MB存储空间来存储梯度
• 大语言模型各层输出值：在前向传播过程中，大语言模型的每一层均会输出一个张量。这些张量同样需要存储到内存或显存中，以用于在反向传播时计算梯度。大语言模型各层输出张量所需存储空间大小取决于输入数据中tokens的数量以及batch_size的大小。一般来说，存储大语言模型各层输出值所需存储空间会比大语言模型参数存储所需空间大得多
• 优化器参数：某些优化器（如Adam）需要额外存储每个参数梯度的均值等数据。以Adam优化器为例，其需要额外占用$621.83\times 2=1243.66\text{MB}$存储空间
• 缓存和其他临时数据：大语言模型训练过程中的缓存及临时变量的值

ChatGPT共包含175B参数，现在流行的大语言模型的参数量都是数十亿起步。可想而知，训练大语言模型的计算代价有多大，对计算设备的要求有多高。根据Lambda Labs的评估数据，在单个V100 datacenter GPU训练GPT-3需要355年，在单个消费级RTX 8000 GPU上训练GPT-3需要665年。

5. 结束语

多少事，从来急，天地转，光阴迫。一万年太久，只争朝夕！从2024年2月4日落笔从零开始实现大语言模型系列专栏第一篇文章从零开始实现大语言模型（一）：概述，至今天2024年5月5日，完成了第一阶段——构建大语言模型部分内容。

天若有情天亦老，人间正道是沧桑。世界上任何一件伟大的事情，都不是轻轻松松，敲锣打鼓就能实现的。

宜将剩勇追穷寇，不可沽名学霸王。开始写下一个阶段——预训练大语言模型部分内容了！