4 Implementing a GPT model from Scratch To Generate Text

• 本章节包含

  1. 编写一个类似于GPT的大型语言模型（LLM），这个模型可以被训练来生成类似人类的文本。
  2. Normalizing layer activations to stabilize neural network training
  3. 在深度神经网络中添加shortcut connections，以更有效地训练模型
  4. 实现 Transformer 模块以创建不同规模的 GPT 模型
  5. 计算 GPT 模型的参数数量及其存储需求

  在上一章中，学习了多头注意力机制并对其进行了编码，它是LLMs的核心组件之一。在本章中，将编写 LLM 的其他构建块，并将它们组装成类似 GPT 的模型

  ​

  ---

4.1 Coding an LLM architecture

• 诸如GPT和Llama等模型,基于原始Transformer架构中的decoder部分,因此，这些LLM通常被称为"decoder-like" LLMs，与传统的深度学习模型相比，LLM规模更大，这主要归因于它们庞大的参数数量，而非代码量。因为它的许多组件都是重复的，下图提供了类似 GPT LLM 的自上而下视图

  ​

  本章将详细构建一个最小规模的GPT-2模型（1.24亿参数），并展示如何加载预训练权重以兼容更大规模的模型。

• 1.24亿参数GPT-2模型的配置细节包括：

  GPT_CONFIG_124M = {
      "vocab_size": 50257,    # Vocabulary size
      "context_length": 1024, # Context length
      "emb_dim": 768,         # Embedding dimension
      "n_heads": 12,          # Number of attention heads
      "n_layers": 12,         # Number of layers
      "drop_rate": 0.1,       # Dropout rate
      "qkv_bias": False       # Query-Key-Value bias
  }
  

  我们使用简短的变量名以避免后续代码行过长

  1. "vocab_size" 词汇表大小，由 BPE tokenizer 支持，值为 50,257。
  2. "context_length" 模型的最大输入标记数量，通过 positional embeddings 实现。
  3. "emb_dim" token输入的嵌入大小，将每个token转换为 768 维向量。
  4. "n_heads" 多头注意力机制中的注意力头数量。
  5. "n_layers" 是模型中 transformer 块的数量
  6. "drop_rate" 是 dropout 机制的强度，第 3 章讨论过；0.1 表示在训练期间丢弃 10% 的隐藏单元以缓解过拟合
  7. "qkv_bias" 决定多头注意力机制（第 3 章）中的 Linear 层在计算查询（Q）、键（K）和值（V）张量时是否包含偏置向量；我们将禁用此选项，这是现代 LLMs 的标准做法；然而，我们将在第 5 章将 OpenAI 的预训练 GPT-2 权重加载到我们的重新实现时重新讨论这一点。

• 下图中的方框展示了我们为实现最终 GPT 架构所需处理的各个概念的顺序。我们将从第一步开始，即一个我们称为 DummyGPTModel 的 GPT 骨架占位符：

  ​

  import torch
  import torch.nn as nn
  
  
  class DummyGPTModel(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"])
          
          # Use a placeholder for TransformerBlock
          self.trf_blocks = nn.Sequential(
              *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])])
          
          # Use a placeholder for LayerNorm
          self.final_norm = DummyLayerNorm(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
          x = self.drop_emb(x)
          x = self.trf_blocks(x)
          x = self.final_norm(x)
          logits = self.out_head(x)
          return logits
  
  
  class DummyTransformerBlock(nn.Module):
      def __init__(self, cfg):
          super().__init__()
          # A simple placeholder
  
      def forward(self, x):
          # This block does nothing and just returns its input.
          return x
  
  
  class DummyLayerNorm(nn.Module):
      def __init__(self, normalized_shape, eps=1e-5):
          super().__init__()
          # The parameters here are just to mimic the LayerNorm interface.
  
      def forward(self, x):
          # This layer does nothing and just returns its input.
          return x
  

  1. DummyGPTModel：简化版的 GPT 类模型，使用 PyTorch 的 nn.Module 实现。
  2. 模型组件：包括标记嵌入、位置嵌入、丢弃层、变换器块、层归一化和线性输出层。
  3. 配置字典：配置通过 Python 字典传入,如 GPT_CONFIG_124M，用于传递模型配置。
  4. forward 方法：描述数据从输入到输出的完整流程。计算嵌入 → 应用 dropout → 通过 transformer blocks 处理 → 应用归一化 → 生成 logits。
  5. 占位符：DummyLayerNorm 和 DummyTransformerBlock 是待实现的组件。

• 数据流动：下图提供了 GPT 模型中数据流动的高层次概述。

  ​

  使用 tiktoken 分词器对由 GPT 模型的两个文本输入组成的批次进行分词：

  import tiktoken
  
  tokenizer = tiktoken.get_encoding("gpt2")
  
  batch = []
  
  txt1 = "Every effort moves you"
  txt2 = "Every day holds a"
  
  batch.append(torch.tensor(tokenizer.encode(txt1)))
  batch.append(torch.tensor(tokenizer.encode(txt2)))
  batch = torch.stack(batch, dim=0)
  print(batch)
  
  """输出"""
  tensor([[6109, 3626, 6100,  345],
          [6109, 1110, 6622,  257]]
        )
  

  接下来，我们初始化一个包含 1.24 亿参数的 DummyGPTModel 实例，并将 tokenized batch 输入其中。

  torch.manual_seed(123)
  model = DummyGPTModel(GPT_CONFIG_124M)
  
  logits = model(batch)
  print("Output shape:", logits.shape)
  print(logits)
  
  """输出"""
  Output shape: torch.Size([2, 4, 50257])
  tensor([[[-0.9289,  0.2748, -0.7557,  ..., -1.6070,  0.2702, -0.5888],
           [-0.4476,  0.1726,  0.5354,  ..., -0.3932,  1.5285,  0.8557],
           [ 0.5680,  1.6053, -0.2155,  ...,  1.1624,  0.1380,  0.7425],
           [ 0.0447,  2.4787, -0.8843,  ...,  1.3219, -0.0864, -0.5856]],
  
          [[-1.5474, -0.0542, -1.0571,  ..., -1.8061, -0.4494, -0.6747],
           [-0.8422,  0.8243, -0.1098,  ..., -0.1434,  0.2079,  1.2046],
           [ 0.1355,  1.1858, -0.1453,  ...,  0.0869, -0.1590,  0.1552],
           [ 0.1666, -0.8138,  0.2307,  ...,  2.5035, -0.3055, -0.3083]]],
         grad_fn=<UnsafeViewBackward0>)
  

  1. 输出张量：输出张量有两行，分别对应两个文本样本。每个文本样本由 4 个标记组成；每个标记是一个 50,257 维的向量，这与标记器的词汇表大小一致。
  2. 嵌入维度：50,257 维对应词汇表中的唯一标记，后处理阶段将其转换回 token IDs 并解码为单词。

  ---

4.2 Normalizing activations with layer normalization

通过层归一化（Layer Normalization）对激活值进行归一化处理。

• Layer normalization (LayerNorm)：将激活值中心化到均值为 0，归一化方差为 1，稳定训练并加速收敛。

  应用位置：

  1. transformer block 中的 multi-head attention module 前后。

  2. 最终输出层之前。

  下图提供了LayerNormalization的直观概述

  ​

  从一个小例子看看LayerNormalization发生了什么

  torch.manual_seed(123)
  
  batch_example = torch.randn(2, 5) 
  
  layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
  out = layer(batch_example)
  print(out)
  print(out.shape)
  
  # 计算均值和方差
  mean = out.mean(dim=-1, keepdim=True)
  var = out.var(dim=-1, keepdim=True)
  
  print("Mean:\n", mean)
  print("Variance:\n", var)
  
  out_norm = (out - mean) / torch.sqrt(var)
  print("Normalized layer outputs:\n", out_norm)
  
  mean = out_norm.mean(dim=-1, keepdim=True)
  var = out_norm.var(dim=-1, keepdim=True)
  print("Mean:\n", mean)
  print("Variance:\n", var)
  
  """输出"""
  tensor([[0.2260, 0.3470, 0.0000, 0.2216, 0.0000, 0.0000],
          [0.2133, 0.2394, 0.0000, 0.5198, 0.3297, 0.0000]],
         grad_fn=<ReluBackward0>)
  
  torch.Size([2, 6])
  
  Mean:
   tensor([[0.1324],
          [0.2170]], grad_fn=<MeanBackward1>)
  
  Variance:
   tensor([[0.0231],
          [0.0398]], grad_fn=<VarBackward0>)
  
  Normalized layer outputs:
   tensor([[ 0.6159,  1.4126, -0.8719,  0.5872, -0.8719, -0.8719],
          [-0.0189,  0.1121, -1.0876,  1.5173,  0.5647, -1.0876]],
         grad_fn=<DivBackward0>)
  
  Mean:
   tensor([[9.9341e-09],
          [0.0000e+00]], grad_fn=<MeanBackward1>)
  Variance:
   tensor([[1.0000],
          [1.0000]], grad_fn=<VarBackward0>)
  

  归一化会独立应用于两个输入（行）中的每一个；使用 dim=-1 表示在最后一个维度（在本例中为特征维度）上进行计算，而不是在行维度上进行计算。

  ​

  关闭科学计数法

  torch.set_printoptions(sci_mode=False) #关闭科学计数法
  print("Mean:\n", mean)
  print("Variance:\n", var)
  
  """输出"""
  Mean:
   tensor([[    0.0000],
          [    0.0000]], grad_fn=<MeanBackward1>)
  Variance:
   tensor([[1.0000],
          [1.0000]], grad_fn=<VarBackward0>)
  

• LayerNorm 类实现：基于归一化思路，实现一个 LayerNorm 类，稍后我们可以在 GPT 模型中使用它

  class LayerNorm(nn.Module):
      def __init__(self, emb_dim):
          super().__init__()
          self.eps = 1e-5
          self.scale = nn.Parameter(torch.ones(emb_dim))
          self.shift = nn.Parameter(torch.zeros(emb_dim))
  
      def forward(self, x):
          mean = x.mean(dim=-1, keepdim=True)
          var = x.var(dim=-1, keepdim=True, unbiased=False)
          norm_x = (x - mean) / torch.sqrt(var + self.eps)
          return self.scale * norm_x + self.shift
  

  层归一化公式（上面的例子中$\gamma =1$ 、 $\beta =0$、$ϵ=0$）
  $$LayerNorm(x_i)=\gamma \cdot \frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$
  其中

  1. $\mu 、{\sigma }^2$ 分别x在layer维度上的均值和方差

  2. $\gamma 、\beta$ 是可学习的缩放平移参数

  3. $ϵ$ 是一个小常数，用于防止除零错误。

  scale 和 shift：可训练参数，用于在归一化后调整数据的缩放和偏移。

  有偏方差：在上述方差计算中，设置 unbiased=False,意味着使用公式 $\frac{{\sum }_i(x-\overline{x})}{n}$，不包含贝塞尔校正。其中 n 是样本大小（此处为特征或列的数量）；该公式不包含贝塞尔校正（即在分母中使用 n-1），因此提供的是方差的有偏估计。（对于 LLMs，嵌入维度 n 非常大，使用 n 和 n-1 之间的差异可以忽略不计，GPT-2 是在归一化层中使用有偏方差进行训练的，因此为了与后续章节中加载的预训练权重兼容，我们也采用了这一设置。）

  ln = LayerNorm(emb_dim=5)
  out_ln = ln(batch_example)
  mean = out_ln.mean(dim=-1, keepdim=True)
  var = out_ln.var(dim=-1, unbiased=False, keepdim=True)
  
  print("Mean:\n", mean)
  print("Variance:\n", var)
  
  """输出"""
  Mean:
   tensor([[    -0.0000],
          [     0.0000]], grad_fn=<MeanBackward1>)
  Variance:
   tensor([[1.0000],
          [1.0000]], grad_fn=<VarBackward0>)
  

• 所以、本节至此，我们介绍了实现GPT架构所需的构建块之一，如下图中打勾的部分

  ​

---

4.3 Implementing a feed forward network with GELU activations

• 本节即将实现子模块，用于transformer block（变换器块）的一部分。为此，我们需要从激活函数开始。

  深度学习中，ReLU因其简单和有效而被广泛使用。但在大语言模型中，还使用了GELU和SwiGLU这两种更复杂、平滑的激活函数，它们结合了高斯和sigmoid门控，提升了模型性能，与ReLU的简单分段线性不同。

• GELU（Hendrycks 和 Gimpel，2016）可以通过多种方式实现；其精确版本定义为 GELU(x)=x⋅Φ(x)，其中 Φ(x) 是标准高斯分布的累积分布函数。

  在实践中，通常会实现一种计算成本更低的近似版本：
  $$\text{GELU}(x)\approx 0.5\cdot x\cdot \left(1+\tanh \left[\sqrt{\frac{2}{\pi }}\cdot \left(x+0.044715\cdot x^3\right)\right]\right)$$
  （原始的 GPT-2 模型也是使用此近似版本进行训练的）。

  class GELU(nn.Module):
      def __init__(self):
          super().__init__()
  
      def forward(self, x):
          return 0.5 * x * (1 + torch.tanh(
              torch.sqrt(torch.tensor(2.0 / torch.pi)) * 
              (x + 0.044715 * torch.pow(x, 3))
          ))
      
  import matplotlib.pyplot as plt
  
  gelu, relu = GELU(), nn.ReLU()
  
  # Some sample data
  x = torch.linspace(-3, 3, 100)
  y_gelu, y_relu = gelu(x), relu(x)
  
  plt.figure(figsize=(8, 3))
  for i, (y, label) in enumerate(zip([y_gelu, y_relu], ["GELU", "ReLU"]), 1):
      plt.subplot(1, 2, i)
      plt.plot(x, y)
      plt.title(f"{label} activation function")
      plt.xlabel("x")
      plt.ylabel(f"{label}(x)")
      plt.grid(True)
  
  plt.tight_layout()
  plt.show()
  

  ​

  如上图所示

  1. ReLU：分段线性函数，正输入直接输出，负输入输出零。

     ReLU的局限性： 零处有尖锐拐角，可能增加优化难度；对负输入输出零，限制了负输入神经元的作用。

  2. GELU：平滑非线性函数，近似 ReLU，对负值具有非零梯度（除了在约-0.75处之外）

     GELU的优势：平滑性带来更好的优化特性；对负输入输出较小的非零值，使负输入神经元仍能贡献学习过程。

• 接下来让我们使用 GELU 函数来实现小型神经网络模块 FeedForward，稍后我们将在 LLM 的转换器块中使用它：

  class FeedForward(nn.Module):
      def __init__(self, cfg):
          super().__init__()
          self.layers = nn.Sequential(
              nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
              GELU(),
              nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
          )
  
      def forward(self, x):
          return self.layers(x)
  

  上述的前馈模块是包含两个线性层和一个GELU激活函数的小神经网络，在1.24亿参数的GPT模型中，用于处理嵌入大小为768的令牌批次。

  print(GPT_CONFIG_124M["emb_dim"])
  
  """输出"""
  768
  

  ​

  初始化一个新的前馈网络（FeedForward）模块，其 token 嵌入大小为 768，并向其输入一个包含 2 个样本且每个样本有 3 个 token 的批次输入。我们可以看到，输出张量的形状与输入张量的形状相同：

  ffn = FeedForward(GPT_CONFIG_124M)
  
  # input shape: [batch_size, num_token, emb_size]
  x = torch.rand(2, 3, 768) 
  out = ffn(x)
  print(out.shape)
  
  """输出"""
  torch.Size([2, 3, 768])
  

  本节的前馈模块对模型学习和泛化至关重要。它通过内部扩展嵌入维度到更高空间，如下图所示，然后应用GELU激活，最后收缩回原维度，以探索更丰富的表示空间。

  前馈神经网络中层输出的扩展和收缩的图示。首先，输入值从 768 个值扩大到 4 倍，达到 3072 个值。然后，第二层将 3072 个值压缩回 768 维表示。

• 本节至此，现在已经实现了 下图中LLM 的大部分构建块（打勾的部分）

  ​

---

4.4 Adding shortcut connections

• 接下来，让我们讨论 shortcut connections（快捷连接）背后的概念，也称为 skip connections（跳跃连接）或 residual connections（残差连接）。最初，在残差网络（ResNet）中提出，用于缓解梯度消失问题。梯度消失问题指的是梯度（在训练过程中指导权重更新）在向后传播通过各层时逐渐变小，导致难以有效训练较早的层，如下图所示

  ​

  对比一个由 5 层组成的深度神经网络，左侧没有快捷连接，右侧带有快捷连接。快捷连接涉及将某一层的输入与其输出相加，从而有效地创建一条绕过某些层的替代路径

  工作原理：

  1. 创建更短的梯度路径，跳过中间层。
  2. 通过将某一层的输出与后面某一层的输出相加实现。

• 前向方法中添加快捷连接

  class ExampleDeepNeuralNetwork(nn.Module):
      def __init__(self, layer_sizes, use_shortcut):
          super().__init__()
          self.use_shortcut = use_shortcut
          self.layers = nn.ModuleList([
              nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), GELU()),
              nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), GELU()),
              nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), GELU()),
              nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), GELU()),
              nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), GELU())
          ])
  
      def forward(self, x):
          for layer in self.layers:
              # Compute the output of the current layer
              layer_output = layer(x)
              # Check if shortcut can be applied
              if self.use_shortcut and x.shape == layer_output.shape:
                  x = x + layer_output
              else:
                  x = layer_output
          return x
  

  该代码实现了一个包含 5 层的深度神经网络，每层由一个 Linear layer（线性层）和一个 GELU activation function（GELU 激活函数）组成。在前向传播过程中，我们迭代地将输入传递到各层，如果 self.use_shortcut 属性设置为 True，则可以选择性地添加上面图中的 shortcut connections（快捷连接）。

  def print_gradients(model, x):
      # Forward pass
      output = model(x)
      target = torch.tensor([[0.]])
  
      # Calculate loss based on how close the target
      # and output are
      loss = nn.MSELoss()
      loss = loss(output, target)
      
      # Backward pass to calculate the gradients
      loss.backward()
  
      for name, param in model.named_parameters():
          if 'weight' in name:
              # Print the mean absolute gradient of the weights
              print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")
  

  接下来，我们实现一个计算模型向后传递中的梯度的函数：

  def print_gradients(model, x):
      # Forward pass
      output = model(x)
      target = torch.tensor([[0.]])
  
      # Calculate loss based on how close the target
      # and output are
      loss = nn.MSELoss()
      loss = loss(output, target)
      
      # Backward pass to calculate the gradients
      loss.backward()
  
      for name, param in model.named_parameters():
          if 'weight' in name:
              # Print the mean absolute gradient of the weights
              print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")
  

  在前面的代码中，我们定义了一个损失函数来计算模型输出与目标值（如 0）的接近程度，并通过调用 loss.backward() 自动计算每一层的损失梯度。使用 model.named_parameters() 可以遍历权重参数，例如对于 3×3 的权重矩阵，计算其 3×3 梯度值的平均绝对梯度，从而得到每一层的单一梯度值，便于比较各层梯度。.backward() 方法的优势在于自动完成梯度计算，无需手动实现数学过程，极大地简化了深度神经网络的训练和使用。

  接下来我们首先打印没有使用shortcut的梯度

  # 未使用shortcut
  layer_sizes = [3, 3, 3, 3, 3, 1]  
  sample_input = torch.tensor([[1., 0., -1.]])
  
  torch.manual_seed(123)
  model_without_shortcut = ExampleDeepNeuralNetwork(
      layer_sizes, use_shortcut=False
  )
  print_gradients(model_without_shortcut, sample_input)
  
  """输出"""
  layers.0.0.weight has gradient mean of 0.00020173587836325169
  layers.1.0.weight has gradient mean of 0.0001201116101583466
  layers.2.0.weight has gradient mean of 0.0007152041653171182
  layers.3.0.weight has gradient mean of 0.001398873864673078
  layers.4.0.weight has gradient mean of 0.005049646366387606
  

  接着打印使用shortcut的梯度

  # 使用shortcut
  torch.manual_seed(123)
  model_with_shortcut = ExampleDeepNeuralNetwork(
      layer_sizes, use_shortcut=True
  )
  print_gradients(model_with_shortcut, sample_input)
  
  """输出"""
  layers.0.0.weight has gradient mean of 0.22169792652130127
  layers.1.0.weight has gradient mean of 0.20694106817245483
  layers.2.0.weight has gradient mean of 0.32896995544433594
  layers.3.0.weight has gradient mean of 0.2665732502937317
  layers.4.0.weight has gradient mean of 1.3258541822433472
  

  根据上面的输出结果可以看出，shortcut connections（快捷连接）防止了梯度在早期层（如 layer.0）中消失，确保梯度的有效传播。

---

4.5 Connecting attention and linear layers in a transformer block

• 本节将实现 transformer 块，这是 GPT 和其他 LLM 架构的核心组件。在 1.24 亿参数的 GPT-2 中，该块重复多次，集成了多头注意力、层归一化、dropout、前馈层和 GELU 激活函数等概念，如下图所示。下一节会将其整合到 GPT 架构中。

  ​

• 创建TransformerBlock

  import tiktoken
  import torch
  import torch.nn as nn
  from torch.utils.data import Dataset, DataLoader
  
  class MultiHeadAttention(nn.Module):
      def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
          super().__init__()
          assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
  
          self.d_out = d_out
          self.num_heads = num_heads
          self.head_dim = d_out // num_heads  # Reduce the projection dim to match desired output dim
  
          self.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.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
          self.dropout = nn.Dropout(dropout)
          self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))
  
      def forward(self, x):
          b, num_tokens, d_in = x.shape
  
          keys = self.W_key(x)  # Shape: (b, num_tokens, d_out)
          queries = self.W_query(x)
          values = self.W_value(x)
  
          # We implicitly split the matrix by adding a `num_heads` dimension
          # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
          keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
          values = values.view(b, num_tokens, self.num_heads, self.head_dim)
          queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
  
          # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
          keys = keys.transpose(1, 2)
          queries = queries.transpose(1, 2)
          values = values.transpose(1, 2)
  
          # Compute scaled dot-product attention (aka self-attention) with a causal mask
          attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head
  
          # Original mask truncated to the number of tokens and converted to boolean
          mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
  
          # Use the mask to fill attention scores
          attn_scores.masked_fill_(mask_bool, -torch.inf)
  
          attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
          attn_weights = self.dropout(attn_weights)
  
          # Shape: (b, num_tokens, num_heads, head_dim)
          context_vec = (attn_weights @ values).transpose(1, 2)
  
          # Combine heads, where self.d_out = self.num_heads * self.head_dim
          context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
          context_vec = self.out_proj(context_vec)  # optional projection
  
          return context_vec
  
      
  GPT_CONFIG_124M = {
      "vocab_size": 50257,    # Vocabulary size
      "context_length": 1024, # Context length
      "emb_dim": 768,         # Embedding dimension
      "n_heads": 12,          # Number of attention heads
      "n_layers": 12,         # Number of layers
      "drop_rate": 0.1,       # Dropout rate
      "qkv_bias": False       # Query-Key-Value bias
  }
  
  class LayerNorm(nn.Module):
      def __init__(self, emb_dim):
          super().__init__()
          self.eps = 1e-5
          self.scale = nn.Parameter(torch.ones(emb_dim))
          self.shift = nn.Parameter(torch.zeros(emb_dim))
  
      def forward(self, x):
          mean = x.mean(dim=-1, keepdim=True)
          var = x.var(dim=-1, keepdim=True, unbiased=False)
          norm_x = (x - mean) / torch.sqrt(var + self.eps)
          return self.scale * norm_x + self.shift
  
  class GELU(nn.Module):
      def __init__(self):
          super().__init__()
  
      def forward(self, x):
          return 0.5 * x * (1 + torch.tanh(
              torch.sqrt(torch.tensor(2.0 / torch.pi)) * 
              (x + 0.044715 * torch.pow(x, 3))
          ))
      
  class FeedForward(nn.Module):
      def __init__(self, cfg):
          super().__init__()
          self.layers = nn.Sequential(
              nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
              GELU(),
              nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
          )
  
      def forward(self, x):
          return self.layers(x)
  
  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"],
              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"])
  
      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
  
  

  使用我们之前定义的 GPT_CONFIG_124M 字典，让我们实例化一个转换器块并为其提供一些示例数据：

  # 实例化
  torch.manual_seed(123)
  
  x = torch.rand(2, 4, 768)  # Shape: [batch_size, num_tokens, emb_dim]
  block = TransformerBlock(GPT_CONFIG_124M)
  output = block(x)
  
  print("Input shape:", x.shape)
  print("Output shape:", output.shape)
  
  """输出"""
  Input shape: torch.Size([2, 4, 768])
  Output shape: torch.Size([2, 4, 768])
  

  transformer block在其输出中维护输入维度，这表明transformer 架构处理数据序列而不改变它们在整个网络中的形状，通过保留输入序列的形状（长度和特征大小）并重新编码每个输出向量以整合全局上下文信息，使其能够有效应用于各种序列到序列任务，同时保持输入与输出的一对一关系。

• 本节至此，现在已经实现 GPT 所需的所有构建块

  ​

---

4.6 Coding the GPT model

• 本章从宏观视角介绍了 DummyGPTModel，使用占位符表示其构建模块，随后用真实的 TransformerBlock 和 LayerNorm 类替换占位符，组装出完整的 1.24 亿参数 GPT-2 模型，并计划在后续章节进行预训练和加载 OpenAI 的预训练权重，同时通过下图 展示了结合本章所有概念的 GPT-2 整体结构。通过将变换器块插入到本章开头的架构中并重复 12 次（以 124M GPT-2 模型为例），我们构建了一个完整且可用的 GPT 架构。

  ​

  从底部开始，tokenized text 首先被转换为 token embeddings，然后通过 positional embeddings 进行增强。这些信息组合成一个张量，随后通过一系列 transformer 块（如中心部分所示，每个块包含多头注意力机制和前馈神经网络层，并应用了 dropout 和层归一化），这些块堆叠在一起，重复 12 次，我们通过 GPT_CONFIG_124M 字典中的“n_layers”条目指定。（在拥有 15.42 亿个参数的最大 GPT-2 模型中，该transformer块重复了 36 次）。

• 上图架构的对应代码实现

  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"])
          
          # 创建 TransformerBlock 模块的顺序堆栈
          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
      
  

  使用 124M 参数模型的配置，我们现在可以用随机初始权重实例化这个 GPT 模型

  # 初始化实例化GPT模型
  torch.manual_seed(123)
  tokenizer = tiktoken.get_encoding("gpt2")
  
  batch = []
  
  txt1 = "Every effort moves you"
  txt2 = "Every day holds a"
  
  batch.append(torch.tensor(tokenizer.encode(txt1)))
  batch.append(torch.tensor(tokenizer.encode(txt2)))
  batch = torch.stack(batch, dim=0)
  print(batch)
  
  
  model = GPTModel(GPT_CONFIG_124M)
  
  out = model(batch)
  print("Input batch:\n", batch)
  print("\nOutput shape:", out.shape)
  print(out)
  
  
  """输出"""
  Input batch:
   tensor([[6109, 3626, 6100,  345],
          [6109, 1110, 6622,  257]])
  
  Output shape: torch.Size([2, 4, 50257])
  tensor([[[ 0.1381,  0.0077, -0.1963,  ..., -0.0222, -0.1060,  0.1717],
           [ 0.3865, -0.8408, -0.6564,  ..., -0.5163,  0.2369, -0.3357],
           [ 0.6989, -0.1829, -0.1631,  ...,  0.1472, -0.6504, -0.0056],
           [-0.4290,  0.1669, -0.1258,  ...,  1.1579,  0.5303, -0.5549]],
  
          [[ 0.1094, -0.2894, -0.1467,  ..., -0.0557,  0.2911, -0.2824],
           [ 0.0882, -0.3552, -0.3527,  ...,  1.2930,  0.0053,  0.1898],
           [ 0.6091,  0.4702, -0.4094,  ...,  0.7688,  0.3787, -0.1974],
           [-0.0612, -0.0737,  0.4751,  ...,  1.2463, -0.3834,  0.0609]]],
         grad_fn=<UnsafeViewBackward0>)
  

  如我们所见，输出张量的形状为 [2, 4, 50257]，因为我们输入了 2 个文本，每个文本包含 4 个 token。最后一个维度 50,257 对应于 tokenizer 的词汇表大小。在下一节中，我们将了解如何将这些 50,257 维的输出向量转换回 token。

• 不过，关于其大小需要简要说明：我们之前将其称为 1.24 亿参数模型；我们可以通过以下方式再次确认这一数字：

  使用 numel() 方法（“元素数量”的缩写），我们可以收集模型参数张量中的参数总数：

  total_params = sum(p.numel() for p in model.parameters())
  print(f"Total number of parameters: {total_params:,}")
  
  """输出"""
  Total number of parameters: 163,009,536
  

  模型参数数量为 163M 而非 124M，原因是未应用权重绑定（weight tying），即 GPT-2 中将token embedding层重用作输出层以减少参数；嵌入层将 50,257 维 one-hot 编码标记投影到 768 维嵌入表示，而输出层将其投影回 50,257 维以转换回单词，两者参数数量一致，需进一步验证模型参数数量为 124M。

  print("Token embedding layer shape:", model.tok_emb.weight.shape)
  print("Output layer shape:", model.out_head.weight.shape)
  
  """输出"""
  Token embedding layer shape: torch.Size([50257, 768])
  Output layer shape: torch.Size([50257, 768])
  

  相应地，如果我们减去输出层的参数数量，就会得到一个 124M 参数的模型：

  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:,}")
  
  """输出"""
  Number of trainable parameters considering weight tying: 124,412,160
  

  即$ 163,009,536 - 50257*768 = 124412160$ ,该模型现在只有 1.24 亿个参数，与 GPT-2 模型的原始大小相匹配。

  在实践中，不使用权重共享训练模型更为简便，因此本节未实现权重共享。后续章节将重新考虑权重共享，并在加载预训练权重时应用。此外，计算模型的内存需求也是一个重要的参考点。

• 计算模型内存需求

  # 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")
  
  """输出"""
  Total size of the model: 621.83 MB
  

  通过计算 GPTModel 对象中 1.63 亿参数的内存需求，并假设每个参数为 32 位浮点数，占用 4 字节，我们发现模型的总大小为 621.83 MB，这说明了即使是相对较小的 LLMs 也需要较大的存储空间。

• 在本节中，我们实现了 GPTModel 架构，并看到它输出了形状为 [batch_size, num_tokens, vocab_size] 的数值张量。
  在下一节中，我们将编写代码将这些输出张量转换为文本。

---

4.7 Generating text

• 本节，我们将实现将 GPT 模型的张量输出转换回文本的代码，回顾下，像我们上面实现的 GPT 模型这样的 LLMs 被用于一次生成一个单词。

  ​

  GPT 模型从输出张量到生成文本的过程涉及几个步骤，如下图 所示。这些步骤包括解码输出张量、根据概率分布选择标记以及将这些标记转换为人类可读的文本。

  ​

  GPT 模型通过逐步生成下一个标记的过程，从输入上下文中构建连贯的文本，具体步骤包括输出潜在标记、计算概率分布、确定标记 ID、解码为文本并更新输入序列。

  在实践中，我们会多次迭代重复此过程，直到达到用户指定的生成令牌数量

• 在代码中，我们可以按如下方式实现令牌生成过程

  def generate_text_simple(model, idx, max_new_tokens, context_size):
      # idx is (batch, n_tokens) array of indices in the current context
      for _ in range(max_new_tokens):
          
          # Crop current context if it exceeds the supported context size
          # E.g., if LLM supports only 5 tokens, and the context size is 10
          # then only the last 5 tokens are used as context
          idx_cond = idx[:, -context_size:]
          
          # Get the predictions
          with torch.no_grad():
              logits = model(idx_cond)
          
          # Focus only on the last time step
          # (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
          logits = logits[:, -1, :]  
  
          # Apply softmax to get probabilities
          probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)
  
          # Get the idx of the vocab entry with the highest probability value
          idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)
  
          # Append sampled index to the running sequence
          idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)
  
      return idx
  

  代码片段演示了使用 PyTorch 的语言模型生成循环的简单实现。它实现了一个迭代过程，迭代生成指定数量的新标记，裁剪当前上下文以适合模型的最大上下文大小，计算预测，然后根据最高概率预测选择下一个标记。

  ​

  在 generate_text_simple 函数中，我们使用 softmax 将 logits 转换为概率分布，并通过 torch.argmax 找到最大值位置。由于 softmax 是单调的，最大值位置在 logits 和 softmax 输出中一致，因此 softmax 步骤可以省略，直接对 logits 使用 torch.argmax 即可。但我们保留了这一步骤，以完整展示从 logits 到概率的转换过程，帮助理解模型如何通过贪婪解码选择最可能的下一个 token。
  在下一章实现 GPT 训练时，我们将引入更多采样技术，通过调整 softmax 输出，使模型不总是选择最可能的 token，从而增加生成文本的多样性和创造性。

• 尝试使用 “Hello, I am”，作为输入

  start_context = "Hello, I am"
  
  encoded = tokenizer.encode(start_context)
  print("encoded:", encoded)
  
  encoded_tensor = torch.tensor(encoded).unsqueeze(0)
  print("encoded_tensor.shape:", encoded_tensor.shape)
  
  """输出"""
  encoded: [15496, 11, 314, 716]
  encoded_tensor.shape: torch.Size([1, 4])
  

  eval() 模式，该模式禁用仅在训练期间使用的随机组件（例如 dropout），并在编码的输入张量上使用generate_text_simple 函数：

  # 使用generate_text_simple函数
  model.eval() # disable dropout
  
  out = generate_text_simple(
      model=model,
      idx=encoded_tensor, 
      max_new_tokens=6, 
      context_size=GPT_CONFIG_124M["context_length"]
  )
  
  print("Output:", out)
  print("Output length:", len(out[0]))
  
  """输出"""
  Output: tensor([[15496,    11,   314,   716, 27018, 24086, 47843, 30961, 42348,  7267]])
  Output length: 10
  

  使用tokenizer的decode方法，我们可以将ID转换为文本。

  decoded_text = tokenizer.decode(out.squeeze(0).tolist())
  print(decoded_text)
  
  """输出"""
  Hello, I am Featureiman Byeswickattribute argue
  

  正如我们所看到的，根据前面的输出，模型生成了乱码，这与本节开始的图中所示的连贯文本完全不同。发生了什么？该模型无法生成连贯文本的原因是我们还没有对其进行训练。到目前为止，我们刚刚实现了 GPT 架构并使用初始随机权重初始化了 GPT 模型实例。

---

4.8 Summary

• Layer normalization stabilizes training by ensuring that each layer’s outputs have a consistent mean and variance.
• Shortcut connections are connections that skip one or more layers by feeding the output of one layer directly to a deeper layer, which helps mitigate the vanishing gradient problem when training deep neural networks, such as LLMs.
• Transformer blocks are a core structural component of GPT models,combining masked multi-head attention modules with fully connected feed-forward networks that use the GELU activation function.
• GPT models are LLMs with many repeated transformer blocks that have millions to billions of parameters.
• GPT models come in various sizes, for example, 124, 345, 762, and 1542 million parameters, which we can implement with the same GPTModel Python class.
• The text generation capability of a GPT-like LLM involves decoding output tensors into human-readable text by sequentially predicting one token at a time based on a given input context.
• Without training, a GPT model generates incoherent text, which underscores the importance of model training for coherent text generation, which is the topic of subsequent chapters.

• 总结下本章节代码，包含本章节中实现的GPT模型（gpt.py）

  # This file collects all the relevant code that we covered thus far
  # throughout Chapters 2-4.
  # This file can be run as a standalone script.
  
  import tiktoken
  import torch
  import torch.nn as nn
  from torch.utils.data import Dataset, DataLoader
  
  #####################################
  # Chapter 2
  #####################################
  
  
  class GPTDatasetV1(Dataset):
      def __init__(self, txt, tokenizer, max_length, stride):
          self.input_ids = []
          self.target_ids = []
  
          # Tokenize the entire text
          token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})
  
          # Use a sliding window to chunk the book into overlapping sequences of max_length
          for i in range(0, len(token_ids) - max_length, stride):
              input_chunk = token_ids[i:i + max_length]
              target_chunk = token_ids[i + 1: i + max_length + 1]
              self.input_ids.append(torch.tensor(input_chunk))
              self.target_ids.append(torch.tensor(target_chunk))
  
      def __len__(self):
          return len(self.input_ids)
  
      def __getitem__(self, idx):
          return self.input_ids[idx], self.target_ids[idx]
  
  
  def create_dataloader_v1(txt, batch_size=4, max_length=256,
                           stride=128, shuffle=True, drop_last=True, num_workers=0):
      # Initialize the tokenizer
      tokenizer = tiktoken.get_encoding("gpt2")
  
      # Create dataset
      dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)
  
      # Create dataloader
      dataloader = DataLoader(
          dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)
  
      return dataloader
  
  
  #####################################
  # Chapter 3
  #####################################
  class MultiHeadAttention(nn.Module):
      def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
          super().__init__()
          assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
  
          self.d_out = d_out
          self.num_heads = num_heads
          self.head_dim = d_out // num_heads  # Reduce the projection dim to match desired output dim
  
          self.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.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
          self.dropout = nn.Dropout(dropout)
          self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1))
  
      def forward(self, x):
          b, num_tokens, d_in = x.shape
  
          keys = self.W_key(x)  # Shape: (b, num_tokens, d_out)
          queries = self.W_query(x)
          values = self.W_value(x)
  
          # We implicitly split the matrix by adding a `num_heads` dimension
          # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
          keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
          values = values.view(b, num_tokens, self.num_heads, self.head_dim)
          queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
  
          # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
          keys = keys.transpose(1, 2)
          queries = queries.transpose(1, 2)
          values = values.transpose(1, 2)
  
          # Compute scaled dot-product attention (aka self-attention) with a causal mask
          attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head
  
          # Original mask truncated to the number of tokens and converted to boolean
          mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
  
          # Use the mask to fill attention scores
          attn_scores.masked_fill_(mask_bool, -torch.inf)
  
          attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
          attn_weights = self.dropout(attn_weights)
  
          # Shape: (b, num_tokens, num_heads, head_dim)
          context_vec = (attn_weights @ values).transpose(1, 2)
  
          # Combine heads, where self.d_out = self.num_heads * self.head_dim
          context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
          context_vec = self.out_proj(context_vec)  # optional projection
  
          return context_vec
  
  
  #####################################
  # Chapter 4
  #####################################
  class LayerNorm(nn.Module):
      def __init__(self, emb_dim):
          super().__init__()
          self.eps = 1e-5
          self.scale = nn.Parameter(torch.ones(emb_dim))
          self.shift = nn.Parameter(torch.zeros(emb_dim))
  
      def forward(self, x):
          mean = x.mean(dim=-1, keepdim=True)
          var = x.var(dim=-1, keepdim=True, unbiased=False)
          norm_x = (x - mean) / torch.sqrt(var + self.eps)
          return self.scale * norm_x + self.shift
  
  
  class GELU(nn.Module):
      def __init__(self):
          super().__init__()
  
      def forward(self, x):
          return 0.5 * x * (1 + torch.tanh(
              torch.sqrt(torch.tensor(2.0 / torch.pi)) *
              (x + 0.044715 * torch.pow(x, 3))
          ))
  
  
  class FeedForward(nn.Module):
      def __init__(self, cfg):
          super().__init__()
          self.layers = nn.Sequential(
              nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
              GELU(),
              nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
          )
  
      def forward(self, x):
          return self.layers(x)
  
  
  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"],
              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"])
  
      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"])
  
          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
  
  
  def generate_text_simple(model, idx, max_new_tokens, context_size):
      # idx is (B, T) array of indices in the current context
      for _ in range(max_new_tokens):
  
          # Crop current context if it exceeds the supported context size
          # E.g., if LLM supports only 5 tokens, and the context size is 10
          # then only the last 5 tokens are used as context
          idx_cond = idx[:, -context_size:]
  
          # Get the predictions
          with torch.no_grad():
              logits = model(idx_cond)
  
          # Focus only on the last time step
          # (batch, n_token, vocab_size) becomes (batch, vocab_size)
          logits = logits[:, -1, :]
  
          # Get the idx of the vocab entry with the highest logits value
          idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch, 1)
  
          # Append sampled index to the running sequence
          idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)
  
      return idx
  
  
  def main():
      GPT_CONFIG_124M = {
          "vocab_size": 50257,     # Vocabulary size
          "context_length": 1024,  # Context length
          "emb_dim": 768,          # Embedding dimension
          "n_heads": 12,           # Number of attention heads
          "n_layers": 12,          # Number of layers
          "drop_rate": 0.1,        # Dropout rate
          "qkv_bias": False        # Query-Key-Value bias
      }
  
      torch.manual_seed(123)
      model = GPTModel(GPT_CONFIG_124M)
      model.eval()  # disable dropout
  
      start_context = "Hello, I am"
  
      tokenizer = tiktoken.get_encoding("gpt2")
      encoded = tokenizer.encode(start_context)
      encoded_tensor = torch.tensor(encoded).unsqueeze(0)
  
      print(f"\n{50*'='}\n{22*' '}IN\n{50*'='}")
      print("\nInput text:", start_context)
      print("Encoded input text:", encoded)
      print("encoded_tensor.shape:", encoded_tensor.shape)
  
      out = generate_text_simple(
          model=model,
          idx=encoded_tensor,
          max_new_tokens=10,
          context_size=GPT_CONFIG_124M["context_length"]
      )
      decoded_text = tokenizer.decode(out.squeeze(0).tolist())
  
      print(f"\n\n{50*'='}\n{22*' '}OUT\n{50*'='}")
      print("\nOutput:", out)
      print("Output length:", len(out[0]))
      print("Output text:", decoded_text)
  
  
  if __name__ == "__main__":
      main()
  

---