事前学習モデルを事後学習しようとしたら入力埋め込みの勾配ノルムが infになって学習できなかったお話し

ことの経緯

llm-jp/llm-jp-4-8b-baseや、このモデルにCPT（継続事前学習）したikedachin/llm-jp-4-8b-base_cpt_v3にSFTをしようとしたら勾配異常が発生しました。

grad normが延々とinf・・・学習できへん😭

チャットテンプレートで使用するspecial tokenを使うと勾配異常が発生してしまうことがわかったので、原因の調査と対策を行います。

CPT時に数%のSFTデータを混ぜてみたり、LRを下げてみたりもしましたが効果なし。まぁ、そうだよね。
ChatGPTに相談しながらチャットテンプレートで使用するトークンの埋め込みを調査してみると、5つのトークンの入力埋め込みがほぼゼロになっていました。怪しい・・・

初期化なしでは最初のステップで勾配ノルムが inf になった一方、これらの埋め込みを再初期化して短い学習を比較したところ3ステップを完了できました。

この記事では、調査方法、実測結果、試した初期化方法をまとめ、CPT有り無しの今治弁モデルの評価結果まで記載したいと思います。特定のトークンの埋め込みベクトルの初期化については少数の合成データによる短い実験・評価について記載していきます。

環境と対象

使った機材はGB10搭載のThinkStation PGX。我が家では「ピジ子」と呼ばれています。
フル稼働状態なので、週に一回の清掃が欠かせません。すごいほこりが出てきます。w

調査対象は、llm-jp/llm-jp-4-8b-base や、継続事前学習した ikedachin/llm-jp-4-8b-base_cpt_v3 です。実は、CPTで学習した内容をどのくらい覚えているのかという調査をしていた時にこの課題に遭遇しました。
CPT前、後のモデルでも起きたことは全く同じでした。
なお、CPT、SFTともにフルパラメータの学習を行っています。

SFT(Supervised Fine Tuning)では、以下のチャットテンプレート用トークンを使用していました。

<|start|> <|end|> <|message|> <|channel|> <|return|> <|call|>

モデルの主な状態は次のとおりです。

項目	値
語彙数	196,608
入力埋め込み	(196608, 4096)、bfloat16
出力層	(196608, 4096)、入力埋め込みとは別の重み
EOS	<|return|>
PAD	<|endoftext|>

このモデルでは入力埋め込みと出力層が重みを共有していないことがわかったため、今回は入力側だけを確認していきます。別のモデルでは重みが共有されている場合があるため、確認が必要です。

• なんで気にしてこの点を書いたかは「ゼロから作るDeepLearning6 LLM編」の110ページを読むとわかる！

前提：nn.Embedding は token ID をキーにベクトルを引く表

今回の「入力埋め込み」は、PyTorchの nn.Embedding が持つ 学習可能な行列 です。各 token ID に行が1つずつ対応しており、IDを行番号として、その行のベクトルを取り出します。 辞書のような対応表と考えると分かりやすいですが、実体は数値の行列です。

今回のモデルでは、行列の形は (196608, 4096) です。196,608個のIDそれぞれに、4,096次元のベクトルが割り当てられています。

例えば <|start|> のIDは 10 だったので、入力として使うベクトルは次のように取得できます。

embedding = model.get_input_embeddings()
token_id = tokenizer.convert_tokens_to_ids("<|start|>")  # このモデルでは10
vector = embedding.weight[token_id]  # weight[10]という1行、形は(4096,)

処理の流れを表すと、次のようになります。

つまり、「特殊トークンの入力埋め込みを初期化する」 とは、そのIDに対応する 埋め込みベクトルをある値に置き換えることです。トークンの文字列やIDを変更したり、語彙を追加したりする処理ではありません。

なお、IDと行は1対1で対応しますが、異なる行に同じベクトルを設定することもできます。また、同じIDの入力ベクトルは出現位置によらず同じ行から取得されます。その後、モデル内部の処理によって文脈に応じた表現になります。

1. トークンと埋め込みの状態を確認する

まず、テンプレートに書いた文字列が実際に単一トークンとして扱われているかを確認します。併せて入力埋め込みと出力層の行のノルムを表示します。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "ikedachin/llm-jp-4-8b-base_cpt_v3"
tokens = [
    "<|start|>", "<|end|>", "<|message|>",
    "<|channel|>", "<|return|>", "<|call|>",
]

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")

vocab = tokenizer.get_vocab()
input_weight = model.get_input_embeddings().weight
output_layer = model.get_output_embeddings()
output_weight = output_layer.weight if output_layer is not None else None

print("input:", tuple(input_weight.shape))
if output_weight is not None:
    print("output:", tuple(output_weight.shape))
    print("重み共有:", input_weight.data_ptr() == output_weight.data_ptr())

with torch.no_grad():
    for token in tokens:
        token_id = vocab.get(token)
        encoded = tokenizer.encode(token, add_special_tokens=False)
        print(f"\n{token}: vocab_id={token_id}, encoded={encoded}")
        print("special登録:", token_id in tokenizer.all_special_ids)

        if token_id is not None and 0 <= token_id < input_weight.shape[0]:
            vector = input_weight[token_id].float()
            print(
                f"input norm={vector.norm().item():.6g}, "
                f"mean={vector.mean().item():.6g}, "
                f"std={vector.std(unbiased=False).item():.6g}"
            )

        if (
            output_weight is not None
            and token_id is not None
            and 0 <= token_id < output_weight.shape[0]
        ):
            print("output norm:", output_weight[token_id].float().norm().item())

今回の調査では、6つの文字列はいずれも単一トークンとしてエンコードされました。<|return|> 以外は all_special_ids に登録されていませんでしたが、テンプレート内ではそれぞれ単一トークンとして扱われていました。

確認結果

通常トークンから4,096行をサンプリングしたところ、入力埋め込みのノルムの中央値は約 1.32 でした。それに対して、対象トークンの結果は以下のとおりでした。

トークン	入力ノルム	入力成分の標準偏差	出力ノルム
<|start|>	計算上 0	約 1.18e-37	約 0.535
<|end|>	計算上 0	約 1.18e-37	約 0.535
<|message|>	計算上 0	約 1.18e-37	約 0.535
<|channel|>	計算上 0	約 1.18e-37	約 0.535
<|call|>	計算上 0	約 1.18e-37	約 0.535
<|return|>	約 0.521	約 0.00814	約 1.397

5トークンの入力ノルムは 0 と表示されましたが、値が非常に小さくてノルム計算中の二乗するところでfloat32の表現可能範囲の下限値を切る「アンダーフロー」が発生したと考えられます。

EOSの <|return|> は他の5つと異なる状態だったため、再初期化の対象から外しました。おそらく、継続事前学習のときに繋がりのない文章を<|return|>で区切っていたのでしょう。というか、私もCPTをするときにトピックが変わるところには<|return|>を入れて学習していました。

2. 追加事前学習の前後で比較する

次に、追加事前学習前の llm-jp/llm-jp-4-8b-base と、追加事前学習後のモデルで同じトークンの埋め込みベクトルを比較しました。前後のトークナイザーで対象文字列とIDをそれぞれ確認し、各モデルの対応する行を比較しています。

その結果、5トークンの入力埋め込みは継続事前学習前からほぼゼロで、継続事前学習後も各成分が完全一致していました。<|return|> の入力行も前後で完全一致でした。一方、5トークンの出力側には小さな変化があり、比較元に対する相対差は約 0.0068 でした。

この結果だけでは、5トークンが事前学習中にほとんど出現しなかったことが想像されます。そこで、これらの埋め込みベクトルを初期化して比較することにしました。

3. 入力埋め込みの初期化方法

これはChatGPTと相談しながら決めました。今回試したのは、通常トークンから参照行を抽出し、その平均ベクトルに小さな乱数を加える方法です。5トークンをそれぞれ異なるベクトルにし、ノルムを参照行の中央値に合わせます。

これは実験用の初期化方法であり、トークンの意味を初期値に与えるものではありません。役割はその後の学習で獲得する必要があります。また、すべてのモデルに適切とは限りません。

以下をモデルとトークナイザーの読み込み後、学習前に実行します。

import math
import random
import torch


def initialize_chat_embeddings(model, tokenizer, sample_size=4096, seed=42):
    tokens = [
        "<|start|>", "<|end|>", "<|message|>",
        "<|channel|>", "<|call|>",
    ]
    weight = model.get_input_embeddings().weight
    output = model.get_output_embeddings()

    if output is not None and weight.data_ptr() == output.weight.data_ptr():
        raise ValueError("入力と出力が重み共有されているため、入力側のみ変更できません")

    vocab = tokenizer.get_vocab()
    protected = {
        token_id
        for token_id in (
            tokenizer.bos_token_id, tokenizer.eos_token_id,
            tokenizer.pad_token_id, tokenizer.unk_token_id,
        )
        if token_id is not None
    }
    ids = []
    for token in tokens:
        token_id = vocab.get(token)
        if (
            token_id is None
            or not 0 <= token_id < weight.shape[0]
            or tokenizer.encode(token, add_special_tokens=False) != [token_id]
        ):
            raise ValueError(f"単一トークンの行を確認できません: {token}")
        if token_id in protected:
            raise ValueError(f"BOS/EOS/PAD/UNKとIDが重複しています: {token}")
        ids.append(token_id)
    if len(set(ids)) != len(ids):
        raise ValueError("対象トークンのIDが重複しています")

    excluded = set(tokenizer.all_special_ids) | set(ids)
    candidates = sorted(
        token_id for token_id in set(vocab.values())
        if 0 <= token_id < weight.shape[0] and token_id not in excluded
    )
    if len(candidates) < 2:
        raise ValueError("参照トークンが不足しています")

    sampled = random.Random(seed).sample(candidates, min(sample_size, len(candidates)))
    refs = weight[sampled].detach().float().cpu()
    if not torch.isfinite(refs).all():
        raise ValueError("参照行に非有限値があります")
    mean = refs.mean(dim=0)
    target_norm = refs.norm(dim=1).median().item()
    if not math.isfinite(target_norm) or target_norm == 0:
        raise ValueError("参照行のノルムが不正です")

    generator = torch.Generator(device="cpu").manual_seed(seed)
    print(f"平均norm={mean.norm().item():.6g}, 参照norm中央値={target_norm:.6g}")
    with torch.no_grad():
        for token, token_id in zip(tokens, ids):
            noise = torch.randn(mean.shape, generator=generator)
            candidate = mean + noise * (target_norm / math.sqrt(mean.numel()))
            candidate *= target_norm / candidate.norm()
            weight[token_id].copy_(candidate.to(device=weight.device, dtype=weight.dtype))
            print(f"{token}: norm={weight[token_id].float().norm().item():.6g}")


initialize_chat_embeddings(model, tokenizer)

実験では、4,096行のベクトルの各要素の平均のノルムが約 0.113、各トークンID毎のノルムの中央値が約 1.320 でした。再初期化後は5行ともノルムが約 1.320 になりました。

既存のモデルで学習済みのトークンを再初期化すると、埋め込みベクトルに保持されていた情報を完全に失います。対象は、トークンの分割や埋め込みの状態を確認したうえで選びます。EOSは学習されていそうだったので、これは初期化せず、他のスペシャルトークンを初期化しました。

4. 短いテスト学習で初期化前後を比較する

同じモデルを条件ごとに読み直し、同じ乱数シードと合成データを使って比較しました。フルパラメータ学習で、各条件の更新回数は最大3回です。学習率は 1e-5、勾配クリッピングの上限は 1.0 に設定しました。

条件	最初のloss	最初の全体勾配ノルム	完了した更新
再初期化なし	8.7768	inf を含むため算出不可	0 / 3
再初期化あり	7.7907	約 144.7	3 / 3

初期化なしでは、最初のステップで入力埋め込み全体と <|start|> 行の勾配ノルムが inf となり、その条件を中止しました。初期化ありでは <|start|> 行の勾配ノルムが最初のステップで約 20.4 となり、3ステップとも有限の勾配ノルムで更新できました。

初期化ありでもクリッピング前の全体勾配ノルムは上限の 1.0 を超えており、各ステップでクリッピングされています。<|call|> は合成データに登場しなかったため、その入力行の勾配は 0 でした。

ここでの inf は勾配ノルムの計算結果です。勾配成分そのものに inf や nan が含まれていたか、成分は有限でもノルム計算中にオーバーフローしたかは、このログだけでは区別できません。

この実験から分かったこと

今回の合成データを使った短い比較では、5トークンの入力埋め込みを再初期化することで、初期化なしでは更新できなかった最初のステップを通過し、3回の更新を完了できました。ほぼゼロの入力埋め込みが勾配の異常に関係していることを支持する結果です。

5. 学習してみる

使用したモデルと学習セットです

学習
model:

• ikedachin/llm-jp-4-8b-base_cpt_v3

data:

• ikedachin/imabari_wiki_qa_v4_reasoning_effort_llmjp4
• ikedachin/std_lang_wiki_qa_v4_validated_w_reasoning_effort_llmjp4
• ikedachin/imabari_wiki_qa_v5_difficult_reasoning_effort_llmjp4
• ikedachin/std_lang_wiki_qa_v5_difficult_reasoning_effort_llmjp4
• llm-jp/AnswerCarefully

お！学習できてそうやん

学習できてそうやん！（２回目）

6. 最後に

ということで、無事に学習することができました。

これで学習可能になったので、今度こそ、ベースモデルと、継続事前学習（CPT）したモデルを比較して新しく学習した記憶がどのくらい残っているか比較してみたいと思います。

実は、今回のきっかけとなった理由はもう一つあって、@katsuki_onoさんのこの記事を読んだからなんですよ。

これを読むと、フレームワークに頼りすぎず、確かに自分でコードを書いた方がいいと思う。
@katsuki_ono さんの目的と僕の目的は違うけれど、自分の手の中にコントロール可能な領域をしっかり持つというのは本当に大事なんですね。めちゃくちゃ実感しました。
きっかけを頂いた@katsuki_ono さんに感謝っ！ありがとうっ！

ではまたっ！