實作GPT模型生成文字技術詳解

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, embedding_dim, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert embedding_dim % num_heads == 0, "Embedding dimension must be divisible by the number of heads"
        
        self.num_heads = num_heads
        self.embedding_dim = embedding_dim
        self.dropout = dropout
        
        self.head_dim = embedding_dim // num_heads
        
        # 定義查詢、鍵和值的線性層
        self.query_linear = nn.Linear(embedding_dim, embedding_dim)
        self.key_linear = nn.Linear(embedding_dim, embedding_dim)
        self.value_linear = nn.Linear(embedding_dim, embedding_dim)
        
        # 定義輸出的線性層
        self.output_linear = nn.Linear(embedding_dim, embedding_dim)
        
        # 定義dropout層
        self.dropout_layer = nn.Dropout(dropout)
        
    def forward(self, query, key, value, attention_mask=None):
        # 取得批次大小
        batch_size = query.size(0)
        
        # 線性變換
        query = self.query_linear(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        key = self.key_linear(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        value = self.value_linear(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        
        # 計算注意力得分
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        if attention_mask is not None:
            scores = scores + attention_mask
        
        # 計算注意力權重
        attention_weights = F.softmax(scores, dim=-1)
        
        # 應用dropout
        attention_weights = self.dropout_layer(attention_weights)
        
        # 計算輸出
        output = torch.matmul(attention_weights, value).transpose(1, 2).contiguous().view(batch_size, -1, self.embedding_dim)
        
        # 線性變換輸出
        output = self.output_linear(output)
        
        return output

#### 內容解密：
1. **多頭注意力機制初始化**：我們定義了一個名為`MultiHeadAttention`的類別，用於初始化多頭注意力模組。它接收三個引數：注意力頭的數量、嵌入維度和dropout率。
2. **斷言檢查**：檢查嵌入維度是否能被注意力頭數量整除，以確保每個頭的維度是整數。
3. **線性層定義**：定義了四個線性層，分別用於查詢、鍵、值和輸出的線性變換。
4. **前向傳播**：實作了多頭注意力的前向傳播過程，包括線性變換、注意力得分計算、注意力權重計算和輸出計算。
5. **注意力掩碼**：支援可選的注意力掩碼，用於遮蔽未來的標記或填充標記。
6. **Dropout**：在注意力權重上應用dropout，以防止過擬合。

# 初始化GPT-2大小的注意力模組
num_heads = 12
embedding_dim = 768

multi_head_attention = MultiHeadAttention(num_heads, embedding_dim)

#### 內容解密：
1. **引數設定**：設定了與最小GPT-2模型相比對的引數：12個注意力頭和768維的嵌入大小。
2. **例項化多頭注意力模組**：使用設定的引數例項化了一個`MultiHeadAttention`物件。

## 從零開始實作GPT模型以生成文字

本章節將探討如何從零開始實作一個GPT模型，以實作文字生成的功能。首先，我們將介紹GPT模型的基本架構，接著逐步實作各個元件，最終組裝成完整的GPT模型。

### GPT-2與GPT-3的比較

本章節主要關注GPT-2模型，因為OpenAI已公開其預訓練模型的權重，我們將在後續章節中載入這些權重。GPT-3在模型架構上與GPT-2基本相同，但規模更大，引數量從15億增加到1,750億，並在更多資料上進行訓練。然而，GPT-3的權重尚未公開。對於學習如何實作大語言模型（LLM）來說，GPT-2是一個更好的選擇，因為它可以在單台筆記型電腦上執行，而GPT-3則需要GPU叢集進行訓練和推斷。

### GPT模型的架構

GPT模型主要由嵌入層（embedding layers）、轉換器區塊（transformer blocks）以及輸出層（output layers）組成。其中，轉換器區塊是GPT類別LLM的關鍵元件。模型的目標是每次生成一個單詞的文字。

#### GPT模型組態

我們使用以下Python字典來指定小型GPT-2模型的組態：
```python
GPT_CONFIG_124M = {
    "vocab_size": 50257,  # 詞彙大小
    "context_length": 1024,  # 上下文長度
    "emb_dim": 768,  # 嵌入維度
    "n_heads": 12,  # 注意力頭數量
    "n_layers": 12,  # 層數
    "drop_rate": 0.1,  # dropout率
    "qkv_bias": False  # Query-Key-Value偏差
}

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"])
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])]
        )
        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__()
        
    def forward(self, x):
        return x

class DummyLayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-5):
        super().__init__()
        
    def forward(self, x):
        return x

class DummyGPTModel(nn.Module):
    # 模型的實作細節...

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)

torch.manual_seed(123)
model = DummyGPTModel(GPT_CONFIG_124M)
logits = model(batch)
print("Output shape:", logits.shape)
print(logits)

class LayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-5):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        
    def forward(self, x):
        # 計算均值和方差，並進行標準化
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_normalized = (x - mean) / torch.sqrt(var + self.eps)
        return x_normalized * self.weight + self.bias