0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@nakatada-lab(Nakamura)

学習lossは正常なのに出力が壊れていた ― PyTorch LayerNormのFP32分散オーバーフローを再現する

0
Last updated at Posted at 2026-11-02

はじめに

独自の重み共有再帰Transformer（7Bパラメータ、単一RTX 3090）を日本語テキストで学習しているときに、NaNもinfも出ないのに、ある深さから先の出力が入力にまったく依存しなくなるという現象に遭遇しました。

原因は、最終段の LayerNorm の分散計算がFP32の上限を超えていたことでした。LayerNormは「値が大きすぎる入力を正規化してくれる層」だと思われがちですが、正規化する前の分散計算そのものが溢れると、出力は静かに定数になります。

この記事では、

  1. 実際のモデルで何が観測されたか
  2. モデルなしで同じ現象を再現する最小スクリプト（CPUで約1分）
  3. なぜ気付きにくいのか、どう検出し、どう直したか

をまとめます。

この記事の結論

  • PyTorchの nn.LayerNorm（FP32）は、1行の「平均からの偏差の二乗和」がFP32の最大値（約3.4×10^38）を超えると、その行の出力が bias だけ（入力に依存しない定数）になります
  • このとき NaNは出ません。ロジットは全入力で同じになり、CEは「それらしい有限の値」を返します
  • 要素1個の二乗はまだ上限の約65分の1でも、行内で合計した時点で溢れます。隠れ次元が大きいほど早く起きます
  • LayerNormの計算だけをFP64にすると、同じ重み・同じ入力で崩壊は消えました

1. 実際のモデルで観測されたこと

対象は、同じブロックを何度も繰り返し適用する重み共有再帰Transformerです（7.0B有効パラメータ、知識蒸留、200Mトークン学習時点）。再帰の深さが増えるほど、隠れ状態の大きさが深さ方向に累積していきます。

200Mトークン時点のチェックポイントで、最終段の LayerNorm（model.norm）に入る直前の隠れ状態を GPU 上で調べた結果です。

観測項目 値
d=8 での LayerNorm 直前の二乗の最大値 2.70×10^38（FP32の最大値 3.4×10^38 に近い）
d=8 での rsqrt(分散) 7.7×10^-19（実質 0）
d=8 でのロジットのバッチ間の広がり 0.000（入力に依存しない出力）
同じ隠れ状態を FP64 の LayerNorm で計算したとき（d=9） 分散の最大値 1.03×10^37、rsqrt が 0 にならない

隠れ状態そのものは発散していません（FP64で計算すれば正常に正規化できます）。FP32のLayerNormの中の分散計算だけが溢れていた、という切り分けです。

その後、model.norm の分散と rsqrt の計算だけを FP64 に置き換えて、0トークンから同じ条件で学習し直しました（0〜200Mトークン）。

指標 FP32 LayerNorm FP64 LayerNorm
崩壊の発生回数 6 0
崩壊せずに使えた再帰の深さ 8段まで 15段まで
200Mトークン時点の検証CE 10.673 9.806

※ FP64版にはFP64化以外の変更も一部含まれているため、CEの差のすべてがFP64化の効果だとは言えません。崩壊が消えたこと自体は、上の同一チェックポイントでの比較（FP32とFP64で同じ隠れ状態を正規化）で確認しています。

なお、その後の学習では、16段の深さで安定して学習できています。

2. モデルなしで再現する

「LayerNormに大きな値を入れたらどうなるか」だけを確かめる最小スクリプトです。入力の大きさ（スケール）だけを変え、同じ重みの FP32 版と FP64 版の LayerNorm を比べます。

実行環境

  • WSL2（Ubuntu）または Linux
  • Python 3 / PyTorch（下の結果は PyTorch 2.15.0、CPU で取得）
python3 ln_overflow_repro.py --json-out result.json | tee result.log

GPUで試す場合は --device cuda を付けてください。CPUで1分前後です。

スクリプト全文

ln_overflow_repro.py
#!/usr/bin/env python3
"""
ln_overflow_repro.py
PyTorch LayerNorm の FP32 分散オーバーフローを最小構成で再現する。

  1. 入力の大きさだけを変えて LayerNorm を通し、出力が
     「正常」→「NaN も出さずに全要素が同じ値（静かな崩壊）」→「NaN」
     と変わる範囲を測る
  2. 静かな崩壊の範囲で、ロジットが入力に依存しなくなり、
     それでも CE が有限の「それらしい値」になることを示す
  3. 同じ重みで LayerNorm だけ FP64 で計算すると崩壊が消えることを示す

実行例（WSL / Linux）:
  python3 ln_overflow_repro.py
  python3 ln_overflow_repro.py --device cuda --hidden 4096
"""

import argparse
import json
import math
import sys

import torch
import torch.nn as nn
import torch.nn.functional as F

TAG = "LNREPRO"
SCRIPT_VERSION = "1.0.0"


class FP64LayerNorm(nn.LayerNorm):
    """分散・rsqrt を含む LayerNorm 全体を FP64 で計算し、出力を元の dtype に戻す。"""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight.double() if self.weight is not None else None
        b = self.bias.double() if self.bias is not None else None
        y = F.layer_norm(x.double(), self.normalized_shape, w, b, self.eps)
        return y.to(x.dtype)


def classify(y: torch.Tensor, ln: nn.LayerNorm):
    """LayerNorm 出力の状態を分類し、崩壊した行の割合を返す。

    LayerNorm の出力は y = z * weight + bias。z（正規化部分）が全要素 0 の行は、
    rsqrt(var) が 0 になり入力の情報が消えた行（出力が bias だけになる）。
    """
    if not torch.isfinite(y).all():
        return "NAN_INF", float("nan")
    z = (y.double() - ln.bias.double()) / ln.weight.double()
    row_dead = z.abs().amax(dim=-1) == 0
    frac = row_dead.float().mean().item()
    if frac == 1.0:
        return "SILENT_COLLAPSE", frac
    if frac > 0.0:
        return "PARTIAL", frac
    return "OK", frac


def logit_stats(h: torch.Tensor, head: torch.Tensor, targets: torch.Tensor):
    """ロジットの入力依存性（バッチ間の広がり）と CE を返す。"""
    logits = h.float() @ head.T
    # 同じ位置でバッチ間にどれだけ違いがあるか。0 なら入力に依存していない
    spread = (logits.max(dim=0).values - logits.min(dim=0).values).max().item()
    if torch.isfinite(logits).all():
        ce = F.cross_entropy(logits, targets).item()
    else:
        ce = float("nan")
    return spread, ce


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
    ap.add_argument("--hidden", type=int, default=1024)
    ap.add_argument("--vocab", type=int, default=32000)
    ap.add_argument("--batch", type=int, default=8)
    ap.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2])
    ap.add_argument("--exp-min", type=float, default=15.0)
    ap.add_argument("--exp-max", type=float, default=23.0)
    ap.add_argument("--exp-step", type=float, default=0.25)
    ap.add_argument("--json-out", default="")
    args = ap.parse_args()

    if args.device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device
    if device == "cuda" and not torch.cuda.is_available():
        print(f"[{TAG}.error] CUDA が使えません", file=sys.stderr)
        return 1

    dev_name = torch.cuda.get_device_name(0) if device == "cuda" else "cpu"
    print(f"[{TAG}.init] version={SCRIPT_VERSION} torch={torch.__version__} "
          f"device={device} ({dev_name}) hidden={args.hidden} vocab={args.vocab} "
          f"batch={args.batch} seeds={args.seeds} FP32_MAX={torch.finfo(torch.float32).max:.3e}")

    n_steps = int(round((args.exp_max - args.exp_min) / args.exp_step)) + 1
    exps = [args.exp_min + i * args.exp_step for i in range(n_steps)]

    results = []
    for seed in args.seeds:
        torch.manual_seed(seed)
        base = torch.randn(args.batch, args.hidden, device=device)
        targets = torch.randint(0, args.vocab, (args.batch,), device=device)
        # 学習済みモデルを模して、LayerNorm の weight/bias を 0/1 以外にする
        ln32 = nn.LayerNorm(args.hidden).to(device)
        with torch.no_grad():
            ln32.weight.copy_(1.0 + 0.1 * torch.randn(args.hidden, device=device))
            ln32.bias.copy_(0.1 * torch.randn(args.hidden, device=device))
        ln64 = FP64LayerNorm(args.hidden).to(device)
        ln64.load_state_dict(ln32.state_dict())
        head = torch.randn(args.vocab, args.hidden, device=device) / math.sqrt(args.hidden)

        for e in exps:
            x = base * (10.0 ** e)
            # 測定用の量は FP64 で計算する（測定側が溢れないように）
            x_c = x.double() - x.double().mean(dim=-1, keepdim=True)
            sq_max = (x_c * x_c).max().item()
            # 行ごとの二乗和の最大値。分散計算の途中でこの規模の合計が現れる
            sumsq_max = (x_c * x_c).sum(dim=-1).max().item()
            with torch.no_grad():
                y32 = ln32(x)
                y64 = ln64(x)
            st32, dead32 = classify(y32, ln32)
            st64, dead64 = classify(y64, ln64)
            sp32, ce32 = logit_stats(y32, head, targets)
            sp64, ce64 = logit_stats(y64, head, targets)
            rec = {
                "seed": seed, "exp": e, "absmax": x.abs().max().item(),
                "sq_max": sq_max, "sumsq_max": sumsq_max,
                "fp32_state": st32, "fp32_dead_rows": dead32, "fp32_spread": sp32, "fp32_ce": ce32,
                "fp64_state": st64, "fp64_dead_rows": dead64, "fp64_spread": sp64, "fp64_ce": ce64,
            }
            results.append(rec)
            print(f"[{TAG}.sweep] seed={seed} scale=1e{e:5.2f} absmax={rec['absmax']:.2e} "
                  f"sq_max={sq_max:.2e} sumsq_max={sumsq_max:.2e} "
                  f"| FP32 {st32:15s} dead={dead32:.3f} spread={sp32:.4f} CE={ce32:.4f} "
                  f"| FP64 {st64:15s} spread={sp64:.4f} CE={ce64:.4f}")

    # seed ごとに、静かな崩壊が始まる/NaN が始まる最小スケールをまとめる
    print(f"[{TAG}.summary] ---")
    summary = []
    for seed in args.seeds:
        rs = [r for r in results if r["seed"] == seed]
        first_partial = next((r["exp"] for r in rs if r["fp32_state"] == "PARTIAL"), None)
        first_silent = next((r["exp"] for r in rs if r["fp32_state"] == "SILENT_COLLAPSE"), None)
        first_nan = next((r["exp"] for r in rs if r["fp32_state"] == "NAN_INF"), None)
        fp64_bad = sum(1 for r in rs if r["fp64_state"] != "OK")
        silent_ce = sorted({round(r["fp32_ce"], 4) for r in rs if r["fp32_state"] == "SILENT_COLLAPSE"})
        s = {"seed": seed, "first_partial_exp": first_partial, "first_silent_exp": first_silent,
             "first_nan_exp": first_nan, "fp64_not_ok_count": fp64_bad,
             "silent_ce_values": silent_ce}
        summary.append(s)
        print(f"[{TAG}.summary] seed={seed} FP32 一部の行が崩壊=1e{first_partial} "
              f"全行が静かに崩壊=1e{first_silent} NaN=1e{first_nan} "
              f"静かな崩壊中のCE={silent_ce} FP64 で OK 以外={fp64_bad}/{len(rs)}")
    print(f"[{TAG}.summary] 参考: ln(vocab)={math.log(args.vocab):.4f}")

    # 崩壊した行 = 「行内の偏差二乗和が FP32 の最大値を超えた行」かを、細かい刻みで照合する
    fp32_max = torch.finfo(torch.float32).max
    for seed in args.seeds:
        torch.manual_seed(seed)
        base = torch.randn(args.batch, args.hidden, device=device)
        ln = nn.LayerNorm(args.hidden).to(device)
        n_rows = 0
        n_mismatch = 0
        e = args.exp_min
        while e <= args.exp_max:
            x = base * (10.0 ** e)
            with torch.no_grad():
                y = ln(x)
            if not torch.isfinite(y).all():
                break
            x_c = x.double() - x.double().mean(dim=-1, keepdim=True)
            over = (x_c * x_c).sum(dim=-1) > fp32_max
            dead = y.abs().amax(dim=-1) == 0
            n_mismatch += (dead != over).sum().item()
            n_rows += args.batch
            e += 0.01
        print(f"[{TAG}.check] seed={seed} 照合した行={n_rows} "
              f"不一致(崩壊 と 偏差二乗和>FP32_MAX)={n_mismatch}")

    if args.json_out:
        meta = {"version": SCRIPT_VERSION, "torch": torch.__version__, "device": device,
                "device_name": dev_name, "args": vars(args)}
        with open(args.json_out, "w", encoding="utf-8") as f:
            json.dump({"meta": meta, "summary": summary, "results": results},
                      f, ensure_ascii=False, indent=2)
        print(f"[{TAG}.saved] {args.json_out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())

実行結果（CPU、PyTorch 2.15.0、hidden=1024、vocab=32000）

seed=0 の結果から、状態が変わる付近を抜き出しました。scale は入力に掛けた倍率、「偏差二乗和」は1行あたりの平均からの偏差の二乗和の最大値（測定用にFP64で計算）です。

scale 入力の最大絶対値 要素の二乗の最大値 偏差二乗和 FP32 LayerNorm 崩壊した行 ロジットの広がり CE（FP32） CE（FP64）
1e17.50 1.30e+18 1.66e+36 1.11e+38 正常 0% 6.8034 11.3182 11.3182
1e17.75 2.31e+18 5.24e+36 3.50e+38 一部崩壊 25% 6.8034 11.0039 11.3182
1e18.00 4.10e+18 1.66e+37 1.11e+39 静かな崩壊 100% 0.0000 10.3897 11.3182
1e19.00 4.10e+19 1.66e+39 1.11e+41 静かな崩壊 100% 0.0000 10.3897 11.3182
1e20.00 4.10e+20 1.66e+41 1.11e+43 静かな崩壊 100% 0.0000 10.3897 11.3182
1e20.25 7.29e+20 5.24e+41 3.50e+43 NaN - nan nan 11.3182

3つのseedのまとめと、照合の結果です（実行ログそのまま）。

[LNREPRO.summary] seed=0 FP32 一部の行が崩壊=1e17.75 全行が静かに崩壊=1e18.0 NaN=1e20.25 静かな崩壊中のCE=[10.3897] FP64 で OK 以外=0/33
[LNREPRO.summary] seed=1 FP32 一部の行が崩壊=1e17.75 全行が静かに崩壊=1e18.0 NaN=1e20.25 静かな崩壊中のCE=[10.367] FP64 で OK 以外=0/33
[LNREPRO.summary] seed=2 FP32 一部の行が崩壊=1eNone 全行が静かに崩壊=1e18.0 NaN=1e20.25 静かな崩壊中のCE=[10.3897] FP64 で OK 以外=0/33
[LNREPRO.summary] 参考: ln(vocab)=10.3735
[LNREPRO.check] seed=0 照合した行=4080 不一致(崩壊 と 偏差二乗和>FP32_MAX)=0
[LNREPRO.check] seed=1 照合した行=4160 不一致(崩壊 と 偏差二乗和>FP32_MAX)=0
[LNREPRO.check] seed=2 照合した行=4008 不一致(崩壊 と 偏差二乗和>FP32_MAX)=0

読み取れることは次のとおりです。

  • 崩壊の条件は「1行の偏差二乗和 > FP32の最大値」と完全に一致しました（3 seed、合計12,248行で不一致0）
  • 崩壊が始まった 1e17.75 では、要素1個の二乗（5.24×10^36）は上限の約65分の1しかありません。溢れているのは合計のほうです
  • 1e18.00〜1e20.00 の約100倍の範囲で、NaNを出さずに全行が崩壊しています。この間、ロジットはどの入力でも同じ（広がり0）で、CEは有限です
  • FP64版は全99点（3 seed × 33点）で正常で、CEも変わりませんでした

3. なぜ起きるのか

LayerNormは、1行（1トークン）の隠れベクトル $x$ について次を計算します。

$$
\mu = \frac{1}{H}\sum_i x_i,\quad
\sigma^2 = \frac{1}{H}\sum_i (x_i - \mu)^2,\quad
y_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}, w_i + b_i
$$

FP32で $\sum_i (x_i-\mu)^2$ が約3.4×10^38を超えると、分散が inf になります。すると

$$
\frac{1}{\sqrt{\infty}} = 0 \quad\Rightarrow\quad y_i = 0 \cdot w_i + b_i = b_i
$$

となり、出力は bias だけ、つまり入力に関係ない定数ベクトルになります。inf が出力まで伝わらないので、NaNチェックには引っかかりません。

ポイントは「合計」で溢れることです。$H$ 個の要素の二乗を足すので、隠れ次元 $H$ が大きいほど、要素1個の大きさが小さいうちに溢れます（上の実験は $H=1024$ です）。

さらに入力を大きくすると、CPUでは 1e20.25 から NaN になりました。一方、GPU（RTX 3090、PyTorch 2.6.0+cu124）では、試した範囲の上限（1e23）まで NaN は一度も出ず、全行が静かに崩壊したままでした。環境によっては、NaN で気付く機会すらありません。

4. なぜ気付きにくいのか

  1. NaNが出ない：torch.isnan や detect_anomaly では検出できません
  2. CEがそれらしい値になる：崩壊中のロジットは全入力で同じなので、CEは語彙サイズで決まる値の近く（このおもちゃ例では ln(32000)=10.37 の近く）に落ち着き、学習ログ上は「高めだが普通の値」に見えます
  3. 崩壊したほうがCEが低く見えることもある：上の表では、正常時のCE（11.32）より崩壊時のCE（10.39）のほうが低くなっています。これは重みがランダムなおもちゃ例の性質ですが、「CEが悪化していないから大丈夫」とは言えないことを示しています
  4. 混合精度でも避けられない：PyTorchの autocast では、LayerNorm は FP32 で計算される演算に分類されています。bf16で学習していても、LayerNormの中はFP32の上限に縛られます

5. 検出のしかた

正規化部分がゼロになった行を数える

崩壊した行では、出力から bias を引いた部分（正規化部分）がすべて0になります。forward hookで数えられます。

import torch
import torch.nn as nn


def attach_ln_collapse_monitor(ln: nn.LayerNorm, name: str, stats: dict):
    """LayerNorm の出力のうち、正規化部分が全要素 0 になった行の割合を記録する。"""

    def hook(module, inputs, output):
        with torch.no_grad():
            y = output.detach().double()
            if module.bias is not None:
                y = y - module.bias.double()
            if module.weight is not None:
                y = y / module.weight.double()
            # 最後の次元以外をまとめて「行」として扱う
            rows = y.reshape(-1, y.shape[-1])
            dead = (rows.abs().amax(dim=-1) == 0).float().mean().item()
            stats[name] = dead

    return ln.register_forward_hook(hook)


# 使い方の例
# stats = {}
# handle = attach_ln_collapse_monitor(model.norm, "final_norm", stats)
# ... forward ...
# print(stats)   # {'final_norm': 0.0} なら正常、0 より大きければ崩壊した行がある

※ weight に0の要素があると割り算で inf/nan になります。学習済みモデルで0がありうる場合は、0の要素を除いて判定してください。

ロジットのバッチ間の広がりを見る

異なる入力に対してロジットがまったく同じなら、その手前で情報が消えています。評価時に「同じ位置のロジットのバッチ間の最大差」を記録しておくと、0になった時点で気付けます（再現スクリプトの spread がこれです）。

6. 対処：LayerNormだけFP64で計算する

実際のモデルでは、最終段の LayerNorm の計算だけを FP64 にしました。入出力の dtype はそのままで、中の計算だけ精度を上げる形です。

import torch
import torch.nn as nn
import torch.nn.functional as F


class FP64LayerNorm(nn.LayerNorm):
    """LayerNorm を FP64 で計算し、出力を入力と同じ dtype に戻す。"""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight.double() if self.weight is not None else None
        b = self.bias.double() if self.bias is not None else None
        y = F.layer_norm(x.double(), self.normalized_shape, w, b, self.eps)
        return y.to(x.dtype)


# 既存の LayerNorm を置き換える例（重みはそのまま引き継ぐ）
# new_norm = FP64LayerNorm(model.norm.normalized_shape, eps=model.norm.eps)
# new_norm.load_state_dict(model.norm.state_dict())
# model.norm = new_norm.to(next(model.parameters()).device)

FP64の最大値は約1.8×10^308なので、FP32では溢れる大きさでも正規化できます。再現スクリプトでも、実際のモデルでも、この置き換えで崩壊が消えました。

注意点：

  • FP64の演算はGPUによっては大幅に遅くなります（GeForceは特に遅い）。実際のモデルでは最終段の1か所だけに適用しました。全層に適用した場合の速度は測っていません
  • これは「隠れ状態が大きくなること」自体を止めるものではありません。値がさらに大きくなってFP64でも足りなくなる可能性や、別の場所（FP32のまま残っている演算）で溢れる可能性は残ります
  • LayerNormは入力を定数倍しても（epsを除けば）結果が変わらないので、正規化の前に入力を縮小する方法も考えられますが、本記事では検証していません

7. この記事の範囲と限界

  • 再現スクリプトはCPU（PyTorch 2.15.0）とGPU（RTX 3090、PyTorch 2.6.0+cu124）で実行しました。崩壊が始まるスケール（一部の行が1e17.75、全行が1e18.0）と「崩壊した行 = 偏差二乗和がFP32の最大値を超えた行」という対応は両方で一致しました（不一致は CPU 0/12,248行、GPU 0/19,200行）。異なったのは NaN になるスケールで、CPU は 1e20.25、GPU は 1e23 まで NaN になりませんでした。CPUとGPUでは乱数の作り方が違うため、同じseedでも入力は同一ではありません
  • 実際のモデルでの観測は、単一のモデル・単一GPU・日本語データでのものです
  • 再現スクリプトのヘッドや重みはランダムです。崩壊時のCEの値そのものに一般的な意味はありません

まとめ

  • LayerNormは「大きな値を正規化してくれる層」ですが、正規化の前の分散計算が FP32 の上限を超えると、NaNを出さずに出力が定数になります
  • 崩壊の条件は「1行の偏差二乗和 > FP32の最大値」で、再現実験では12,248行すべてで一致しました
  • NaNチェックやCEの監視だけでは見落とします。正規化部分がゼロになった行の割合やロジットのバッチ間の広がりを見るのが確実です
  • 隠れ状態が深さ方向に大きくなるモデル（再帰型、非常に深いモデル、残差が累積するモデルなど）では、一度確認しておく価値があると思います

参考

  • 中村忠行「精度差分観測による重み共有再帰Transformerの学習ダイナミクス解明」Jxiv preprint, 2026. doi:10.51094/jxiv.6476（本記事の「1. 実際のモデルで観測されたこと」の数値の出典）
  • 中村忠行「単一GPUによる大規模言語モデルの効率的事前学習: 重み共有アーキテクチャと通常Transformerの比較実験」Jxiv preprint, 2026. doi:10.51094/jxiv.5809
  • PyTorch Documentation: LayerNorm
  • PyTorch Documentation: Automatic Mixed Precision package (torch.amp)
  • Micikevicius, P. et al. (2018). Mixed Precision Training. ICLR.

2026/11/02 追記：GPU（RTX 3090）での実行結果を「3.」と「7.」に追記しました。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?