MiniMind の設定を縮めてモデルを作り、層の決定を数字で確かめる

目標

MiniMindConfigを128次元・4層・Qヘッド4・KVヘッド2・語彙1024に縮小してモデルを作り、パラメーター数を式とモデルで合わせます。GQAのKVキャッシュの節約、RMSNormの出力、RoPEの相対位置の性質、学習前の損失、MoEの全体・アクティブなパラメーターを、MiniMindのコードで直接測ります。

なぜ重要なのか

設定ファイル1枚が、モデルの大きさとサービングのコストを決めます。パラメーターを式で数えられれば、共有された重みや正規化を数え落としていないかがすぐにわかり、KVキャッシュをトークンあたりのバイト数で計算できれば、1枚のカードで何人を受け入れられるかを計算できます。また、最近のモデルが使っている処方(RMSNorm・GQA・RoPE・SwiGLU・埋め込みの共有)には、それぞれ理由があります。このラボは、説明を信じる代わりに、MiniMindの実際のコードで、その性質を数字で確認します。平均を引かない正規化、半分になったK・V、距離だけが残る回転です。

ステップ

小さな設定を保存してください(hidden_size 128、num_hidden_layers 4、vocab_size 1024、num_attention_heads 4、num_key_value_heads 2、max_position_embeddings 512。出力先: /root/mm/arch/config.json)。
パラメーター数を式で計算し、モデルでも数えて、intermediate_size・per_layer・formula・modelとして書いてください(出力先: /root/mm/arch/count.json)。
1つ目の層のアテンションのq_proj・k_projの出力の大きさと、fp32でトークン1つをキャッシュするときのバイト数を、GQAのままの場合と、KVヘッドがQヘッドの数と同じ場合の両方で書いてください(出力先: /root/mm/arch/gqa.json)。
torch.manual_seed(0)のあと、randn(4, 128)*5+3を1つ目の層のinput_layernormに入れて、入力・出力のRMSと出力の平均を書いてください(出力先: /root/mm/arch/rmsnorm.json)。
torch.manual_seed(0)のあと、qとkを1つずつ生成し、MiniMindのapply_rotary_pos_embで(3,7)・(103,107)・(3,50)の位置に置いて、内積を書いてください(出力先: /root/mm/arch/rope.json)。
mmkit.seed_all(0)で作った学習前のモデルの損失を、検証データの先頭16×128トークンで測り、loss・ln_vocabとして書いてください(出力先: /root/mm/arch/init_loss.json)。
同じ設定にuse_moe=Trueを加えたモデルの、全体・アクティブなパラメーターを書いてください(出力先: /root/mm/arch/moe.json)。
## 파라미터는 어디에 ## GQA 와 KV 캐시 ## RoPE 와 RMSNormの3つのセクションを書き、ステップ2のパラメーター数とステップ3のGQAのトークンあたりのキャッシュのバイト数を入れてください(出力先: /root/mm/arch/report.md)。韓国語の見出しは、順に「パラメーターはどこに」「GQAとKVキャッシュ」「RoPEとRMSNorm」という意味です。

参考

モデルはimport mmkit; mmkit.new_model(설정 dict)で作ります。内部では、/opt/minimind/model/model_minimind.pyのMiniMindForCausalLM(MiniMindConfig(**설정))を呼んでいるだけです(2つのコードのプレースホルダーは、どちらも設定です)。コードを直接開いて見てください。
層の部品の名前: model.model.layers[i].self_attn.q_proj、.input_layernorm、.mlp.gate_proj、model.model.freqs_cos。
よくある間違い: 共有されたlm_headをもう一度数えること、q_norm・k_norm(head_dimの大きさ)を数え落とすこと、KVキャッシュを複製(repeat_kv)後の大きさで数えること。
原典: model_minimind.py · RMSNorm · RoFormer(RoPE) · GQA · GLU Variants

小さな設定を1枚書く

hidden_size 128、num_hidden_layers 4、vocab_size 1024、num_attention_heads 4、num_key_value_heads 2、max_position_embeddings 512を書いてください(出力先: /root/mm/arch/config.json)。残りの値は、MiniMindConfigのデフォルトを使います。

MiniMind-3は768・8・6400・8・4です。層の数と次元を減らすと、パラメーターが次元の2乗で減って、CPUでも数分で学習できます。採点ツールは、このファイルで実際にモデルを作ってみます。

パラメーターを式で数える

パラメーター数を式で計算し(formula)、モデルでも数えて(model)、intermediate_size・per_layer・formula・modelとして書いてください(出力先: /root/mm/arch/count.json)。2つの値が同じになる必要があります。

1層は、q・k・v・oの射影と、q_norm・k_norm(それぞれhead_dim)と、gate・up・down(それぞれhidden×intermediate)と、2つのRMSNorm(それぞれhidden)の合計です。ここに、埋め込み(語彙×hidden、出力層と共有なので1回だけ)と最後のRMSNormを足します。intermediate_sizeはceil(hidden·π/64)·64です。

GQAが減らすものを測る

1つ目の層のアテンションのq_proj.out_features・k_proj.out_features・n_repと、fp32でトークン1つをキャッシュするバイト数(K・Vの2つ×層×KVヘッド×head_dim×4)を、GQAのままの場合(kv_bytes_per_token_gqa)と、KVヘッドがQヘッドの数と同じ場合(kv_bytes_per_token_mha)の両方で書いてください(出力先: /root/mm/arch/gqa.json)。

MiniMindのAttentionは、K・Vをnum_key_value_heads個だけ作り、計算するときにrepeat_kvでn_rep回複製します。キャッシュ(past_kv)には、複製前のK・Vが入ります。

RMSNormは平均を引かないことを確かめる

torch.manual_seed(0)のあと、x = torch.randn(4, 128) * 5 + 3を作って、1つ目の層のinput_layernormに入れ、入力と出力のRMS(行ごとのsqrt(mean(x²))の平均)と、出力全体の平均を、rms_before・rms_after・mean_afterとして書いてください(出力先: /root/mm/arch/rmsnorm.json)。

MiniMindのRMSNormは、x * rsqrt(mean(x²) + eps)に重み(最初は1)を掛けます。RMSは1に戻りますが、平均を引かないので、入力の+3が出力に痕跡として残ります。LayerNormなら、平均が0になります。

RoPEは距離だけを残すことを確かめる

torch.manual_seed(0)のあと、q = torch.randn(1,1,1,head_dim)とk = torch.randn(1,1,1,head_dim)を取り出し、(3,7)・(103,107)・(3,50)のそれぞれについて、モデルのfreqs_cos・freqs_sinとapply_rotary_pos_embで、qを最初の位置、kを2つ目の位置に回して内積を取った値を、dot_3_7・dot_103_107・dot_3_50として書いてください(出力先: /root/mm/arch/rope.json)。

apply_rotary_pos_emb(q, k, cos[p:p+1], sin[p:p+1])は、qとkを同じ位置に回すので、qとkを別々に呼び出して、それぞれの位置に回してください。距離が4で同じ2組の内積は、小数点以下5桁まで同じになる必要があります。

学習前の損失がln(語彙)になることを見る

mmkit.seed_all(0)でシードを固定してからモデルを新しく作り、/opt/mm/ref/val.npyの先頭16×128トークン(view(16, 128))の損失(model(x, labels=x).loss)を、loss・ln_vocabとして書いてください(出力先: /root/mm/arch/init_loss.json)。

ランダムに初期化されたモデルは、すべてのトークンにほぼ同じ確率(1/1024)を与えます。そのクロスエントロピーがln 1024 ≈ 6.93です。この値から遠ければ、初期化かラベルの処理が間違っています。

MoEは大きいが、1トークンが使う分は少ない

ステップ1の設定にuse_moe=Trueを加えたモデル(デフォルトでエキスパート4個・トークンあたり1個)の全体のパラメーターと、トークン1つが実際に使うパラメーター(全体−エキスパート1つの大きさ×エキスパートの数+エキスパート1つの大きさ×トークンあたりのエキスパート数)を、num_experts・top_k・total_params・active_paramsとして書いてください(出力先: /root/mm/arch/moe.json)。

エキスパート1つの大きさは、名前にmlp.experts.0.が入ったパラメーターを足せばよいです(層ごとに1つずつなので、すべての層のエキスパート0の合計)。MiniMindのtrainer_utils.get_model_paramsが、同じ計算で「198M-A64M」を出力します。

構成を説明する報告を書く

## 파라미터는 어디에 ## GQA 와 KV 캐시 ## RoPE 와 RMSNormの3つのセクションを書き、ステップ2のパラメーター数(model)とステップ3のkv_bytes_per_token_gqaを数字で入れてください(出力先: /root/mm/arch/report.md)。韓国語の見出しは、順に「パラメーターはどこに」「GQAとKVキャッシュ」「RoPEとRMSNorm」という意味です。

最初のセクションにはFFN・アテンション・埋め込みがそれぞれ何%か、2つ目のセクションにはGQAがキャッシュを何分の1に減らしたか、3つ目のセクションにはステップ4・5で見た性質を、1行ずつ書けばよいです。