言語モデルを作る

blog

はじめに

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

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

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

目標

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

とりあえずやってみた

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

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

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

トーカナイザー

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

train_tokenizer.py 
from tokenizers import Tokenizer,normalizers,pre_tokenizers,decoders
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from datasets import load_dataset
import huggingface_hub
import os
import string
import re
import unicodedata

huggingface_hub.login(token=os.environ['HF_TOKEN'])

tokenizer = Tokenizer(BPE(unk_token='<unk>', end_of_word_suffix="</w>", byte_fallback=True))
#tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
tokenizer.normalizer = normalizers.NFKC()
tokenizer.decoder = decoders.Sequence(
    [
        decoders.ByteFallback(),
        decoders.BPEDecoder(suffix=tokenizer.model.end_of_word_suffix),
    ]
)

def clean_text(text):
    # Unicode normalize
    text = unicodedata.normalize("NFKC", text)

    # Normalize line endings
    text = text.replace('\r\n', '\n')
    text = text.replace('\r', '\n')

    # Remove huge blank blocks
    text = re.sub(r'\n{4,}', '\n\n\n', text)

    # Remove separator-only lines
    text = re.sub(r'(?m)^[-_=]{8,}$', '', text)

    # Remove control chars (except \n and \t)
    text = re.sub(r'[\x00-\x08\x0B-\x1F\x7F]', '', text)

    return text

trainer = BpeTrainer(special_tokens=["<unk>", "<bos>", "<eos>", "<pad>", "<mask>", "<|fim_prefix|>", "<|fim_middle|>", "<|fim_suffix|>", "<|fim_pad|>", "<|file_separator|>", "<start_of_image>", "<end_of_image>", "<image_soft_token>"]+[f"<0x{i:02X}>" for i in range(256)], limit_alphabet=4000, vocab_size=64000, max_token_length=10, end_of_word_suffix=tokenizer.model.end_of_word_suffix)
dataset_names = [('agentlans/high-quality-english-sentences', None, None, 'text', 0),('bigcode/the-stack', None, 'data/c', 'content', 10000), ('bigcode/the-stack', None, 'data/c++', 'content', 10000), ('bigcode/the-stack', None, 'data/c-sharp', 'content', 10000), ('bigcode/the-stack', None, 'data/java', 'content', 10000), ('bigcode/the-stack', None, 'data/python', 'content', 10000), ('bigcode/the-stack', None, 'data/ruby', 'content', 10000), ('bigcode/the-stack', None, 'data/go', 'content', 10000)]
#dataset_names = [('agentlans/high-quality-english-sentences', None, None, 'text', 0)]
seed=1141

def dataset_iter():
    for name in dataset_names:
        dataset = None
        key = None
        print(name)
        if None != name[2]:
           count = 0
           dataset = load_dataset(name[0], split='train', data_dir=name[2], trust_remote_code=True, streaming=True)
           for item in dataset:
               count += 1
               if 0<name[4] and name[4]<count:
                   break
               yield clean_text(item[name[3]])
        else:
           count = 0
           dataset = load_dataset(name[0], name[1], split='train', trust_remote_code=True, streaming=True)
           for item in dataset:
               count += 1
               if 0<name[4] and name[4]<count:
                   break
               yield clean_text(item[name[3]])
               #yield remove_non_ascii(item[name[3]])

tokenizer.train_from_iterator(dataset_iter(), trainer)
tokenizer.save('tokenizer.json')

text = '''
ﾜｶﾞﾊｲは㈱である. 吾輩は猫である。名前はまだない。
'''

tokens = tokenizer.encode(text)
print(tokens.tokens)
print(tokenizer.decode(tokens.ids, skip_special_tokens=False))

アーキテクチャ

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拡張の方はほとんど完成していたりする. こちらも, 独自の特徴の案を試してみたい.