概要
個人的な備忘録を兼ねたPyTorchの基本的な解説とまとめです。GPTタイプ （自己回帰型の言語モデル）の事前学習の演習で学んだことをまとめてみました。実際に小型のGPTモデルで事前学習してみました。GPT2の構造を修正して、「GPTのLlama化」に取り組んでいきます。
HuggingFaceのKVキャッシュに関するページを眺めていると、登場してきたのがstatic KV cacheです。ということで、今回は静的KVキャッシュについて演習してみました。内容的に43回とかなり重なるのでKVキャッシュ部分のみ掲載しています。相対的な位置埋め込みと合わせるとキャッシュ部分のコードが読みにくくなるので今回もRoPEなしで実装してみました。データ読み込み、学習部分、推論など前回とほぼ重複するので静的KVキャッシュ部分を中心に記載してあります。
Llama化シリーズで扱いたい内容
- RMSNorm・SwiGLU：第42回
- 動的KVキャッシュ：第43回
- KVキャッシュ (Static版)：今回 HuggingFaceのKVキャッシュ解説ページ
- RoPE (Rotary Positional Embeddings) ：次回
演習用のファイル
- トークナイザー: livedoor_tokenizer_10k.json
- Livedoorニュースコーパスをデータとしてbytelevel BPEで学習した10Kトークナイザー
- データ：LivedoorニュースコーパスのHOMMEカテゴリー
- コード: sample_44.ipynb
1. 静的KVキャッシュ (Static KV Cache) による改良
GPT2タイプのモデル以降、基本構造を保ちつつ効率性を高めるために様々な改良が進んでいます。前回からテキスト生成時の計算コストを低下させる方法のKVキャッシュ (Key Value Cache)についてまとめいます。KVキャッシュは、Tranformerに登場するKey・Valueの一部を保存、再利用することで生成する文章を変更することなく計算を高速化できる方法となります。
1.1 動的KVキャッシュのコード
第43回で検討した動的KVキャッシュのコードのポイントを確認しておきます。KVをキャッシュする部分が下記コードの(3)番目となります。テンソルを事前に作成せず、torch.catを使い１個ずつ繋げる形で実装しています。トークンが生成されるたびに情報が追加されていく形となります。
dynamic 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 を動的に連結していく
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)、生成時：(１, 過去の系列長)
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
ベクトルkとvをその都度追記してテンソルKとVの形状が変更されるためdynamic KV cache （動的KVキャッシュ）と呼ばれています。
cat部分
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
この部分、先にKやVの最大系列長で領域を確保して、キャッシュする値に入れ替えていく方法にすれば、いろいろ効率が良さそうです。今回はstatic KVキャッシュについて少し深めてみます。
1.2 静的KVキャッシュの発想
「春 は 曙 やうやう 白く...」を例に文章生成をについて再考してみます。GPTタイプのモデルで事前学習が終了した段階の状態だとします。次のトークンを予測できる状態です。「春/は」までの２トークンを受けて、次のトークン「曙」を予測する概要を図１と図２で表してみました。
図１：動的KVキャッシュ
「春/は/曙」というトークン１〜トークン３を入力してトークン４を予測する状況で考えてみます。次の単語までの流れはだいたい次のようになります。
- 「春/は/曙」を分散表現行列「X₁ X₂ X₃」に変更
- 「X₁ X₂ X₃」からクエリ行列「Q₁ Q₂ Q₃」、キー行列「K₁ K₂ K₃」、バリュー行列「V₁ V₂ V₃」を作成
- softmax(QK)Vを使いTransformerブロックの出力「A₁ A₂ A₃」を計算
- 最後に全結合層で出力次元を語彙数の整えて、次のトークンの予測ベクトルを計算
- Token₃のベクトルの最大成分（確率が最も高い値）が次のトークンであるToken₄「やうやう」が確定
トークン４を予測するとき、トークン3を求めるときに計算した「Q₁ Q₂」、「K₁ K₂」、「V₁ V₂」を再び計算しています。この重複計算を行わずに値をプールし、K₃、V₃を追記して使うという発想が動的KVキャッシュでした。赤い枠で囲まれた部分をキャッシュしておき、K₃やV₃を追加して次のトークンを求める形になります。トークンを生成するたびにテンソルKとVの形状が生成時に1個ずつ増えていくことになります。図１の状況に対応します。
テンソルKとVの形状を徐々に大きくしていくのではなく、最初にKとVの形状を固定して、計算に利用しない部分をマスクするという発想が静的KVキャッシュの基本となります。
図２：静的KVキャッシュ
静的キャッシュの発想は図２のように
- 最初にテンソルKとVのサイズを固定
- 本来利用しない部分を [mask] して計算
- [mask] 部分に新しい情報（ K₃とV₃ ）をコピー
となります。
事前にKとVのサイズを決めるために利用するメモリーが大きくなることが予想されますが、KとVの形状が一定のためtorch.compileを利用して高速化できると言われています。
2. 静的KVキャッシュの実装
Q、K、Vの計算をするのはTransformerブロックの中でもSelfAttention計算部分になります。MultiHeadAttentionの計算を記述したCausalSelfAttentionのクラスを静的KVキャッシュ対応に改良していきます。基本的に図2をコード化すれば良いので主要な変更箇所は次の3点になります。
- 事前に大きなテンソルを作成
- [mask] の作成
- torch.catの部分を、事前作成したテンソルの該当箇所置き換えに変更
動的KVキャッシュのときはAttentionクラス部分に直接記述したのですが、今回はStaticKVCacheクラスを作成して、テンソル領域の確保、アップデート、初期化などのメソッドを準備しておきます。
RMSNorm + SwiGLU + Static KV cache
import math
import torch.nn.functional as F
class StaticKVCache:
"""
Static KV Cache
1. 事前に、全層の K/V を (batch, nhead, max_cache_len, head_dim) で確保
2. Attention には常に max_cache_len の長さ、未使用部分はマスク
3. 書き込みは cache_position の位置へ index_copy_ で 置き換え
"""
def __init__(self, config, batch_size, device, dtype, max_cache_len=None):
self.max_batch_size = batch_size
self.max_cache_len = max_cache_len or config.seq_len
head_dim = config.d_model // config.nhead
shape = (batch_size, config.nhead, self.max_cache_len, head_dim)
# (1) 0で予めKVキャッシュの領域をテンソルとして確保
# forループ：すべてのTransformer層で作成
self.key_cache = []
self.value_cache = []
for _ in range(config.num_layers):
# empty ではなく zeros
# 未書き込みの領域も Attention に入るぞ 計算量が減ることはない
k = torch.zeros(shape, device=device, dtype=dtype)
v = torch.zeros(shape, device=device, dtype=dtype)
self.key_cache.append(k)
self.value_cache.append(v)
# (2) index_copy_() で、k、v cacheのテンソルの指定位置に値を上書き
# layer_idxのキャッシュを戻り値とする
def update(self, k, v, layer_idx, cache_position):
"""k, v: (bs, nhead, T, head_dim) を cache_position の位置へ書き込み"""
k_out = self.key_cache[layer_idx]
v_out = self.value_cache[layer_idx]
k_out.index_copy_(2, cache_position, k) # dim=2 (系列方向) に書き込み
v_out.index_copy_(2, cache_position, v)
return k_out, v_out
# KとVのキャッシュを０で初期化（新しく作成ではなくて上書き）
def reset(self):
# 同じオブジェクト（同じアドレス）を使い回すため、zero_ でクリアする
for k, v in zip(self.key_cache, self.value_cache):
k.zero_()
v.zero_()
class CausalSelfAttention(nn.Module):
def __init__(self, config: ModelConfig, layer_idx: int):
super().__init__()
assert config.d_model % config.nhead == 0
self.config = config
self.layer_idx = layer_idx # Transformerのlayer index
self.nhead = config.nhead
self.head_dim = config.d_model // config.nhead
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, cache_position=None, attn_mask=None):
bs, seq_len, d_model = x.shape
# q, k, v を計算
q, k, v = self.qkv(x).split(d_model, dim=2)
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)
# (4)
if kv_cache is not None:
# 推論時（prefill / decode）
# キャッシュへ書き込み、形状 (bs, nhead, max_cache_len, head_dim) の K/V を受け取る
k, v = kv_cache.update(k, v, self.layer_idx, cache_position)
# 未来の位置・未書き込みの位置は attn_mask で隠す is_causal=Trueは使えないぞ〜
y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
else:
# 学習時: キャッシュなしの通常の causal attention
y = F.scaled_dot_product_attention(
q, k, v,
is_causal=True,
dropout_p=(self.attn_dropout_p if self.training else 0.0),
)
y = y.transpose(1, 2).contiguous().view(bs, seq_len, d_model)
y = self.resid_dropout(self.proj(y))
return y
# SwiGLUタイプのFFN
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)) # 線形層でFFNと同様に形状を整える出力
# RMSNormタイプのTransformerブロック
class TransformerBlock(nn.Module):
def __init__(self, config, layer_idx):
super().__init__()
self.attn_norm = nn.RMSNorm(config.d_model, eps=1e-6)
self.attn = CausalSelfAttention(config, layer_idx)
self.ffn_norm = nn.RMSNorm(config.d_model, eps=1e-6)
self.ffn = FFN(config)
def forward(self, x, kv_cache=None, cache_position=None, attn_mask=None):
x = x + self.attn(self.attn_norm(x), kv_cache=kv_cache, cache_position=cache_position, attn_mask=attn_mask) # x + Attention(RMSNorm(x))
x = x + self.ffn(self.ffn_norm(x)) # x + FFN(RMSNorm(x))
return x
説明メモ
- (1) max_cache_len (最大系列長)サイズのテンソルKとVを
torch.zerosで作成します。形状は（バッチサイズ，ヘッド数，最大系列長，特徴量次元）となります。forループでtransformer層の層数だけキャッシュ用のテンソルを作成しています。 - (2) updateメソッド：テンソルKとVの必要箇所へ値を追記する役割。動的KVキャッシュのtorch.catに対応する部分、事前に作成したテンソルのcache_positionにindex_copy_()を利用して値をコピーするだけです。
- (3) resetメソッド：一度作成したオブジェクトをzero初期化して再利用するための役割。
zero_()で上書きします。 - (4) kv_cacheの有無で判定。キャッシュがない場合（学習時）は、通常の因果マスク付き自己注意機構となります。scaled_dot_product_attentionのオプションであるis_causal=Trueが使えます。キャッシュがある場合（推論時）は、後述するマスク（attn_mask）を使います。マスクを指定する場合は、is_causal=Trueが使えません。
ネットワーク全体を記述します。この段階でマスクも作成します。
class DNN(nn.Module):
def __init__(self, config: ModelConfig):
super().__init__()
self.config = config
# 絶対位置 埋め込み (RoPEは次回以降)
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)
# causal transformer layer
# 各層のlayer_idxを指定、キャッシュ更新などで利用
self.transformer_blocks = nn.ModuleList([TransformerBlock(config, i) for i in range(config.num_layers)])
# 最後の出力に向けた正規化と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_()
# static KV cache化対応した部分
def forward(self, x, kv_cache=None, cache_position=None):
bs, seq_len = x.shape
# (1) 学習時
if kv_cache is None:
assert seq_len <= self.config.seq_len, "位置埋め込みの範囲外 seq_lenを超えています"
positions = torch.arange(seq_len, device=x.device)
attn_mask = None # 学習時はmaskなし is_causal=True部分で作成
# (2) 推論時
else:
# 【重要部分】キャッシュ利用時のマスク作成
positions = cache_position # cache_position：現在位置
key_positions = torch.arange(kv_cache.max_cache_len, device=x.device)
attn_mask = key_positions[None, :] <= cache_position[:, None]
attn_mask = attn_mask[None, None, :, :] # 形状：(1, 1, T, max_cache_len) or reshape(1,1,*attn_mask)
# 位置埋め込み
x = self.token_embedding(x) + self.pos_embedding(positions).unsqueeze(0)
x = self.dropout(x)
# キャッシュは上書き更新される
for block in self.transformer_blocks:
x = block(x, kv_cache=kv_cache, cache_position=cache_position, attn_mask=attn_mask)
# 推論時は最後のtokenのlogitsだけ使えばよい（decode時は seq_len=1 なので何もしない）
if kv_cache is not None:
x = x[:, -1:] # x[:, -1].unsqueeze(1)と同じ、わかりやす方でOK
logits = self.fc(self.norm(x))
return logits
説明メモ
- (1) 学習時はこれまでのコードと同一。さり気なく、assertを入れてみました😆
- (2) 推論時（kv_cashe is not None）の処理。
- prefillの状態では、cache_position = [0, 1, ..., T-1]
- decodeの状態では、cache_position = [現在位置]
-
key_position[None,:] <= cache_position[:,None]実際にマスク作成する部分
key_position.reshape(1,-1) <= cache_popsition.reshape(-1,1)と同じ意味でtorch.tensorのブロードキャストの性質を使ってマスクを作成します。現在位置が2の場合、次のような形でTrue/Falseマスクが作成されます。
マスクの具体例key_position[None,:] <= cache_position[:,None]で比較されるテンソル [[0, 1, 2, 3, 4], [[0, 0, 0, 0, 0], [0, 1, 2, 3, 4], <= [1, 1, 1, 1, 1], [0, 1, 2, 3, 4]] [2, 2, 2, 2, 2]] 作成されるAttention Mask [[True, False, False, False, False], [True, True, False, False, False], [True, True, True, False, False]]
3. ネットワークの全体像のまとめ
【クリックで展開】ネットワーク構造の全体像となります。
import torch
import torch.nn as nn
from tokenizers import Tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# --- ModelConfigは単純にモデルの数値を記述したクラス
class ModelConfig:
def __init__(self, tokenizer):
# モデル構造
self.vocab_size = tokenizer.get_vocab_size()
self.seq_len = 256
self.d_model = 256
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
# --- ここからがKVキャッシュの内容
import math
import torch.nn.functional as F
class StaticKVCache:
"""
Static KV Cache
1. 事前に、全層の K/V を (batch, nhead, max_cache_len, head_dim) で確保
2. Attention には常に max_cache_len の長さ、未使用部分はマスクで隠す
3. 書き込みは cache_position (LongTensor) の位置へ index_copy_ で in-place に行う
"""
def __init__(self, config, batch_size, device, dtype, max_cache_len=None):
self.max_batch_size = batch_size
self.max_cache_len = max_cache_len or config.seq_len
head_dim = config.d_model // config.nhead
shape = (batch_size, config.nhead, self.max_cache_len, head_dim)
# (1) 0で予めKVキャッシュの領域をテンソルとして確保
self.key_cache = []
self.value_cache = []
for _ in range(config.num_layers):
k = torch.zeros(shape, device=device, dtype=dtype)
v = torch.zeros(shape, device=device, dtype=dtype)
self.key_cache.append(k)
self.value_cache.append(v)
# (2) index_copy_() で、k、v cacheのテンソルの指定位置に値を上書き
def update(self, k, v, layer_idx, cache_position):
k_out = self.key_cache[layer_idx]
v_out = self.value_cache[layer_idx]
k_out.index_copy_(2, cache_position, k) # dim=2 (系列方向) に書き込み
v_out.index_copy_(2, cache_position, v)
return k_out, v_out
# KとVのキャッシュを０で初期化（新しく作成ではなくて上書き）
def reset(self):
for k, v in zip(self.key_cache, self.value_cache):
k.zero_()
v.zero_()
class CausalSelfAttention(nn.Module):
def __init__(self, config: ModelConfig, layer_idx: int):
super().__init__()
assert config.d_model % config.nhead == 0
self.config = config
self.layer_idx = layer_idx # Transformerのlayer index
self.nhead = config.nhead
self.head_dim = config.d_model // config.nhead
# q,k,v をまとめて1回のmatmulで計算
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, cache_position=None, attn_mask=None):
bs, seq_len, d_model = x.shape
# q, k, v を計算
q, k, v = self.qkv(x).split(d_model, dim=2)
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)
# (4)
if kv_cache is not None:
# 推論時（prefill / decode）
k, v = kv_cache.update(k, v, self.layer_idx, cache_position)
y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
else:
# 学習時: キャッシュなしの通常の causal attention
y = F.scaled_dot_product_attention(
q, k, v,
is_causal=True,
dropout_p=(self.attn_dropout_p if self.training else 0.0),
)
y = y.transpose(1, 2).contiguous().view(bs, seq_len, d_model)
y = self.resid_dropout(self.proj(y))
return y
# SwiGLUタイプのFFN
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)) # 線形層でFFNと同様に形状を整える出力
# RMSNormタイプのTransformerブロック
class TransformerBlock(nn.Module):
def __init__(self, config, layer_idx):
super().__init__()
self.attn_norm = nn.RMSNorm(config.d_model, eps=1e-6)
self.attn = CausalSelfAttention(config, layer_idx)
self.ffn_norm = nn.RMSNorm(config.d_model, eps=1e-6)
self.ffn = FFN(config)
def forward(self, x, kv_cache=None, cache_position=None, attn_mask=None):
x = x + self.attn(self.attn_norm(x), kv_cache=kv_cache, cache_position=cache_position, attn_mask=attn_mask) # x + Attention(RMSNorm(x))
x = x + self.ffn(self.ffn_norm(x)) # x + FFN(RMSNorm(x))
return x
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)
# causal transformer layer
# 各層のlayer_idxを指定、キャッシュ更新などで利用
self.transformer_blocks = nn.ModuleList([TransformerBlock(config, i) for i in range(config.num_layers)])
# 最後の出力に向けた正規化と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_()
# 重要部分
def forward(self, x, kv_cache=None, cache_position=None):
bs, seq_len = x.shape
if kv_cache is None:
# 学習時
assert seq_len <= self.config.seq_len, "位置埋め込みの範囲外 seq_lenを超えています"
positions = torch.arange(seq_len, device=x.device)
attn_mask = None
else:
positions = cache_position
key_positions = torch.arange(kv_cache.max_cache_len, device=x.device)
attn_mask = key_positions[None, :] <= cache_position[:, None]
attn_mask = attn_mask[None, None, :, :] # (1, 1, T, max_cache_len) reshape(1,1,*attn_mask)
# 位置埋め込み
x = self.token_embedding(x) + self.pos_embedding(positions).unsqueeze(0)
x = self.dropout(x)
# キャッシュは上書き更新される
for block in self.transformer_blocks:
x = block(x, kv_cache=kv_cache, cache_position=cache_position, attn_mask=attn_mask)
# 推論時は最後のtokenのlogitsだけ使えばよい（decode時は seq_len=1 なので何もしない）
if kv_cache is not None:
x = x[:, -1:]
logits = self.fc(self.norm(x))
return logits
config = ModelConfig(tokenizer)
model = DNN(config).to(device)
ネットワークの中身が変更されているので、概要を表示させても変化はありませんが、具体例として掲載してみました。実際に動作確認させる場合はgithubのコードを使用してください。torchinforで概要を表示してみましょう。
from torchinfo import summary
summary(
model,
input_size=(1, config.seq_len), # (batch, seq_len)
dtypes=[torch.long], # 埋め込みの入力は整数インデックス
depth=2,
)
表示結果はこれまでと同じです。埋め込み層、Transformer層、全結合層で語彙数の次元へという流れになります。
===============================================================================================
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
===============================================================================================
4. 生成
KVキャッシュの影響を受けるのは生成時です。文章を生成する関数も静的KVキャッシュへの変更による影響が現れます。Greedyによる文章生成なのでlogits = model(入力文のID) を求めて、argmaxなトークンを加えていく形になります。キャッシュがないテキスト生成関数に、StaticKVCacheでキャッシュ領域を確保、Prefill、Decodeと処理を追加していきます。
@torch.inference_mode()
def generate_text(model, input_ids, max_new_tokens=100, eos_token_id=None,
kv_cache=None):
model.eval()
input_ids = input_ids.to(device)
bs, prompt_len = input_ids.shape
max_len = model.config.seq_len
# 長さのチェック
assert prompt_len < max_len, "プロンプトが seq_len 以上です"
max_new_tokens = min(max_new_tokens, max_len - prompt_len) # 合計長を seq_len 以下に
# autocast の状態に応じて dtype を決定
device_type = input_ids.device.type
if torch.is_autocast_enabled(device_type):
cache_dtype = torch.get_autocast_dtype(device_type)
else:
cache_dtype = next(model.parameters()).dtype
# (1) kv cacheがない場合は領域確保
if kv_cache is None:
kv_cache = StaticKVCache(
model.config,
bs,
device=input_ids.device,
dtype=cache_dtype
)
# (2) すでにキャッシュがある場合は初期化
else:
assert kv_cache.key_cache[0].dtype == cache_dtype
kv_cache.reset()
# (3) Prefill
cache_position = torch.arange(prompt_len, device=input_ids.device)
logits = model(input_ids, kv_cache=kv_cache, cache_position=cache_position)
next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
generated = torch.cat([input_ids, next_token], dim=1)
# (4) Decode
cache_position = torch.tensor([prompt_len], device=input_ids.device)
for _ in range(max_new_tokens - 1):
if eos_token_id is not None and (next_token == eos_token_id).all():
break
# 最新の1トークンだけ入力
logits = model(next_token, kv_cache=kv_cache, cache_position=cache_position)
cache_position += 1
next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
generated = torch.cat([generated, next_token], dim=1)
return generated
説明メモ
- (1) キャッシュがない場合、固定サイズの領域を確保する StaticKVCache を作成します。
- (2) すでにキャッシュが存在する場合は内容を初期化して再利用します。
- (3) Prefill：最初のトークンを生成。
- (4) Decode：2トークン目以降の生成。終端記号 or max_new_tokensまで文字を生成します。
cache_position = torch.tensor([prompt_len])： 次に書き込むキャッシュの位置を、プロンプト直後のインデックス（prompt_len）にセット！
ベタですがcache_position += 1で位置を順番にずらしていきます。
次回
埋め込み部分を修正していきます。回転させる埋め込み方法であるRoPE (Rotary Positional Embeddings) に取り組んで行きます。
目次ページ
参考
静的KVキャッシュに関する参考ページ。
-
Hugging Faceの解説ページ
Cache strategies -
PyTorchのBlog (2023)に登場
Accelerating Generative AI with PyTorch II: GPT, Fast