从零开始逐步指导开发者构建自己的大型语言模型（LLM）学习笔记- 第4章 GPT 的大语言模型架构 练习题

Chapter 4 Exercise solutions

In [1]:

from importlib.metadata import version

import torch
print("torch version:", version("torch"))
torch version: 2.4.0

Exercise 4.1: Parameters in the feed forward versus attention module

In [2]:

from gpt import TransformerBlock

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}

block = TransformerBlock(GPT_CONFIG_124M)

• The results above are for a single transformer block
• 可以选择乘以 12 以涵盖 1.24 亿参数的 GPT 模型中的所有变压器模块。

Exercise 4.2: Initialize larger GPT models

• GPT2-small (the 124M configuration we already implemented):

  • "emb_dim" = 768
  • "n_layers" = 12
  • "n_heads" = 12

• GPT2-medium:

  • "emb_dim" = 1024
  • "n_layers" = 24
  • "n_heads" = 16

• GPT2-large:

  • "emb_dim" = 1280
  • "n_layers" = 36
  • "n_heads" = 20

• GPT2-XL:

  • "emb_dim" = 1600
  • "n_layers" = 48
  • "n_heads" = 25

In [5]:

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}


def get_config(base_config, model_name="gpt2-small"):
    GPT_CONFIG = base_config.copy()

    if model_name == "gpt2-small":
        GPT_CONFIG["emb_dim"] = 768
        GPT_CONFIG["n_layers"] = 12
        GPT_CONFIG["n_heads"] = 12

    elif model_name == "gpt2-medium":
        GPT_CONFIG["emb_dim"] = 1024
        GPT_CONFIG["n_layers"] = 24
        GPT_CONFIG["n_heads"] = 16

    elif model_name == "gpt2-large":
        GPT_CONFIG["emb_dim"] = 1280
        GPT_CONFIG["n_layers"] = 36
        GPT_CONFIG["n_heads"] = 20

    elif model_name == "gpt2-xl":
        GPT_CONFIG["emb_dim"] = 1600
        GPT_CONFIG["n_layers"] = 48
        GPT_CONFIG["n_heads"] = 25

    else:
        raise ValueError(f"Incorrect model name {model_name}")

    return GPT_CONFIG


def calculate_size(model): # based on chapter code
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total number of parameters: {total_params:,}")

    total_params_gpt2 =  total_params - sum(p.numel() for p in model.out_head.parameters())
    print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")
    
    # Calculate the total size in bytes (assuming float32, 4 bytes per parameter)
    total_size_bytes = total_params * 4
    
    # Convert to megabytes
    total_size_mb = total_size_bytes / (1024 * 1024)
    
    print(f"Total size of the model: {total_size_mb:.2f} MB")

calculate_size 的函数，该函数用于计算模型的参数数量和模型的总大小（以MB为单位）。以下是对这段代码的详细解释：

1. def calculate_size(model): # based on chapter code

   • 这是函数的定义，它接受一个参数 model，这个参数是一个PyTorch模型。# based on chapter code 是一个注释，说明这个函数是基于某一章节的代码编写的。

2. total_params = sum(p.numel() for p in model.parameters())

   • 这行代码计算模型中所有参数的总数。model.parameters() 返回一个迭代器，遍历模型的所有参数。p.numel() 是一个方法，用于返回张量 p 中的元素数量。sum() 函数用于将所有参数的元素数量相加，得到总的参数数量。

3. print(f"Total number of parameters: {total_params:,}")

   • 这行代码打印出模型的总参数数量。f"{total_params:,}" 是一个格式化字符串，其中 total_params 是要格式化的变量，: 表示格式化选项，, 表示使用逗号作为千位分隔符。

4. total_params_gpt2 = total_params - sum(p.numel() for p in model.out_head.parameters())

   • 这行代码计算考虑权重绑定（weight tying）后的可训练参数数量。在某些语言模型中，输出层的权重矩阵与输入层的嵌入矩阵共享参数，这称为权重绑定。model.out_head.parameters() 返回输出层的参数迭代器。通过减去输出层的参数数量，可以得到不包括输出层权重的可训练参数数量。

5. print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")

   • 这行代码打印出考虑权重绑定后的可训练参数数量。

6. total_size_bytes = total_params * 4

   • 这行代码计算模型的总大小（以字节为单位）。假设每个参数使用32位浮点数（4字节）表示，因此总大小等于总参数数量乘以4。

7. total_size_mb = total_size_bytes / (1024 * 1024)

   • 这行代码将总大小从字节转换为兆字节（MB）。1024 * 1024 是1MB的字节数。

8. print(f"Total size of the model: {total_size_mb:.2f} MB")

   • 这行代码打印出模型的总大小（以MB为单位）。{total_size_mb:.2f} 是一个格式化字符串，其中 total_size_mb 是要格式化的变量，:.2f 表示保留两位小数。

 

In [6]:

from gpt import GPTModel


for model_abbrev in ("small", "medium", "large", "xl"):
    model_name = f"gpt2-{model_abbrev}"
    CONFIG = get_config(GPT_CONFIG_124M, model_name=model_name)
    model = GPTModel(CONFIG)
    print(f"\n\n{model_name}:")
    calculate_size(model)

gpt2-small:
Total number of parameters: 163,009,536
Number of trainable parameters considering weight tying: 124,412,160
Total size of the model: 621.83 MB


gpt2-medium:
Total number of parameters: 406,212,608
Number of trainable parameters considering weight tying: 354,749,440
Total size of the model: 1549.58 MB


gpt2-large:
Total number of parameters: 838,220,800
Number of trainable parameters considering weight tying: 773,891,840
Total size of the model: 3197.56 MB


gpt2-xl:
Total number of parameters: 1,637,792,000
Number of trainable parameters considering weight tying: 1,557,380,800
Total size of the model: 6247.68 MB

Exercise 4.3: Using separate dropout parameters

In [7]:

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate_emb": 0.1,        # NEW: dropout for embedding layers
    "drop_rate_attn": 0.1,       # NEW: dropout for multi-head attention  
    "drop_rate_shortcut": 0.1,   # NEW: dropout for shortcut connections  
    "qkv_bias": False
}

In [8]:

import torch.nn as nn
from gpt import MultiHeadAttention, LayerNorm, FeedForward


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"], 
            dropout=cfg["drop_rate_attn"], # NEW: dropout for multi-head attention
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate_shortcut"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed-forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x


class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate_emb"]) # NEW: dropout for embedding layers

        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

GPTModel 类的 forward 方法，它定义了模型的前向传播过程。以下是对这段代码的详细解释：

1. def forward(self, in_idx):

   • 这是 forward 方法的定义，它接受一个输入张量 in_idx，表示输入的标记索引。

2. batch_size, seq_len = in_idx.shape

   • 这行代码获取输入张量 in_idx 的形状，并将其分解为 batch_size（批次大小）和 seq_len（序列长度）。

3. tok_embeds = self.tok_emb(in_idx)

   • 这行代码通过调用 self.tok_emb（一个嵌入层）对输入的标记索引 in_idx 进行嵌入操作，得到标记嵌入 tok_embeds。嵌入层将每个标记索引映射到一个固定大小的向量表示。

4. pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))

   • 这行代码生成位置嵌入 pos_embeds。torch.arange(seq_len, device=in_idx.device) 创建一个从 0 到 seq_len-1 的张量，表示序列中每个位置的索引。然后，通过调用 self.pos_emb（另一个嵌入层）对这些位置索引进行嵌入操作，得到位置嵌入。

5. x = tok_embeds + pos_embeds # Shape [batch_size, num_tokens, emb_size]

   • 这行代码将标记嵌入 tok_embeds 和位置嵌入 pos_embeds 相加，得到最终的输入表示 x。这个表示包含了每个标记的嵌入信息和其在序列中的位置信息。

6. x = self.drop_emb(x)

   • 这行代码对输入表示 x 应用一个丢弃层（dropout layer），以防止过拟合。丢弃层在训练过程中随机将一些元素置为零，从而增加模型的鲁棒性。

7. x = self.trf_blocks(x)

   • 这行代码将输入表示 x 传递给一系列的 TransformerBlock 模块，这些模块构成了模型的主体部分。每个 TransformerBlock 包含一个多头注意力机制（MultiHeadAttention）和一个前馈神经网络（FeedForward），以及相应的归一化层（LayerNorm）和丢弃层（Dropout）。

8. x = self.final_norm(x)

   • 这行代码对经过 TransformerBlock 处理后的表示 x 应用最后一个归一化层，进一步稳定模型的训练。

9. logits = self.out_head(x)

   • 这行代码将归一化后的表示 x 传递给一个线性层 self.out_head，将其映射到词汇表大小的输出空间，得到每个标记的对数几率（logits）。

10. return logits

    • 最后，forward 方法返回计算得到的对数几率 logits，这些对数几率可以用于后续的损失计算和预测。

In [9]:

import torch

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)