TT Lab
はじめる
学ぶ 学習パス コース

MiniMind — 小さな言語モデルを最初から最後まで自分で学習する

MiniMind の generate で KV キャッシュとサンプリングを数字で測る

TT Labで続きを見る

目標

MiniMindのgenerateで、KVキャッシュがあるときとないときに、グリーディ生成が同じトークンを出すかを確認し、forwardに入ったトークン数を式と合わせ、時間を測ります。キャッシュの実際のバイト数を式と合わせ、温度とtop-pが何を変えるかを、SFTモデルで数えます。

なぜ重要なのか

生成のコストは、キャッシュが決めます。キャッシュがないと、新しいトークンを1つ作るたびに、ここまでの全体を計算し直し、あれば1トークンだけを計算します。その差は、長さが長いほど、2乗で開きます。その代わり、キャッシュはメモリを食い、その大きさは、設定(層・KVヘッド・head_dim)で正確に計算できます。サービングで、1枚のカードで何人を受け入れるかが、この数字です。 選び方は、キャッシュとは別のものです。キャッシュは「同じ結果を安く」で、サンプリングは「どんな結果を」です。2つを混ぜて考えると、「キャッシュをオンにしたら答えが変わった」のような、間違った診断をすることになります。

ステップ

  1. 基準の事前学習モデルで、検証文書の先頭6個をつなげて200トークンに切ったプロンプト(先頭にbos)の後ろを、64トークン、グリーディ生成で、キャッシュありとなしで生成して、同じかを書いてください(出力先: /root/mm/infer/same.json)。
  2. 同じ2回の生成で、forwardに入ったトークン数を数えて書いてください(出力先: /root/mm/infer/count.json)。
  3. 2つの方式の生成時間を(1回ウォームアップしてから3回測って、中央値を)書いてください(出力先: /root/mm/infer/time.json)。
  4. トークン100個を入れたあとのKVキャッシュの大きさを測り、式と合わせて書いてください(出力先: /root/mm/infer/kvbytes.json)。
  5. SFTモデルで、韓国語の質問「ガラム村の特産物は何ですか」に、温度0.3・1.0・1.5で30回ずつ答えさせて、異なる答えの数を書いてください(出力先: /root/mm/infer/temp.json)。
  6. 事前学習モデルに[bos] 민수는(韓国語で「ミンスは」を意味する語です)を入れたあとの次のトークン分布で、MiniMindのtop-pの規則により、p=0.5・0.9・0.99で生き残る候補の数を書いてください(出力先: /root/mm/infer/topp.json)。
  7. ## KV 캐시 ## 온도 ## top-pの3つのセクションを書き、ステップ3のspeedupとステップ4のキャッシュのバイト数を入れてください(出力先: /root/mm/infer/report.md)。韓国語の見出しは、1つ目と2つ目が、順に「KVキャッシュ」「温度」という意味です。

参考

キャッシュがあってもなくても同じ答えになる

スクリプト(/root/mm/infer/cache.py)を書いて、/opt/mm/ref/pretrain.pthを読み込み、/opt/mm/data/pretrain_val.jsonlの先頭6個の文書のtextを空白でつなげて切ったトークンの先頭200個に、bos(1)を付けたプロンプトの後ろを、64トークン、グリーディ(do_sample=False, top_k=0, top_p=1.0, eos_token_id=None)で、use_cache=TrueとFalseの2回生成してください。新しいトークンが同じか(same_ids)と、キャッシュ側の新しいトークン(ids)を、prompt_len・new_tokensとあわせて書いてください(出力先: /root/mm/infer/same.json)。

因果マスクのおかげで、前のトークンのK・Vは、後ろにトークンが付いても変わりません。そのため、取り出して使っても結果が同じです。ここで違うなら、キャッシュではなく、条件(サンプリング・乱数)が違っています。

forwardに入ったトークンを数える

ステップ1の2回の生成で、model.forwardに入ったinput_idsの長さをすべて足して、prompt_len・new_tokens・fed_with_cache・fed_without_cacheとして書いてください(出力先: /root/mm/infer/count.json)。

キャッシュがあれば、最初の1回でプロンプト全体、それ以降は1トークンずつです。なければ、1回ごとにここまでの全体を入れ直します。採点ツールは、書かれた2つの値を、この2つの式で計算した値と比べます。

時間で見る

キャッシュありとなしの64トークンの生成を、1回ウォームアップしてから、それぞれ3回測って、中央値を、with_cache_s・without_cache_s・speedup(なし÷あり)として書いてください(出力先: /root/mm/infer/time.json)。

time.perf_counter()で測ります。最初の実行は、メモリの割り当てや準備のために遅いので、除きます。トークン数では50倍以上の差があるのに、時間はそれほど差が出ません。小さなモデルでは、1回ごとの固定コスト(Pythonやカーネルの呼び出し)が大きいからです。

キャッシュが何バイトかを調べる

基準の事前学習モデルに、ランダムなトークン100個(torch.randint(3, 어휘, (1, 100), generator=시드 0)。プレースホルダーは、順に語彙とシードです)をuse_cache=Trueで入れて、返ってきたpast_key_valuesのすべてのテンソルのバイト数の合計(measured_bytes)と、式で計算した値(formula_bytes = 2 × 層 × トークン × KVヘッド × head_dim × 4)、最初の層のKの形(k_shape)を、tokensとあわせて書いてください(出力先: /root/mm/infer/kvbytes.json)。

MiniMindのAttentionは、past_kv = (xk, xv)を、repeat_kvの前に作ります。そのため、Kの形は(バッチ, 長さ, KVヘッド2, 32)で、GQAで減らした分だけ、キャッシュも小さくなります。

温度を上げるとどうなるかを見る

SFTモデル(/opt/mm/ref/sft.pth)に、韓国語の質問「ガラム村の特産物は何ですか」を、チャット形式(mmkit.chat_text(…, add_generation_prompt=True))で入れ、温度0.3・1.0・1.5ごとにtorch.manual_seed(0)をかけたあと、do_sample=True, top_k=0, top_p=1.0, max_new_tokens=24, num_return_sequences=30で生成して、<|im_end|>の手前までに切った、異なる答えの数を、{"0.3": {"distinct": n}, "1.0": …, "1.5": …}の形で書いてください(出力先: /root/mm/infer/temp.json)。

温度ごとに同じ乱数から出発してはじめて、温度の効果だけを見られます。温度1.5では、30個がほとんどすべて違い、文字が崩れた答えも出ます。まれなトークンが選ばれると、小さなモデルは元に戻れません。

top-pが残す候補を数える

基準の事前学習モデルに[bos] + '민수는'(韓国語で「ミンスは」を意味する語です)を入れた最後の位置のロジットで、MiniMindのgenerateのtop-pの規則(確率の降順の累積和がpを超える位置から切るが、1つずらしてpを初めて超えるトークンは残し、先頭の1つは常に残す)により、p=0.5・0.9・0.99で生き残る候補の数を、{"0.5": n, "0.9": n, "0.99": n}の形で書いてください(出力先: /root/mm/infer/topp.json)。

/opt/minimind/model/model_minimind.pyのgenerateのtop_pの部分の3行を、そのまま写せばよいです。「ミンスは」の後ろには、動詞・場所など、複数の方向があるので、候補が多く残ります。尖った位置(例: 「村の特産物は」の後ろ)と比べてみてください。

推論のコストと選び方の記録を残す

## KV 캐시 ## 온도 ## top-pの3つのセクションを書き、ステップ3のspeedupとステップ4のmeasured_bytesを数字で入れてください(出力先: /root/mm/infer/report.md)。韓国語の見出しは、1つ目と2つ目が、順に「KVキャッシュ」「温度」という意味です。

キャッシュがトークン数では何倍、時間では何倍だったかを並べて書き、その差の理由を1行付け加えてください。