言語モデルを作る

はじめに

興味だけで言語モデルをフルスクラッチで作ってみる.
Visual StudioのC#のAI補完のように, コンテキストを踏まえたコード補完は非常に便利だ.
AIアージェントやヴァイブコーディングより便利で実用的だと思っている.

Visual StudioのC/C++にはAI補完はない, 外部サーバにデータを投げる拡張はあるのだが（そして実際便利で使っている）, 社内では使えない.

興味も兼ねて言語モデルから自作することにした.

目標

• Fill In the Middle (FIM)に対応した, コード補完専用言語モデル
• ローカルで動作する. 拡張に同梱できる程度のサイズ.
• 気に入らなければユーザが無視すればよい程度の精度
  • 知識も思考も不要, もっともらしい次を予測出来れば充分

とりあえずやってみた

Gemma-3 270mのconfigをベースに, ボキャブラリサイズを6400にした, パラメータ数250M程度の言語モデルを対象に試行錯誤していた.
ボキャブラリサイズを小さくした分, Gemma-3 270mより余裕はあったはずだ.
lossが3.2前後から下がらないし, 大した成果がないように思えた. 半年ぐらい浪費した.

ChatGPTとやり取りしながら進める

私は機械学習は専門ではないので, 言いなり状態だ. Google検索より早いので裏取りもそこそこに, 750M程度のアーキテクチャに落ち着いた.
300M前後だとそもそも英文法を詰め込むにも足りないらしい.

トーカナイザー

ボキャブラリサイズは6400だが, データセットのクリーニングは修正した. 以前はASCII以外は除去するだけだった.
コメントのエンコード効率も考慮してASCII以外も残し, 代わりに連続した改行・記号を圧縮したりしてみた.
英語テキスト50%, C/C++コード50%の割合で構築した. コード補完が目的なので英語以外は不要なのだ.

アーキテクチャ

flash attention対応を入れつつしていると肥大化した.
full attentionが元のGemma-3 270mより大幅に増えているので推論速度が遅くなりそうだ.

config.json

{
  "_sliding_window_pattern": 4,

  "architectures": [
    "Gemma3ForCausalLM"
  ],

  "model_type": "gemma3_text",
  "transformers_version": "5.5.0",

  "vocab_size": 64000,

  "bos_token_id": 2,
  "eos_token_id": 1,
  "pad_token_id": 0,

  "dtype": "bfloat16",

  "hidden_activation": "silu",

  "hidden_size": 1280,
  "intermediate_size": 5120,

  "num_hidden_layers": 28,

  "num_attention_heads": 20,
  "num_key_value_heads": 5,
  "head_dim": 64,

  "attention_bias": false,
  "attention_dropout": 0.0,

  "query_pre_attn_scalar": 64,
  "initializer_range": 0.02,
  "rms_norm_eps": 1e-6,
  "max_position_embeddings": 8192,
  "layer_types": [
    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention",

    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention",

    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention",

    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention",

    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention",

    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention",

    "sliding_attention",
    "sliding_attention",
    "sliding_attention",
    "full_attention"
  ],
  "sliding_window": 2048,
  "rope_parameters": {
    "full_attention": {
      "rope_type": "default",
      "rope_theta": 100000.0
    },

    "sliding_attention": {
      "rope_type": "default",
      "rope_theta": 100000.0
    }
  },
  "attn_logit_softcapping": null,
  "final_logit_softcapping": null,
  "tie_word_embeddings": true,
  "use_cache": true,
  "use_bidirectional_attention": false
}

AutoModel

RTX 4070 Ti Superが１枚なので, VRAM16GiBで戦わなければいけない.
初期はHuggingfaceのTransformersのレシピで学習していたが, Nvidia Nemo AutoModelで8bit学習を導入した.
flash attentionと合わせてVRAM16GiBでもぎりぎり動かせるようになった.
8bit Adamはエラーが出て解決できなかった.

まとめ

途中も途中で, いつ学習が終わるんだ状態だか. とりあえず以前よりlossは下がっている.
Visual Studio拡張の方はほとんど完成していたりする. こちらも, 独自の特徴の案を試してみたい.