🔰PyTorchでニューラルネットワーク基礎 #43 【Llama化編・KVキャッシュ】

概要

個人的な備忘録を兼ねたPyTorchの基本的な解説とまとめです。GPTタイプ （自己回帰型の言語モデル）の事前学習の演習で学んだことをまとめてみました。実際に小型のGPTモデルで事前学習してみました。GPT2の構造を修正して、「GPTのLlama化」に取り組んでいきます。

Llama化シリーズで扱いたい内容

• RMSNorm・SwiGLU：第42回
• KVキャッシュ (Dynamic版)：今回 HuggingFaceのKVキャッシュ解説ページ
• RoPE

演習用のファイル

• トークナイザー: livedoor_tokenizer_10k.json

  • Livedoorニュースコーパスをデータとしてbytelevel BPEで学習した10Kトークナイザー

• 学習データ：livedoorニュースコーパスのlivedoor-hommeカテゴリー

• 事前学習済みモデル：kv_homme_seq_256_bpe_10k.model

  • 系列長256のモデル。seq_len=256で動作します。

• コード: sample_43.ipynb

1. KVキャッシュによる改良

GPT2タイプのモデル以降、基本構造を保ちつつ効率性を高めるために様々な改良が進んでいます。Llama2やLlama3でも導入された方法をいくつか取り入れて演習をしてみました。今回はテキスト生成時の計算コストを低下させる方法のKVキャッシュ (Key Value Cache)についてまとめてみました。KVキャッシュは、Tranformerに登場するKey・Valueの一部を保存、再利用することで生成する文章を変更することなく計算を高速化できる方法となります。

文章の生成側だからコードは大きく変わらないだろうと思っていたのが間違えでした。Transformerブロックの流れを修正する必要がありそうです。

1.1 KVキャッシュの発想

「春　は　曙　やうやう　白く...」を例に文章生成をについて再考してみます。

GPTタイプのモデルで事前学習が終了した段階の状態だとします。次のトークンを予測できる状態です。「春は」までの２トークンを受けて、次のトークン「曙」を予測する概要を図１で表してみました。

トークン３の予測

1. 「春/は」を分散表現行列「X₁ X₂」に変更
2. 「X₁ X₂」からクエリ行列「Q₁ Q₂」、キー行列「K₁ K₂」、バリュー行列「V₁ V₂」を作成
3. softmax(QK)Vを使いTransformerブロックの出力「A₁ A₂」を計算
4. 最後に全結合層で出力次元を語彙数の整えて、次のトークンの予測ベクトルを求めます。
5. Token₂のベクトルの最大成分（確率が最も高い値）が次のトークンであるToken₃「曙」となります。

図１：文章生成の構造

「曙」の次のトークンはどのように求まるのでしょうか？

トークン４の予測

1. 「春/は/曙」を分散表現行列「X₁ X₂ X₃」に変更
2. 「X₁ X₂ X₃」からクエリ行列「Q₁ Q₂ Q₃」、キー行列「K₁ K₂ K₃」、バリュー行列「V₁ V₂ V₃」を作成
3. Attention(Q, K, V)を使いTransformerブロックの出力「A₁ A₂ A₃」を計算
4. 最後に全結合層で出力次元を語彙数の整えて、次のトークンの予測ベクトルを求めます。
5. Token₃のベクトルの最大成分（確率が最も高い値）が次のトークンであるToken₄「やうやう」となります。

１〜５の流れはトークン３の予測とトークン４の予測で等しくなっています。２つの予測から次のようなことがわかります。

• トークン４を予測するとき、トークン3を求めるときに計算した「Q₁ Q₂」、「K₁ K₂」、「V₁ V₂」を再び計算してる。

\begin{align*}
Q_1 & = X_1 W_Q^T,&Q_2 = X_2 W_Q^T\\
K_1 & = X_1 W_K^T,&K_2 = X_2 W_K^T\\
V_1 & = X_1 W_V^T,&V_2 = X_2 W_V^T
\end{align*}

• Token₄ （次のトークン）を予測するにはQ₃ （クエリ行列の最終行）のみで求められる。
• Q₃，K₃，V₃を求めるには、分散表現ベクトルX₃があれば求められる。

「重複計算を避けつつ、クリエ行列の最終行のみを利用して文章生成する」のがKVキャッシュの基本的な考えになります¹。

1.2 KVキャッシュの実装

Q、K、Vの計算をするのはTransformerブロックの中でもSelfAttention計算部分です。MultiHeadAttentionの計算を記述したCausalSelfAttentionのクラスをKVキャッシュ対応に改良していきます。

KVキャッシュ対応の注意機構のクラス

import math
import torch.nn.functional as F
from torch.nn.attention.bias import causal_lower_right

class CausalSelfAttention(nn.Module):
    ```
    よくあるMultiHeadAttentionの実装
        1. 分散表現行列Xを受け取り、Q、K、Vを計算
        2. scaled_dot_product_attention(q, k, v)で Attention(Q,K,V)を求める
    この内容に
        3. KとVをプール
        4. 新規追加と更新
    を加える。
    ```
    # (1) 初期化と基本構造
    def __init__(self, config: ModelConfig):
        super().__init__()
        assert config.d_model % config.nhead == 0
        self.nhead = config.nhead
        self.head_dim = config.d_model // config.nhead

        # q,k,v をまとめて表現するタイプ
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout_p = config.dropout
        self.resid_dropout = nn.Dropout(config.dropout)

    # (2)
    def forward(self, x, kv_cache=None, use_cache=False):
        bs, seq_len, d_model = x.shape
        q, k, v = self.qkv(x).split(d_model, dim=2)
        # (bs, seq_len, d_model) -> (bs, nhead, seq_len, head_dim)
        q = q.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        k = k.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        v = v.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)

        # (3) 過去の k,v を動的に連結していくので Dynamic Cacheと呼ばれることがある
        # q は連結しないので、q_len は seq_len のまま
        if kv_cache is not None:                  # kv_cacheに値がある場合 (not None)
            past_k, past_v = kv_cache             # kv_cacheの値をpast_k, past_vへ代入
            k = torch.cat([past_k, k], dim=2)     # 新しいkとvを追記
            v = torch.cat([past_v, v], dim=2)
        new_cache = (k, v) if use_cache else None # (4) use_cacheフラグがTrueのときは、追記したkとvをnew_cacheへ

        # (5) 右下詰めの因果マスク付きのSDPA
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=causal_lower_right(seq_len, k.size(2)),        # 学習時：(seq_len, seq_len)、生成時：(１, 過去の系列長+1)
            dropout_p=self.attn_dropout_p if self.training else 0.0,
        )
        # (6) projを使い形状を整えて出力へ
        y = y.transpose(1, 2).contiguous().view(bs, seq_len, d_model)
        y = self.resid_dropout(self.proj(y))
        return y, new_cache

KVをキャッシュする部分が(3)、ベタにtorch.catを使い実装してみました。new_cacheという新しい戻り値がある部分、ここに気をつけてネットワークモデルの構造を修正していきます。

説明メモ

• (1) QKVの計算用Linear層、形状を整えるLinear層、ドロップアウトの指定を行います。
• (2) ヘッド毎のQKV計算、ここまでがよくあるMultiHeadAttentionの実装。
• forward()の引数に注意。過去のKVの値を入れるkv_cacheとuse_cacheフラグを準備しました。
• (3) kv_cacheがない場合（学習時など）は、if文を飛ばして次！
  kv_cacheがある場合（生成時２トークン目以降のDecode時）は、if文実行！(2) で計算したkとvをkv_cacheへ追記²
• (3) 文章生成時、１トークン生成されるたびにtorch.catを使って kv_cache に (k, v) を都度追記する形になります。動的KVキャッシュと呼ばれることがあるみたいです。やや効率が悪い気がします³
• (4) use_cache=Trueの場合、(k, v)を新しいkv_cacheとして利用します。
• (5) PyTorchのscaled_dot_product_attentionを使いました。入力データの形状が学習時と文章生成時で異なります。具体的にはクエリーテンソルQの形状です。Decode時は最終行のみ参照すれば良いので系列長が１になります。マスク作成に、is_causal=True を利用すると、新トークンが最初のトークンしか見えないマスクとなってしまいます。正しくマスク作成するために、PyTorchのcausal_lower_rightを使いました⁴。
• (6) (バッチサイズ、系列長、d_model)に並び替えて、Linear層のprojを使い形状を整えます。
• 戻り値にnew_cacheがあります。学習時はnew_cache=Noneのまま。

作成したCausalSelfAttentionクラスを使ってTransformer層を作成すればKVキャッシュ対応のTransformer層になります。戻り値にnew_cacheが加わっているので他の部分も影響を受けます。

残るKVキャッシュによるGPTタイプからの変更箇所は、ネットワーク全体と文章生成部分の2箇所となります。

2. データ

livedoorニュースコーパスのlivedoor hommeカテゴリーに分類されているテキストデータを使います。GPT編から使いまわしているデータセットとなります。

前処理

• 冒頭から3行の削除「URL、日付、タイトルの削除」
• ライセンス的に改変したデータをアップロードできないので各自でお願いします

3. トークナイザー

冒頭３行を削除したlivedoorニュースコーパス全体を使い、ByteLevelのBPEで学習させた、トークン数10,000のトークナイザーを使います。学習したトークナイザーの特殊トークンと語彙数を確認しておきましょう

from tokenizers import Tokenizer

tokenizer_filename = "./tokenizer/livedoor_10k.json"  # 利用するトークナイザー
tokenizer = Tokenizer.from_file(tokenizer_filename) 

print("特殊トークンID:")
print(f"<pad>: {tokenizer.token_to_id('<pad>')}")               # <pad>: 0
print(f"<eod>: {tokenizer.token_to_id('<eod>')}")               # <eod>: 1
print(f"<system>: {tokenizer.token_to_id('<system>')}")         # <system>: 2
print(f"<assistant>: {tokenizer.token_to_id('<assistant>')}")   # <assistant>: 3
print(f"<user>: {tokenizer.token_to_id('<user>')}")             # <user>: 4
print(f"size: {tokenizer.get_vocab_size()}")                    # size: 10,000

特殊トークンはほぼ活用されないのですが、一応、指示チューニングを視野に入れた簡易構成となっています。

4. モデル設定と学習

GPTタイプの学習は、「次の単語予測問題」を解くことです。大枠はKVキャッシュの内容とは無関係なので、これまでの記事と同様の流れとなります。

• 準備 (4.1)
• ネットワークモデルの設定や定義　(4.2)
• Dataset・DataLoaderによるデータの読み込み (4.3)
• 誤差関数と誤差最小化の手法の選択　(4.4)
• 変数更新のループ　(4.5)
• 検証　(4.6)

4.1 準備

利用するライブラリーを読み込みます。

ライブラリーなど

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from tokenizers import Tokenizer
from pathlib import Path

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"{device=}")

利用するデータセットのディレクトリー、トークナイザー、学習したモデルの保存ファイル名を指定します。

ファイル名とトークナイザー

data_dir = "./data/homme_texts/"                           # データのディレクトリ
tokenizer_filename = "./tokenizer/livedoor_10k.json"       # 利用するトークナイザー
pretrain_filename = "model/kv_homme_seq_256_bpe_10k.model" # 保存ファイル名
tokenizer = Tokenizer.from_file(tokenizer_filename)        # トークナイザー

4.2 ネットワークモデルの設定や定義

ModelConfigクラスで、系列長、特殊トークンのID、バッチサイズなどの基本的な数値を定めます。DNNクラスでGPTタイプのネットワークモデルを構成します。

CausalSelfAttentionクラスの変更に伴い、ネットワークの実装部分がやや長くなります。

4.2.1 モデルの設定

モデル設定のクラスです。ネットワーク層で利用される数値を一括して管理するクラスです。学習率なども合わせて記載しているのでこのあたり、規模が大きくなる場合分離したほうが見通しが良くなる可能性があります。

モデル設定クラス

class ModelConfig:
    def __init__(self, tokenizer):
        # モデル構造
        self.vocab_size = tokenizer.get_vocab_size()
        self.seq_len = 256   # 128トークンだとSFT時に少ない
        self.d_model = 256   # 512
        self.nhead = 8
        self.dim_feedforward = 4*self.d_model
        self.num_layers = 6
        self.dropout = 0.1

        # 特殊トークンID
        self.pad_token_id = tokenizer.token_to_id("<pad>")
        self.eod_token_id = tokenizer.token_to_id("<eod>")

        # 学習データに関する設定
        self.context_size = self.seq_len         # 学習できる長さ

        # 学習設定
        self.batch_size = 256
        self.learning_rate = 0.001  # これだとデフォルトと変わらない
        self.num_steps = 2500
        self.max_grad_norm = 1.0
        self.ignore_index = -100
        self.weight_decay = 0.01

ModelConfigクラスは、数値を並べてるだけのクラスです。系列長を表すseq_lenには注意が必要です。小さすぎると文章をうまく生成することができませんし、大きな数値だと、消費されるVRAMも巨大化します。簡易的な、指示チューニングまで考慮すると、256トークン以上は必要と思われます。事前学習して文章の生成結果を確認するだけなら、系列長のseq_len=128として、10分程度 (Colab T4) で学習も終了します。

4.2.2 Transformerブロック

GPT型Transformerブロックに以下を追加しています

1. RMSNormによる正規化
2. SwiGLUを利用したFFNブロック
3. KVキャッシュ対応のCausalSelfAttentionクラス

図２：Transformerブロックの概要

図２の形に従って各ネットワーク層を配置していきます。全体をまとめてTransformerBlockとします。

Transformerブロック

import math
import torch.nn.functional as F
from torch.nn.attention.bias import causal_lower_right   # PyTorch 2.5 以降?

# (1) KVキャッシュ対応版
class CausalSelfAttention(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        assert config.d_model % config.nhead == 0
        self.nhead = config.nhead
        self.head_dim = config.d_model // config.nhead

        # q,k,v をまとめて線形層１つで対応
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout_p = config.dropout
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x, kv_cache=None, use_cache=False):
        bs, seq_len, d_model = x.shape
        q, k, v = self.qkv(x).split(d_model, dim=2)
        # (bs, seq_len, d_model) -> (bs, nhead, seq_len, head_dim)
        q = q.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        k = k.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        v = v.view(bs, seq_len, self.nhead, self.head_dim).transpose(1, 2)

        # 過去の k,v を連結（RoPEを入れる時はここより「前」で q,k を回転させる）
        # q は連結しないので、q_len は seq_len のまま
        if kv_cache is not None:                  # kv_cacheに値がある場合 (not None)
            past_k, past_v = kv_cache             # kv_cacheの値をpast_k, past_vへ代入
            k = torch.cat([past_k, k], dim=2)     # 新しいkとvを追記
            v = torch.cat([past_v, v], dim=2)
        new_cache = (k, v) if use_cache else None # use_cacheフラグがTrueのときは、追記したkとvをnew_cacheへ

        # 右下詰めの因果マスク。学習・プリフィル・デコード
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=causal_lower_right(seq_len, k.size(2)),        # 学習時：(seq_len, seq_len)、生成時：(１, 過去の系列長)
            dropout_p=self.attn_dropout_p if self.training else 0.0,
        )

        # projを使い形状を整えて出力へ
        y = y.transpose(1, 2).contiguous().view(bs, seq_len, d_model)
        y = self.resid_dropout(self.proj(y))
        return y, new_cache

# (2) SwiGLUを使う
class FFN(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden = config.dim_feedforward
        self.w_gate = nn.Linear(config.d_model, hidden, bias=False)
        self.w_up   = nn.Linear(config.d_model, hidden, bias=False)
        self.w_down = nn.Linear(hidden, config.d_model, bias=False)
        self.activation = nn.SiLU()
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        gate = self.activation(self.w_gate(x))    # ゲート側
        up   = self.w_up(x)                       # そのまま通過の特徴量
        return self.dropout(self.w_down(gate*up)) # 線形層でMLPと同様に形状を整える出力

# (3) RMSNormを使う
class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn_norm = nn.RMSNorm(config.d_model, eps=1e-6)
        self.attn = CausalSelfAttention(config)
        self.ffn_norm = nn.RMSNorm(config.d_model, eps=1e-6)
        self.ffn = FFN(config)

    def forward(self, x, kv_cache=None, use_cache=False):
        attn_out, new_cache = self.attn(self.attn_norm(x), kv_cache, use_cache)       # Attention(RMSNorm(x))
        x = x + attn_out                    # x + Attention(RMSNorm(x))
        x = x + self.ffn(self.ffn_norm(x))  # x + MLP(RMSNorm(x))
        return x, new_cache

CausalSelfAttentionクラス、FFNクラスを作成してTransformerブロックと順番に並べていきます。

説明メモ

• (1) 因果マスク付きの自己注意
  よくある$W_q$、$W_k$、$W_v$用のLinearをまとめて作成する形になっています。
  • self.qkvの部分についてです。Q，K，V用のLinearをひとまとめにして扱っています。3 * config.d_modelと3倍されている部分に注目してください。
  • self.qkv(x).splitの部分でQ，K，Vをヘッド数に応じた形で分割します。
  • 注意機構の部分。ここはPyTorchの関数を使っています😉
  • is_causal=Trueオプションではなくて、causal_lower_rightによって「未来のトークンを利用しない」マスクを作成して、自己注意として実装しています。
  • 文章生成時、is_causal=Trueだと左上揃えのマスクとなり、エラー出力もなくおかしな結果となります
• (2) SwiGLUを利用したFFNブロック
  • Transformer層の出力に向けたネットワーク層となります。層を繰り返し使えるようにw_downと命名した全結合層の出力次元がd_modelとなります。
• (3) Transformerブロック
  • 正規化層をRMSNormに変更します。
  • 出力部分のLinear層をSwiGLU対応のFFNブロックへ変更します。

4.2.3 GPTタイプのネットワークモデル

Transformerブロックが作成できたので、「埋め込み層、causalマスク付きのtransformerブロック、全結合層で次の単語予測」という形のネットワーク構造を作れば自己回帰型のモデルが完成となります。KVキャッシュを導入したので適宜if文で分岐させます。

RMSNorm＋SwiGLU版＋KVキャッシュ型

class DNN(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config

        # 埋め込み層
        self.token_embedding = nn.Embedding(num_embeddings=config.vocab_size, embedding_dim=config.d_model)
        self.pos_embedding = nn.Embedding(num_embeddings=config.seq_len, embedding_dim=config.d_model)
        self.dropout = nn.Dropout(config.dropout)

        # (1) transformer ブロック nn.ModuleListを使う
        # self.transformer_blocks = nn.Sequential(*[TransformerBlock(config) for _ in range(config.num_layers)])
        self.transformer_blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])

        # (2) 最後の出力に向けた正規化とFC層　最終的に単語数になる
        self.norm = nn.RMSNorm(config.d_model, eps=1e-6)
        self.fc = nn.Linear(config.d_model, config.vocab_size,bias=False)
        self.apply(self._init_weights)                    # 個別に初期化
        self.fc.weight = self.token_embedding.weight      # 重み共有


        # 残差方向の出力層だけ 1/sqrt(2N) スケール
        std = 0.02 / math.sqrt(2 * config.num_layers)
        for pn, p in self.named_parameters():
            if pn.endswith("proj.weight") or pn.endswith("w_down.weight"):
                nn.init.normal_(p, mean=0.0, std=std)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:          # 全層bias=False なのでスキップ
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.padding_idx is not None:   #  padding_idxは未指定なのでスキップ
                with torch.no_grad():
                    module.weight[module.padding_idx].zero_()

    # (3) cacheの有無で分類
    def forward(self, x, kv_caches=None, use_cache=False):
        bs, seq_len = x.shape

        if kv_caches is None:    # kv_cashesなし：学習時・初回 (Prefill)
            past_len = 0
        else:                    # kv_cashesある時：生成２回目以降
            past_len = kv_caches[0][0].size(2)
        assert past_len + seq_len <= self.config.seq_len, "位置埋め込みの範囲外だぞ〜 seq_lenを超えています"

        positions = torch.arange(past_len, past_len + seq_len, device=x.device)
        tok_emb = self.token_embedding(x)
        pos_emb = self.pos_embedding(positions).unsqueeze(0)
        x = tok_emb + pos_emb           # トークン埋め込み＋位置埋め込み
        x = self.dropout(x)

        new_caches = [] if use_cache else None
        for i, block in enumerate(self.transformer_blocks):
            cache = None if kv_caches is None else kv_caches[i]
            x, new_cache = block(x, kv_cache=cache, use_cache=use_cache)
            if use_cache:
                new_caches.append(new_cache)
        # (4) Prefill時に発動 (use_cacheフラグ=True・系列長が2以上)
        if use_cache and seq_len > 1:
            x = x[:, -1:]

        logits = self.fc(self.norm(x))
        return logits, new_caches

用語

• Prefill：文章生成時の最初（1トークン目）の段階。プロンプト全体を一度に処理して最初のKVキャッシュを作成するステップ
• Decode：文章生成時の2トークン目以降、1トークンずつ生成して、KVキャッシュを追記更新するステップ

説明メモ

• (1) forward部分でキャッシュを利用するかどうかの分岐が入り乱れます。やや記述が煩雑になりますが、nn.ModuleListを使い対応しています。
• (2) 正規化にRMSNorm、GPTタイプの初期化、重み共有を使います。
• (3) transformerブロック部分、キャッシュの有無で分岐します。
  学習時はkv_caches=None、生成時の２トークン目（Decode）からcache=kv_caches[i]として機能、新しく計算したKVを追記する形となります。

# 学習時: new_caches=None, cache=None, append しない
for block in self.transformer_blocks:
    x, _ = block(x, kv_cache=None, use_cache=False)

# デコード時: i層目のキャッシュを読み、i層目の新キャッシュへ追記
for i, block in enumerate(self.transformer_blocks):
    x, new_cache = block(x, kv_cache=kv_caches[i], use_cache=True)
    new_caches.append(new_cache)

• (4) Prefill時にif文が実行されることになります。x[:,-1:]によって、形状 (batch_size=1, seq_len, d_model) から最後のトークンだけを抽出し、形状を (batch_size=1, 1, d_model) に保持します。つまり、クエリ行列（Q行列）の最終行だけ抽出する部分を表現しています。

モデルの確認
tokenizerを指定して、ModelConfigを使って設定を読み込みます。

モデル

config = ModelConfig(tokenizer)
model = DNN(config).to(device)

torchinfoで大まかな構造とパラメータ数を確認しました。

torchinfoによる表示

from torchinfo import summary

summary(
    model,
    input_size=(1, config.seq_len),        # (batch, seq_len)
    dtypes=[torch.long],                   # 埋め込みの入力は整数インデックス
    depth=2,
)

summary部分にひと工夫して入出力形状も表示しました。

ネットワーク構造

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
DNN                                           [1, 256, 10000]           --
├─Embedding: 1-1                              [1, 256, 256]             2,560,000
├─Embedding: 1-2                              [256, 256]                65,536
├─Dropout: 1-3                                [1, 256, 256]             --
├─ModuleList: 1-4                             --                        --
│    └─TransformerBlock: 2-1                  [1, 256, 256]             1,049,088
│    └─TransformerBlock: 2-2                  [1, 256, 256]             1,049,088
│    └─TransformerBlock: 2-3                  [1, 256, 256]             1,049,088
│    └─TransformerBlock: 2-4                  [1, 256, 256]             1,049,088
│    └─TransformerBlock: 2-5                  [1, 256, 256]             1,049,088
│    └─TransformerBlock: 2-6                  [1, 256, 256]             1,049,088
├─RMSNorm: 1-5                                [1, 256, 256]             256
├─Linear: 1-6                                 [1, 256, 10000]           2,560,000
===============================================================================================
Total params: 11,480,320
Trainable params: 11,480,320
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 28.19
===============================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 69.24
Params size (MB): 45.92
Estimated Total Size (MB): 115.16
===============================================================================================

SwiGLUの部分、パラメータ調整（2/3倍してパラメータ数を近づける作業）を行っていないので、Linear層１個分だけパラメータが増加します。

torchinfoによると🤔1000万パラメータとなります🤔 モデルサイズが10M、心配なくらい小さいですしかも、トークン埋め込みの部分に250万パラメータとかなりの部分を使っています。そういえば、トークン埋め込みと最終層の全結合層は重みを共有しているので、実質的な学習パラメータって、740万パラメータくらいかな？とても小さなモデルです。正確なパラメータ数は、model.parameters()を足すほうが良いと思います。大まかにみると、30%程度がトークン埋め込みパラメータとなり、文脈理解の部分は残りの70%に委ねられています。文脈理解度を上げるにはModelConfigのd_modelサイズを大きくする必要がありそうです。

4.3 Dataset・DataLoaderによるデータの読み込み

GPT編で何度も登場しているRandomGPTDatasetクラスとなります。日本語文章のままファイルを読み込み、文字の開始位置をランダム化したIDデータを生成します。

Datasetクラス

class RandomGPTDataset(Dataset):
    def __init__(
        self,
        data_dir,
        tokenizer,
        context_size=128,
        samples_per_epoch=10_000,
    ):
        self.context_size = context_size
        self.samples_per_epoch = samples_per_epoch
        eod_id = tokenizer.token_to_id("<eod>")
        filenames = sorted(Path(data_dir).glob("*.txt"))

        all_ids = []
        for filename in filenames:
            text = filename.read_text(encoding="utf-8")
            document_ids = tokenizer.encode(text, add_special_tokens=False).ids
            all_ids.extend(document_ids)
            all_ids.append(eod_id)

        self.ids = torch.tensor(all_ids, dtype=torch.long)

    def __len__(self):
        # 1 epochあたりに何個のランダム窓を学習するか
        # 10,000個（回？）がデフォルト値
        return self.samples_per_epoch

    def __getitem__(self, idx):
        # idxは使わず、毎回ランダムな開始位置を選ぶ
        start = torch.randint(
            low=0,
            high=len(self.ids) - self.context_size,
            size=(1,),
        ).item()

        x = self.ids[start:start + self.context_size]
        y = self.ids[start + 1:start + self.context_size + 1]

        return {"ids": x, "labels": y}

dataset = RandomGPTDataset(
    data_dir=data_dir,
    tokenizer=tokenizer,
    context_size=config.context_size,
    samples_per_epoch=10_000,
)

dataloader = DataLoader(
    dataset=dataset,
    batch_size=config.batch_size,  # メモリ足りない場合は小さくする
    shuffle=False,
    num_workers=0,
    pin_memory=torch.cuda.is_available(), # GPU使う時True
)

RandomGPTDatasetクラス部分で入力されるデータは毎回シャッフルされます。DataLoaderでのシャッフルはFalseでOKです。

4.4 誤差関数と誤差最小化の手法の選択

次の単語予測なので誤差関数はクロスエントロピー損失、AdamWを使って誤差の最小化を狙います。オプティマイザーについては第41.5回を参照してください。

重み減衰の設定と損失関数

# weight decayの設定
decay, no_decay = [], []
for name, param in model.named_parameters():
    if not param.requires_grad:
        continue
    if param.dim() >= 2:
        decay.append(param)      # Linearの重み、Embedding
    else:
        no_decay.append(param)   # RMSNormのgain。これは 0へ
optim_groups = [
    {"params": decay,    "weight_decay": config.weight_decay},
    {"params": no_decay, "weight_decay": 0.0},
]
# AdamWの設定
optimizer = torch.optim.AdamW(optim_groups, lr=config.learning_rate, betas=(0.9, 0.95), fused=(device.type == "cuda"))
criterion = nn.CrossEntropyLoss(ignore_index=config.ignore_index)    # 今回ignore_indexを利用していないがSFT用につけておく

betasの２つ目、β₂=0.999がPyTorchのデフォルト値です。今回はLLMの慣例に従い0.95にしました。

4.5 変数更新のループ

学習ループです。GPT編以降同一コードです　簡単化のために検証など入っていません。

学習ループ

# step数で管理するためのデータローダー関数
def infinite_loader(dataloader):
    while True:
        for batch in dataloader:
            yield batch

data_iter = infinite_loader(dataloader)  # step数で計測

# ---- 学習ループの設定　tqdmを使わないほうが断然シンプル
from tqdm import tqdm               # tqdmを使い進捗状況表示をしてみた
max_iters = config.num_steps        # step数
total_token = 0                     # 累積トークン数を初期化
pbar = tqdm(range(max_iters))
use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()  # flash attentionを利用の判定

# ---- ここから学習ループ
model.train()                       # trainモードを明示
for step in pbar:
    batch = next(data_iter)
    input_ids = batch["ids"].to(device, non_blocking=True)
    labels = batch["labels"].to(device, non_blocking=True)
    optimizer.zero_grad()
    # GPUが対応している場合は、bf16へ変更してflash attentionを使う
    with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=use_bf16):
        outputs = model(input_ids)                 # [batch, seq_len, vocab_size]
        loss = criterion(
            outputs.view(-1, config.vocab_size),   # [batch * seq_len, vocab_size]
            labels.view(-1),                       # [batch * seq_len]
        )
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.max_grad_norm)  # 勾配クリップ
    optimizer.step()

    # このstepのトークン数 （バッチサイズ×系列長）を加算
    total_token += input_ids.shape[0] * input_ids.shape[1]

    # set_postfixは辞書が引数になる
    pbar.set_postfix({"loss": f"{loss.item():.4f}", "tokens": f"{total_token/1e6:.2f}M"})
    if (step+1) % 500 == 0:
        tqdm.write(f"{step+1}-step:\tloss:{loss.item():.3f}\ttokens:{total_token:,}") 

表示用の設定と形状変更で、複雑に見えますが、基本構造は初回から変わっていません

基本構造

for ループ
    input_ids, labels：入力データと教師データ
    optimizer_zero_grad()
    outputs = model(input_ids)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

input_idsとlabels部分、Flash Attentionを使う部分などやや込み入っていますが、基本構造はとても単純なのです。本来は検証目的のコードも挿入するべきです。モデルサイズ小さいし、動作確認が主目的なので大胆に割愛しています😆

今回もlossが1を下回るくらいになるまで頑張って学習させます。学習データやモデル、語彙サイズの関係からなのか詳細は不明ですが、lossは0.5程度を下回ると思います。文章を全暗記するレベル感です。暗記できていないときはどこかおかしいとなります。

メモリーが足りない！表示や時間がかかるというときは系列長やバッチサイズを環境に合わせて調整してください。

4.6 検証

文章生成、これまで同様のgenerate_textが使えると思ったのですが、ここでも、修正が〜というか、この生成のための仕組みがKVキャッシュでした。

検証と言っても実際に生成される速度を堪能してください。

Greedyでの文章生成

@torch.inference_mode()
def generate_text(model, input_ids, max_new_tokens=100, eos_token_id=None):
    model.eval()
    generated = input_ids.to(device)
    max_len = model.config.seq_len

    # (1) プリフィル：プロンプト全体
    logits, caches = model(generated, use_cache=True)
    next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
    generated = torch.cat([generated, next_token], dim=1)

    # (2) デコードのループ
    for _ in range(max_new_tokens - 1):
        if eos_token_id is not None and next_token.item() == eos_token_id:
            break
        if generated.size(1) >= max_len:
            break
        # (3) デコード：直前の1トークン(最終行)だけ使う
        logits, caches = model(next_token, kv_caches=caches, use_cache=True)
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        generated = torch.cat([generated, next_token], dim=1)

    return generated

説明メモ

• (1) プリフィル (prefill)：KVキャッシュを利用する最初の段階、入力された文章（プロンプト）から次のトークンを計算します。
• (2) デコード (Decode)：作成したKとVのキャッシュを追加、次のトークンを生成するループです。if文で終了判定をします。
• (3) 入力されるトークン（Qに対応）は、直前の１トークン (next_token) だけです。Greedyを使ったテキスト生成なので、argmaxを使います。

sample_prompt.csvを利用して表示

from tokenizers import decoders
import pandas as pd
tokenizer.decoder = decoders.ByteLevel() # BPEのときはこれが必要、文字化けするぞ

def response(prompt):
    model.eval()
    model_input = tokenizer.encode(prompt, add_special_tokens=False)
    input_ids = torch.tensor([model_input.ids], dtype=torch.long, device=device)

    output_ids = generate_text(
        model=model,
        input_ids=input_ids,
        max_new_tokens=config.seq_len,
        eos_token_id=config.eod_token_id,
    )

    # 全文をデコード
    generated_text = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=False)
    print(generated_text)

df = pd.read_csv("./data/sample_prompt.csv")

# ラフに実行時間を計測してみた
import time
start = time.perf_counter()
for num, prompt in enumerate(df["prompt"],1):
    print(f"【{num}番目】")
    response(prompt)
end = time.perf_counter()
print(f"ループ全体の実行時間: {end - start :.3f} 秒")

KVキャッシュ「あり」と「なし」でtimeでラフに生成速度を比較してみました
生成速度向上！と思いきや、系列長が短い場合はKVキャッシュなしの素朴なタイプと対して変わらない、それどころか遅いのではないか。毎回新しい情報を追記 (cat) する部分でどうも時間ロスが起きるようです。動的なKVキャッシュの問題点か？ちなみに系列長を512にすると、「KVキャッシュありのほうが高速」でした。小さなモデルの小さな環境ですのでこの差はあまり参考にならないと思います。ぜひ試して効果を体感してみてください

図３：参考出力例

次回に向けて

新しいKV情報をその都度追加するのではなく、あらかじめ箱？を準備して、必要な箇所のみ書き換えることで、cat部分の時間ロスを節約する static KVキャッシュと呼ばれる方法があるようです。少し時間がかかりそうですが確認してみたい内容です。

目次ページ

参考

• 斎藤 康毅 (2026)
  『ゼロから作るDeep Learning ❻ ―LLM編』オライリー・ジャパン
  • とてもわかり易い。しかも多色刷りで図解が丁寧すぎ。掲載されているコードもエディターのように色分けされていて、驚き

KVキャッシュ関連
提案論文はうまく探せませんでした。徐々に浸透していったのかな？適当に２つほど。

• Noam Shazeer (2019) "Fast Transformer Decoding: One Write-Head is All You Need", arXiv:1911.02150v

  • KVキャッシュと明示的に登場していませんが、keyとvalueのテンソルを再利用してみたいな記載と擬似コードが掲載されています。

• Pope et al. (2022) "Efficiently Scaling Transformer Inference" arXiv:2211.05102

  • この年代になるとKV cacheという用語が一般的になっているようでした。PrefillやDecodeという言葉遣いも登場しています。KVキャッシュの問題点に関する内容を検証しています。

注

1. 文章生成時のみの変更かと思いきや甘かった
   重複計算を避けるために一度計算した$K$や$V$をプールしていくため、Transformerブロックの中身を少し変更する必要がありました。 ↩

2. 一見重複計算してそうなのですが、テキスト生成時forwardのxに予測したトークンのみを指定することで重複計算を避けることが可能となります。 ↩

3. 最初にキャッシュで用いる領域を確保する方法もあるようで、静的KVキャッシュ (static KV cache) と分類されていました。 ↩

4. causal_lower_rightがドキュメントの説明となります。 ↩