入力
  ↓
Masked Multi-Head Self-Attention
  ↓ + 残差接続 + LayerNorm
Feed-Forward Network
  ↓ + 残差接続 + LayerNorm
出力

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


class GPTBlock(nn.Module):
    """GPTのTransformerブロック"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Multi-Head Self-Attention
        self.W_qkv = nn.Linear(d_model, 3 * d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.attn_dropout = nn.Dropout(dropout)

        # Feed-Forward Network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )

        # Layer Normalization (Pre-LN)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()

        # Pre-LN
        h = self.ln1(x)

        # QKV projection
        qkv = self.W_qkv(h)
        Q, K, V = qkv.chunk(3, dim=-1)

        # ヘッド分割
        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores + mask
        attn = F.softmax(scores, dim=-1)
        attn = self.attn_dropout(attn)
        context = torch.matmul(attn, V)

        # ヘッド結合
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        attn_out = self.W_o(context)

        # 残差接続
        x = x + self.dropout(attn_out)

        # FFN (Pre-LN)
        x = x + self.ffn(self.ln2(x))

        return x


class GPT(nn.Module):
    """GPTモデル"""
    def __init__(self, vocab_size, d_model=768, n_heads=12, d_ff=3072,
                 n_layers=12, max_len=1024, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # 埋め込み
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        self.dropout = nn.Dropout(dropout)

        # Transformerブロック
        self.blocks = nn.ModuleList([
            GPTBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
        ])

        # 出力層
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

        # 因果マスクをバッファとして登録
        mask = torch.triu(torch.ones(max_len, max_len), diagonal=1)
        mask = mask.masked_fill(mask == 1, float('-inf'))
        self.register_buffer('causal_mask', mask)

    def forward(self, x):
        batch_size, seq_len = x.size()

        # 位置インデックス
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)

        # 埋め込み
        x = self.token_emb(x) + self.pos_emb(positions)
        x = self.dropout(x)

        # 因果マスク
        mask = self.causal_mask[:seq_len, :seq_len]

        # Transformerブロック
        for block in self.blocks:
            x = block(x, mask)

        # 出力
        x = self.ln_f(x)
        logits = self.head(x)

        return logits


# 動作確認
vocab_size = 50000
model = GPT(vocab_size=vocab_size, d_model=768, n_heads=12, d_ff=3072, n_layers=12)
x = torch.randint(0, vocab_size, (2, 128))
logits = model(x)
print(f"入力形状: {x.shape}")
print(f"出力形状: {logits.shape}")
print(f"パラメータ数: {sum(p.numel() for p in model.parameters()):,}")