LLMの進化の歴史 - GPT2の再構築 -(1)

本文は Andrej Karpathy の 4 時間で GPT-2 を再現する内容に基づいており、個人的に非常に良い動画だと感じました。これは LLM の進化の歴史の終わりを示すものであり、この記事はその文字バージョンの補足です。前の内容については https://blog.nagi.fun のホームページを参照してください。このブロガーは非常に詳細に書いています。

このシリーズは、実装の主要部分、加速実装、分散トレーニングの 3 つの部分に分かれています。

GPT-2 nn.Module の実装#

Config 設定#

@dataclass
class GPTConfig():
    block_size: int=1024     # シーケンス長の制限（コンテキストウィンドウの長さ）
    vocab_size: int=50527    # 語彙のサイズ
    n_layer: int=12          # Transformerの層数
    n_head: int=12           # アテンションヘッドの数
    n_embd: int=768          # 埋め込み次元（各トークンのベクトルの長さ）

@dataclassデコレーターは、GPTConfigという名前の設定クラスを定義します。

（デコレーターが理解できない場合は、自分で CSDN または Zhihu を調べてください）

なぜ dataclass を使用するのか：

•通常のクラスでは手動で__init__メソッドを書く必要がありますが、デコレーターを使うと非常に簡単です。
•明示的な宣言をサポートし、直接pirnt(GPTConfig(n_head=16))でパラメータを直接印刷できます。

BackBone#

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
		        # 単語トークン埋め込み
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            # 単語位置埋め込み
            wpe = nn.Embedding(config.block_size, config.n_embd),
            # 主体ブロック
            h = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
            # 単語トークン埋め込み、
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.attn = CasualSelfAttention(config)
        self.mlp = mlp(config)
    
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

transformer

self.transformer：transformer アーキテクチャの本体

nn.ModuleDict：nn.Module 内の辞書、nn.ModuleDict(ln_f = nn.LayerNorm(config.n_embd) ,)は{ln_f: nn.LayerNorm(config.n_embd)}と理解できます。

wte ：単語から埋め込みへの線形層 [語彙の長さ、埋め込み次元]、単語を特徴ベクトルに変換します。

wpe ：単語位置から埋め込みへの線形層 [シーケンス長、埋め込み次元]、位置情報を特徴ベクトルに変換します。

h ：transformer のエンコーダ主体、各 Block は 1 つのアテンションと 1 つの mlp で構成されています。

ln_f ：LayerNorm、Pre-Norm 後に得られた大きな分散を正規化します。後で説明があります。

lm_head ：最終層の出力で、単語の特徴ベクトルを具体的な単語に変換します。

Block ：Transformer は複数の同じ Block で構成されています。

Tips⚠️：ここでは GPT2 の LN 層がアテンションと mlp の前にあることがわかりますが、上の図の原文の LN 層と残差（先に残差、次に正規化）を接続するのとは異なります。

Karpathyの説明は次のようです：元のモデルは先に残差接続を行い、その後LN正規化処理を行います。これは、接続された残差も正規化されることを示していますが、これは良くありません。純粋な残差の方が良いです。なぜなら、逆伝播の際に勾配が戻るとき、加算はその勾配を2つの分岐に均等に分配するため、勾配が残差の経路を通じて直接入力に流れることを意味します。最適化の観点から見ると、クリーンな残差は望ましいです。

正直なところ、彼の説明は理解できなかったので、関連する内容をネットで検索しました。そして、GPT2のやり方はpre-normと呼ばれ、Attention is all you needのやり方はpost-normと呼ばれることがわかりました。

pre-norm

スー神はこの 2 つの違いについて非常に適切に説明しています。残差接続は $x+F(x)$ であり、もし $x$ の分散が $\sigma^2_1$ で、 $F(x)$ の分散が $\sigma^2_2$ であれば、残差接続後の分散は $\sigma^2_1+\sigma^2_2$ になります。つまり、残差は分散を増幅します。私たちはこの分散を縮小する方法を考える必要があります。素朴な方法は正規化を追加することです。つまり、 $x_{t+1}=Norm(x_t+F(x))$ ですが、この方法は前向き伝播の分散を安定させる一方で、実際には残差の恒等分岐を大幅に弱めてしまい、残差の「訓練しやすさ」の利点を失ってしまいます。通常、ウォームアップを行い、十分に小さな学習率を設定しないと収束しません。また、transformer の構造には 2 つの特徴があります： ウォームアップ段階のハイパーパラメータに敏感で、最適化プロセスの収束速度が遅い。（この部分は筆者もなぜか分かりませんが）、post-norm の状況では収束がさらに難しくなり、トレーニングコストもある程度上昇します。

次に、残差の恒等分岐が弱まった理由（実際には Karpathy が言っているクリーンな残差）を説明します。初期状態で $x$ と $F(x)$ の分散がどちらも 1 であれば、 $x+F(x)$ の分散は 2 になります。そして、Normalization 操作は分散を再び 1 に戻す役割を果たします。これにより、初期段階の Post Norm は次のようになります。

x_{t+1}=\frac{x_t+F(x_t)}{\sqrt{2}}

再帰的に続けると、

x_l=\frac{x_{l-1}}{\sqrt{2}}+\frac{F_{l-1}(x_{l-1})}{\sqrt{2}}=\frac{x_{l-2}}{{2}}+\frac{F_{l-2}(x_{l-2})}{{2}}+\frac{F_{l-1}(x_{l-1})}{{2}^{1/2}}

x_l=\frac{x_{0}}{{2^{l/2}}}+\frac{F_{0}(x_{0})}{{2^{l/2}}}+\frac{F_{1}(x_{1})}{{2}^{(l-1)/2}}+\frac{F_{2}(x_{2})}{{2}^{(l-2)/2}}+...+\frac{F_{l-1}(x_{l-1})}{{2}^{1/2}}

本来残差の意味は、前の入力層に「グリーンチャンネル」を作り、勾配がより直接に逆伝播できるようにすることですが、Post Normではこの「グリーンチャンネル」が大幅に弱められています。前のチャンネルに近づくほど重みが小さくなり、残差接続が何度も行われた後、前の残差は末尾の勾配変化を感知できなくなり、残差は「名存実亡」となり、訓練が難しくなります。

論文は「ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE」を参照してください。

修正された Pre-Norm の形式は次のようになります。

x_{t+1}=x_t+F_t(Norm(x_t))

展開して反復すると：

x_{t+1}=x_t+F_t(Norm(x_t)) = x_{t-1}+F_{t-1}(Norm(x_{t-1}))+F_t(Norm(x_t))

x_{t}=x_{0}+F_{0}(Norm(x_{0}))+F_1(Norm(x_1))+...+F_{l-1}(Norm(x_{l-1}))

各残差経路は平等に重み付けされ、残差の効果は Post Norm よりも明確になります。したがって、最適化も容易になります。もちろん、最終的な $x_l$ の分散は非常に大きくなるため、予測層の前に $x_l$ にも正規化を追加する必要があります。これがln_fです。

Karpathy は、アテンションはトークン間の通信の場所であり、プーリング関数であり、加重和関数であり、 ruduce operation であると述べています。
MLP は各トークンで発生し、トークン間で情報が収集または交換されることはなく、 map operation です。

MLP#

class mlp(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, config.n_embd*4)
        self.Gelu = nn.GELU(approximate='tanh')
        self.c_proj = nn.Linear(config.n_embd*4, config.n_embd)
    
    def forward(self, x):
        x = self.c_fc(x)
        x = self.Gelu(x)
        x = self.c_proj(x)
        return x

非常にシンプルな MLP 線形マッピングで、 [n_embd, 4 * n_embd] から [4 * n_embd, n_embd] へ、間に非線形層 GELU 活性化関数を追加します。GELU の関数グラフは Relu に似ていますが、尾部で微分可能であり、つまり x が 0 未満のときに微分が常に 0 であるという Relu の問題を解決します。この滑らかさがより良い効果を生み出します。

Karpathy はここでなぜ tanh の近似値を使用するのかについて言及しており、これは歴史的な問題であり、tensorflow の時代に正確な GELU は非常に遅かったため、tanh を利用して GELU を近似する関数が開発されました。

GELU

Attention#

class CasualSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, config.n_embd*3)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.n_embd = config.n_embd
        self.n_head = config.n_head
        # 後続の重みモジュールでその役割について説明します
        self.c_proj.NANOGPT_SCALE_INIT = 1 
        self.register_buffer("bias", torch.tril((torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size)))
    
    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        # qkvを取得
        q, k, v = qkv.split(self.n_embd, dim=2)
        # query, key, valueはすべて[B, n_head, T, n_embd//n_head]に分割されます
        query = q.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
        key = k.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
        value = v.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
        # QK^T/d
        att = query @ key.transpose(-1, -2) * (1.0/math.sqrt(key.size(-1)))
        mask_att = att.masked_fill(self.bias[:,:,:T,:T]==0, float('-inf'))
        wei = F.softmax(mask_att, dim=-1)
        out = wei @ value
        out = out.transpose(1,2).contiguous().view(B, T, C)
        out = self.c_proj(out)
        return out

self.c_attn： $W_q,W_k,W_v$ の組み合わせで、入力 $x$ を入力の $Q,K,V$ に変換します。

self.c_proj： $\frac{QK^T}{\sqrt{d_k}}V$ を計算した後の線形層です。

self.n_embd：各トークンの特徴ベクトル空間です。

self.n_head：マルチヘッドアテンションメカニズムのヘッドの数です。

self.bias：ここでの bias はマスクの意味であり、上三角行列です。前のトークンが後のトークンの情報を学習するのを防ぎます。具体的な原理は次の通りです：入力 $x$ に対して：

x=\begin{bmatrix} x_{11} & x_{12} & x_{13} \\ x_{21} & x_{22} & x_{23} \\ x_{31} & x_{32} & x_{33} \end{bmatrix},bias=\begin{bmatrix} 1 & 0 & 0 \\ 1 & 1 & 0 \\ 1 & 1 & 1 \end{bmatrix} ,mask_{att}=\begin{bmatrix} x_{11} & -inf & -inf \\ x_{21} & x_{22} & -inf \\ x_{31} & x_{32} & x_{33} \end{bmatrix}

-inf は後続のsoftmaxプロセスで 0 に近い値に変わり、分類に影響を与えません。

contiguous()：transpose は物理的な順序を変更せず、形式的な順序を変更するだけですが、この関数を使用することで物理的な順序を修正できます。例えば、配列 $[[[1,2][7,8]],[[3,4][5,6]]]\space \space \space \space \space shape=[1,2,2]$ 、 transpose(1,2)後は $[[[1,2][2,3]],[[3,5][4,6]]]\space \space \space \space \space shape=[1,2,2]$ ですが、これらの配列の物理的なストレージは $[1,2,7,8,3,4,5,6]$ です。したがって、transpose 後の配列に対して view 操作を行うとエラーが発生します。

HuggingFace からのダウンロード#

    # class GPT内で
    @classmethod
    def from_pretrained(cls, model_type):
        """Huggingfaceから事前学習済みGPT-2モデルの重みを読み込みます"""
        # 4種類のモデル
        assert model_type in {'gpt2','gpt2-medium','gpt2-large','gpt2-xl'}
        # 読み込んでいるモデルの種類を表示
        print("Loading weights from pretrained gpt:%s"%model_type)
        # 各GPTに対応するハイパーパラメータは異なります
        config_args ={
        'gpt2' : dict(n_layer=12,n_head=12,n_embd=768), # 124M params
        'gpt2-medium' : dict(n_layer=24,n_head=16,n_embd=1024), #350M params
        'gpt2-large' : dict(n_layer=36,n_head=20,n_embd=1280), #774M params
        'gpt2-xl' : dict(n_layer=48,n_head=25,n_embd=1600), #1558M params
        }[model_type]
        # 語彙のサイズは常に50527です
        config_args['vocab_size'] = 50257 
        # 単一のブロックのサイズは常に1024です
        config_args['block_size'] = 1024
        # モデルにハイパーパラメータをインポートします
        config = GPTConfig(**config_args)
        model = GPT(config)
        # sdはモデルパラメータの名前辞書です
        sd = model.state_dict()
        sd_keys = sd.keys()
        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # このマスクを除外します

        # HFから重みをダウンロードし、sd_hfはダウンロードされたモデルパラメータの名前辞書です
        model_hf = GPT2LMHeadModel.from_pretrained(model_type, cache_dir="/home/shong_Tan/project/gpt_2/model_weight", local_files_only=True)
        sd_hf = model_hf.state_dict()

        sd_keys_hf = sd_hf.keys()
        # HFの重みからマスクbiasを除外します
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')]
        # HFの重みからbiasを除外します
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')]
        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
        # sdとhf_sdのパラメータ名が同じ数であることを確認します
        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
        # sdとhf_sdのtransformerブロックの重み名が同じであることを確認します
        for k in sd_keys_hf:
            if any(k.endswith(w) for w in transposed):
                assert sd_hf[k].shape[::-1] == sd[k].shape
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k].t())
            else:
                assert sd_hf[k].shape == sd[k].shape, f"mismatched keys: {sd_hf[k].shape} != {sd[k].shape}"
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k])
        return model

コードのコメントを見ておけば大丈夫です。

Tips⚠️：HF からダウンロードされたlm_head.weight と transformer.wte.weight の形状は同じで、どちらも $[50527, 768]$ です。これは入力埋め込みと出力 logits であり、これらは一貫性がある必要があります。これは、トークンが特徴ベクトルに埋め込まれたときの意味が相互作用の後に出力されるときに、元のトークンに戻ることができることを反映します。同時に、 $50527*768 \approx 40M$ であり、大量のメモリを節約できます。

Forward#

    # class GPT内で
    def forward(self, idx, target):
		    # 入力時の次元は[batch, トークンの長さ]
        B, T = idx.size()
        # トークンの長さはコンテキストを超えてはいけません
        assert T <= self.config.block_size, f"入力コンテキスト長の制限を超えています {T-self.config.block_size} トークン"
        # pos [0,1,2,..,T-1]、デバイスに配置することを忘れずに
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        # 位置埋め込み
        pos = self.transformer.wpe(pos) #(T, n_embd)
        # トークン埋め込み
        tok = self.transformer.wte(idx) #(B, T, n_embd)
        # 加算は (T, n_embd) の次元で数値を加算します
        x = tok + pos
        # transformerブロックを通過
        for block in self.transformer.h:
            x = block(x)
        # 最後の層の正規化
        x = self.transformer.ln_f(x)
        # 線形層の出力
        logits = self.lm_head(x) #(B, T, vocab_size)
        loss = None
        # targetがある場合は、つまりラベルがある場合は訓練を行い、損失関数を計算します。そうでなければ推論だけで済みます。
        if target is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1))
        return logits, loss

# 小試牛刀
num_return_sequences = 5
max_length = 30
model = GPT.from_pretrained('gpt2')
# evalは評価時にドロップアウト層を無効にし、batchnormにも異なる反応を示し、パラメータを凍結します
model.eval()
# モデルをgpuに移動
model.to('cuda')

# 以下はトークン化で、openaiのtiktokenライブラリを使えば大丈夫です。原理を知りたい場合は、記事の冒頭のブロガーのブログを参照することをお勧めします。
import tiktoken
enc = tiktoken.get_encoding('gpt2')
tokens = enc.encode("Hello, I'm a language model, ")
tokens = torch.tensor(tokens, dtype=torch.long) # [8, ]
tokens = tokens.unsqueeze(0).repeat(num_return_sequences, 1) # [5,8]
x = tokens.to('cuda')

while x.size(1) < max_length:
		with torch.no_grad():
				# モデルに入力して結果を得る
				logits, loss  = model(x) # x: [B, T]    logits:[B,T,C]
				# 最後のトークンの予測を取得
				logits = logits[:, -1, :] # [B, 1, C]
				# 最後の次元Cでsoftmaxを取る
				probs = F.softmax(logits, dim=-1) # [B, 1, C]
				# 最後の次元Cからtopkの確率と対応するインデックスを選択
				topk_probs, topk_indices = torch.topk(probs, 50, dim=-1) 
				# topkの確率からランダムに1つの確率を選択
				ix = torch.multinomial(topk_probs, 1)
				# 選ばれた確率に対応するインデックスを見つける
				xcol = torch.gather(topk_indices, -1, ix)
				# 得られた出力トークンをxに加え、再び入力として使用します [B, T+1]
				x = torch.cat((x, xcol), dim=1)

トークン化の形式： "Hello, I'm a language model," を [15496, 11, 314, 1101, 257, 3303, 2746, 11, 220] に変換します。

以下のウェブサイトを参考にして自分で試すことができます。

https://tiktokenizer.vercel.app/

初期化#

Dataset#

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
# これはmacのMチップシリーズです
elif hasattr(torch.backends, "eps") and torch.backends.mps.is_available():
    device = 'mps'
print("Using device: ", device)
import tiktoken
enc = tiktoken.get_encoding('gpt2')
with open('input.txt', 'r') as f:
    data = f.read()
text = data[:1000]
tokens = enc.encode(text)
B, T = 4, 32
buf = torch.tensor(tokens[:B*T+1])
buf.to(device)
# 本質的には前n個の単語からn+1個の単語を推測します
x = buf[:-1].view(B, T)
y = buf[1:].view(B,T)

model.GPT(GPTConfig())
model.to(device)
logits, loss = model(x)
print(loss.item())

ここでlossの値はおおよそ 11 程度です。なぜなら $-log(\frac{1}{50257})\approx 11$ だからです。

単一バッチのトレーニングコード

# AdamWオプティマイザーを使用します。AdamとSGDの違いは自分で調べてください。
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
		# オプティマイザーの勾配をゼロにします
		optimizer.zero_grad()
		# logitsと損失を取得します
		logits, loss = model(x, y)
		# 逆伝播で導関数を求めます
		loss.backward()
		# 導関数を利用して元のパラメータを更新します
		optimizer.step()

Adam オプティマイザーは SGD よりも早く収束します。

Dataloader 関数

class DataLoaderLite():

    def __init__(self, B, T):
        self.B = B
        self.T = T
        # input.txt全体を読み込みます
        with open('input.txt', 'r') as f:
            data = f.read()
        enc = tiktoken.get_encoding('gpt2')
        tokens = enc.encode(data)
        self.tokens = torch.tensor(tokens, dtype=torch.long)
        print(f"load {len(self.tokens)} tokens")
        print(f"1 epoch = {len(self.tokens)//(B*T)} batched")
        # 現在のバッチの位置を定義します
        self.current_position = 0

    def next_batch(self):
        B, T = self.B, self.T
        buf = self.tokens[self.current_position: self.current_position+B*T+1]
        x = buf[:-1].view(B, T)
        y = buf[1:].view(B, T)
        # 各バッチにはB*T個の要素のペアがあります
        self.current_position += B*T
        # バッチがトークンを使い果たした場合、再びtokens[0]に戻ります
        if self.current_position+B*T+1 > len(self.tokens):
            self.current_position = 0
        return x, y

トレーニングコードの修正

train_loader = DataLoaderLite(4, 32 )
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
		optimizer.zero_grad()
    x, y = train_loader.next_batch()
    x, y = x.to(device), y.to(device)
    logits, loss = model(x, y)
    loss.backward()
		optimizer.step()
		print(f"step: {i}, loss: {loss.item()}")

重み#

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, ' '):
                std = std * (2*self.config.n_layer**-0.5)
            torch.nn.init.normal_(module.weight, mean=0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)

        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0, std=0.02)

std = std * (2*self.config.n_layer)**-0.5 ：ここでの分散は残差フローの寄与を考慮したものであり、各残差接続は入力 input に等しい寄与を持つことを示しています。処理するための因子が必要です。 $\frac{1}{\sqrt{2*n_{layer}}}$ 、ここでは Pre-Norm の残差接続によって分散が大きくなりすぎるのを制御しています。ここでの 2 は、各層内で Attention と MLP がそれぞれ 1 回ずつ残差を使用するためです。

std ：std の値の出所も根拠があります。GPT2 の文書によれば、 $\frac{1}{\sqrt{n_{embd}}}$ のあたりが最適です。