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?

@GeneLab_999

flash-attnの代役がVRAMを18GB溶かしていたので87%削った話

0
Last updated at Posted at 2026-10-24

この記事の対象読者

  • PyTorch で公開モデルの推論コードを読んだり手を入れたりしている中級者

この記事で得られること

  • flash-attn が入らない環境で、自分のモデルの代替経路が二次メモリを食っていないかを判定し、直すか回避するかを決められる
  • nn.MultiheadAttention が need_weights=False でも注意行列を丸ごと作る条件と、その一次情報
  • 推論中の正方形テンソルを洗い出す 20 行ほどの検出関数

この記事で扱わないこと

  • flash-attn 自体のビルド方法
  • Xiaomi-Lens-0 の精度評価やモデル構造の詳細
  • アプリ全体の VRAM 最適化 (量子化やオフロード)

はじめに

先日、Xiaomi-Lens-0 という画像セグメンテーション対応の 4B モデルに 1 ファイルだけの修正を送り、翌日マージされました。flash-attn が無い環境で、マスクを作る段の追加メモリが 3 メガピクセル (MP) の画像で 18.02 GB から 2.28 GB に減っています。

書こうと思った動機は、原因箇所に need_weights=False がきちんと書いてあったことです。公式ドキュメントが推奨する設定で、私も最初は素通りしました。正しい設定のまま 18 GB が消えるなら、同じ罠は多くのコードに潜んでいるはずです。

読み終えた方に考えてほしいのは、オプション依存が欠けたときに自分のモデルがどの経路を通るのか、一度でも測ったことがあるか、という点です。

GPU の数字は PR に載せた RTX 5090 32 GB (コンシューマ向けハイエンド構成の代表例) での実測、仕組みの再現は CPU 2 コア・メモリ 7 GB の Linux で行いました。

この記事では、ひとつのたとえを最後まで使います。Attention の計算は、全参加者が全参加者と当たる総当たり戦です。主役の Flash Attention は対戦結果を暗算で流す名人ですが、欠席することがあります。そのとき舞台に立つ代役が、総当たり表を全部紙に印刷するタイプだった、という話です。

1. 何が起きていたか ─ 3MP で 18GB、8MP で 1 回 30.75GiB

冒頭で、主役が欠席すると代役が舞台に立つという構図を置きました。代役がどれだけ高くつくのかを数字で押さえないと、直す価値があるかを判断できません。

このセクションで分かること: 問題の規模と、どこで測った数字なのか

Xiaomi-Lens-0 は Qwen3.5-4B ベースのマルチモーダルモデルで、生成した <|mask|> トークンをマスクに変換します。

マスクを作るデコーダの PatchSelfAttention 層は、Self-Attention の実装を flash / sdpa / mha の 3 つ持ち、同じ重みを共有しています。修正前は、flash-attn が無い、GPU が sm80 より古い、CPU で動かす、のどれかに当てはまると mha が自動で選ばれていました。この自動の代替経路をフォールバックと呼びます。sdpa は環境変数 XIAOMI_LENS_DECODER_ATTN=sdpa を知る人だけが選べました。

PR の実測です。数字は「モデル読み込み後から、マスク生成段 run_forward が追加で使った VRAM のピーク」で、アプリ全体の使用量ではありません。

入力画像 パッチ数 P 修正前 mha 修正後 sdpa
0.5 MP 1,944 0.86 GB 0.46 GB
1 MP 3,952 2.58 GB 0.83 GB
2 MP 7,776 8.53 GB 1.55 GB
3 MP 11,616 18.02 GB 2.28 GB

環境は Windows 11、PyTorch 2.14.1+cu130、Transformers 5.4.0、flash-attn なし。scikit-image 付属のコーヒーの写真を各サイズに縮小し、「the coffee cup」を探させました。同じ量を測るための骨組みです (infer.py と同じ順に呼びます)。

gpu_mask_pass_peak.py
def mask_pass_extra_gb(model, processor, mask_id, image, text, device):
    inputs, _ = build_inputs(processor, image, text, device)
    ids, _, n_masks = generate_completion(
        model, processor, inputs, max_new_tokens=2048, mask_token_id=mask_id)
    assert n_masks > 0
    torch.cuda.synchronize()
    base = torch.cuda.memory_allocated()
    torch.cuda.reset_peak_memory_stats()
    run_forward(model, inputs, ids, mask_id, device)  # ここが mask pass
    torch.cuda.synchronize()
    return (torch.cuda.max_memory_allocated() - base) / 1e9

8 MP (3840×2560) では、mha 経路は 30.75 GiB を一度に確保しようとして OOM、sdpa 経路は完走してリクエスト全体のピーク 15.05 GB、約 6 秒でした。これは ComfyUI 経由の別テスト (PatchSelfAttention はバイト単位で同一のコピー) で測定範囲が違うため、表とは並べないでください。

2. なぜ二次で増えるのか ─ 総当たり表の印刷代

前節で、画像を 6 倍にすると mha 経路の追加メモリが約 21 倍になることを見ました。増え方の形が分かれば、どの解像度で破綻するかを事前に予測できます。

このセクションで分かること: 自分の GPU で何 MP まで載るかの見積もり方

記号を 1 つずつ導入します。

$P$ (パッチ数、整数) は総当たり戦の参加者数です。画像を 16×16 画素に切った区画の数で、おおむね画素数 ÷ 256 です。

$H$ (ヘッド数、整数) は総当たり表の冊数です。このデコーダでは 16 です。

$b$ (1 要素のバイト数、整数) はマス 1 つのインク量です。bf16 なら 2、fp32 なら 4 です。

M_{\text{表}} = H \cdot P^2 \cdot b

この式が言っているのは、参加者が 2 倍になると印刷する紙は 4 倍になる、ということです。3 MP では $16 \times 11616^2 \times 2 \approx 4.32$ GB が表 1 束分で、実測 18.02 GB はおよそ 4 束分が同時に机に載っている計算です。内訳 (スコア行列、softmax 出力、型変換の一時領域など) までは切り分けていません。

記号 たとえでの役割 3 MP での値
$P$ 参加者数 11,616
$H$ 表の冊数 16
$b$ マス 1 つのインク量 2 バイト
$M_{\text{表}}$ 表 1 束の紙の量 約 4.32 GB

8 MP の OOM も同じ式で説明できます。$P$ が約 3.2 万になると表 1 束だけで 30.75 GiB で、32 GB の GPU にはモデル本体と一緒に載りません。PR の 4 点に、mha は $aP^2 + bP$、sdpa は $cP + d$ を最小二乗で当てはめて見積もりを作りました。

fit_scaling.py
import numpy as np

P = np.array([1944, 3952, 7776, 11616], dtype=float)
mha_gb = np.array([0.86, 2.58, 8.53, 18.02])
sdpa_gb = np.array([0.46, 0.83, 1.55, 2.28])

(a, b), *_ = np.linalg.lstsq(np.stack([P**2, P], 1), mha_gb, rcond=None)
(c, d), *_ = np.linalg.lstsq(np.stack([P, np.ones_like(P)], 1), sdpa_gb, rcond=None)

print("| 画像 | P | MHA予測 | SDPA予測 |")
print("|---|---|---|---|")
for mp in (0.5, 1, 2, 3, 4, 6, 8):
    p = mp * 1e6 / 256
    print(f"| {mp} MP | {p:,.0f} | {a*p*p + b*p:.1f} GB | {c*p + d:.2f} GB |")
画像 P mha 予測 sdpa 予測
3 MP 11,719 18.3 GB 2.30 GB
4 MP 15,625 31.6 GB 3.03 GB
8 MP 31,250 120.3 GB 5.97 GB

4 MP 以上は測定範囲外への外挿です。8 MP の mha は、実際には最初の 30.75 GiB の確保で OOM になるので 120 GB には到達しません。

sdpa の 8 MP 予測にモデル本体 (bf16 で約 8 GB) を足すと約 14 GB で、全体ピーク 15.05 GB と大まかに合います。32 GB の GPU なら、mha 経路の実用上限は 3 MP 前後、sdpa 経路なら上限の 8.4 MP まで追加 6 GB 程度、という判断になります。

3. need_weights=False なのに、なぜ表が印刷されたのか

前節で、代役は総当たり表を丸ごと印刷していると分かりました。しかし原因箇所の need_weights=False は、公式ドキュメント上は印刷しない経路に乗るはずの設定です。ここを解かないと「MultiheadAttention は悪」という誤った結論になります。

このセクションで分かること: nn.MultiheadAttention が注意行列を作る正確な条件

修正前の mha 経路は self.attn(normed, normed, normed, need_weights=False, key_padding_mask=...) と呼ばれていました。PyTorch 2.14 のドキュメントには次の一文があります。

Set need_weights=False to use the optimized scaled_dot_product_attention

和訳: 最適化された scaled_dot_product_attention を使うには need_weights=False を設定してください。

同じページには推論用の高速経路 (fast path) に入る条件も並んでいます。要約すると、自己注意、batch_first の 3 次元入力、autograd 無効、.eval() 済み、などです。infer.py は model.eval() と torch.inference_mode() で推論するので、すべて満たします。

問題は高速経路の中身です。CUDA 実装 native_multi_head_attention_cuda は、SDPA に渡す条件の先頭でマスクの有無を見ています。

if (!mask.has_value() && no_seq_len_1_nested &&

マスクが渡されると SDPA には進まず、bmm_nt でスコア行列を作って masked_softmax をかける経路に落ちます。コメントには形が [B, num_head, T, T] と書かれています。この丸ごとメモリに置く動きを実体化と呼びます。そして Xiaomi-Lens-0 のデコーダは、パディングを示す key_padding_mask を常に渡していました。

CPU 上で 4 条件を 1 回ずつ実行し、呼ばれた aten 演算をプロファイラで拾いました。

probe_path.py
import torch
import torch.nn as nn
from torch.profiler import ProfilerActivity, profile

E, H, P = 256, 16, 512  # 経路の確認が目的なので小さめ
KEYS = ("scaled_dot_product", "softmax", "aten::bmm", "native_multi_head")

mha = nn.MultiheadAttention(E, H, batch_first=True)
x = torch.randn(1, P, E)
for mode in ("eval", "train"):
    getattr(mha, mode)()
    for mask in (None, torch.zeros(1, P, dtype=torch.bool)):
        with torch.no_grad(), profile(activities=[ProfilerActivity.CPU]) as prof:
            mha(x, x, x, need_weights=False, key_padding_mask=mask)
        ops = sorted({e.key for e in prof.key_averages() if any(k in e.key for k in KEYS)})
        print(mode, "mask" if mask is not None else "none", ops)
モード マスク 呼ばれた主な演算
eval なし _native_multi_head_attention, bmm, _softmax
eval あり _native_multi_head_attention, bmm, _masked_softmax
train なし scaled_dot_product_attention
train あり scaled_dot_product_attention

eval・マスクありの _masked_softmax には [1, 16, 512, 512] が渡っていました。CPU はマスクが無くても行列を作る点で CUDA と違いますが、マスクありの推論で表を印刷する点は共通です。

一番驚いたのは、学習モードでは SDPA に乗っている点です。同じ層を同じ引数で呼んでも、表を印刷するのは推論時だけです。学習コードや勾配を有効にした単体テストでは見つからず、推論の本番でだけ顔を出す罠でした。

4. CPU で再現する ─ 印刷代は理論値の約 2 倍

前節で、マスクつきの推論では高速経路がスコア行列を実体化することを確かめました。経路が分かっても、それがメモリの主犯かは量を測らないと言えません。

このセクションで分かること: mha と sdpa のピークメモリが P に対してどう伸びるか

forward 1 回ごとに別プロセスで実行し、/proc/self/clear_refs でピーク常駐メモリ (VmHWM) をリセットしてから増分を取ります。

peak_mem.py
import sys

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

E, H = 1024, 16  # head_dim = 64。実モデルの寸法ではない


def sdpa_like_repo(mha, x, kpm):  # Xiaomi-Lens-0 の _sdpa_self_attention と同じ組み立て
    B, P, _ = x.shape
    q, k, v = F.linear(x, mha.in_proj_weight, mha.in_proj_bias).chunk(3, -1)
    q, k, v = (t.view(B, P, H, E // H).transpose(1, 2) for t in (q, k, v))
    bias = torch.zeros(B, 1, 1, P, dtype=x.dtype).masked_fill(
        kpm[:, None, None, :], float("-inf"))
    out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias)
    return F.linear(out.transpose(1, 2).reshape(B, P, E), mha.out_proj.weight, mha.out_proj.bias)


def mib(field):
    return next(int(l.split()[1]) for l in open("/proc/self/status") if l.startswith(field)) / 1024


if __name__ == "__main__":
    path, P = sys.argv[1], int(sys.argv[2])
    torch.manual_seed(0)
    mha = nn.MultiheadAttention(E, H, batch_first=True).eval()
    x, kpm = torch.randn(1, P, E), torch.zeros(1, P, dtype=torch.bool)
    base = mib("VmRSS")
    open("/proc/self/clear_refs", "w").write("5")  # ピーク (VmHWM) をリセット
    with torch.inference_mode():
        if path == "mha":
            mha(x, x, x, need_weights=False, key_padding_mask=kpm)
        else:
            sdpa_like_repo(mha, x, kpm)
    print(f"{path} P={P}: +{mib('VmHWM') - base:.0f} MiB")
P 表 1 束の理論値 mha 高速経路 sdpa
1,024 64 MiB 172 MiB 44 MiB
2,048 256 MiB 580 MiB 64 MiB
3,072 576 MiB 1,245 MiB 84 MiB
4,096 1,024 MiB 2,165 MiB 104 MiB

理論値は $H P^2 b$ (fp32 で $b = 4$)、PyTorch 2.14.1 で 2 回とも同じ値でした。mha は P が 4 倍で約 13 倍に伸び、理論値の約 2 倍 (スコア行列と softmax 出力の 2 束) に近づきます。sdpa は P が 1,024 増えるごとに 20 MiB の直線です。表を印刷せず、ブロックごとに暗算しているからです。

ここまでのまとめ

  • mha 経路の追加メモリは $P^2$ で伸び、3 MP で 18 GB、8 MP では 1 回の確保が 30.75 GiB に達する
  • 原因は need_weights ではなく、推論時の高速経路がマスクつき入力で SDPA を使わないこと
  • CPU でも同じ形の行列が作られ、sdpa は直線でしか伸びないことを再現できた

5. 直し方 ─ 控室にいたもう 1 人の代役を舞台へ

前節で、sdpa なら総当たり表を印刷しないことを量で確かめました。残る問いは、どう直すのが最小で安全か、です。

このセクションで分かること: 修正案の比較と、実際の差分

案 内容 判断
A mha の呼び方を変えて高速経路を外す 勾配の有無に頼る裏技で、版で挙動が変わりうる
B マスクを渡さない パディングがある入力で意味が変わる
C 既存の sdpa 経路をフォールバックにする 同じ重み、実装済み、1 ファイルで済む

C を選びました。控室にはもともと、印刷しない代役が待機していました。環境変数を知らない人にも自動で回るようにしただけです。差分はこうです (ドキュメント文字列は省略し、シグネチャは 1 行にまとめています)。

     def _self_attention(self, normed, key_padding_mask):
-        if _DECODER_ATTN == "sdpa":
-            return self._sdpa_self_attention(normed, key_padding_mask)
+        if _DECODER_ATTN == "mha":
+            return self._mha_self_attention(normed, key_padding_mask)
         if _DECODER_ATTN == "flash":
             output = self._flash_self_attention(normed, key_padding_mask)
             if output is not None:
                 return output
-        return self._mha_self_attention(normed, key_padding_mask)
+        return self._sdpa_self_attention(normed, key_padding_mask)

警告文も「falling back to SDPA」に直しました。XIAOMI_LENS_DECODER_ATTN=mha で旧挙動も選べます。

6. 出力は変わったか ─ 同じ台本、わずかに違う声

前節で、代役を mha から sdpa に替えました。代役が替われば演技は変わりえます。メモリが減っても出力が壊れていては意味がありません。

このセクションで分かること: 修正前後で何が同一で、何がわずかに違うのか

PR では 6 ケース (猫、コーヒーカップ、23 枚の硬貨、ロケット、旗、推論を伴う依頼) で、生成トークン ID、ボックス、すべてのマスクロジットを比べました。

比較 結果
修正後の既定 と 修正前に sdpa を明示 6/6 でビット単位一致
修正後に mha を明示 と 修正前の既定 6/6 でビット単位一致
sdpa と mha のトークン・ボックス 6/6 で同一
sdpa と mha のマスク (閾値 0.5) の IoU 0.9976 〜 0.9998

上 2 行は配線の確認、下 2 行が経路の違いによる差です。マスクの差は境界にわずかに出る程度で、PR では丸めによる差と判断しました。6 ケースは網羅的な品質評価ではありませんが、差の出る場所と大きさは見えます。

mha から SDPA へ移植するときは真偽値マスクの意味が逆です。key_padding_mask は True が無視、SDPA の attn_mask は True が参加です。

mask_semantics.py
import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = k = v = torch.randn(1, 2, 4, 8)
pad = torch.tensor([False, False, False, True]).view(1, 1, 1, 4)  # True = padding
ref = F.scaled_dot_product_attention(q, k, v, attn_mask=torch.zeros(1, 1, 1, 4).masked_fill(pad, float("-inf")))
for name, m in (("反転して渡す", ~pad), ("そのまま渡す", pad)):
    print(name, (F.scaled_dot_product_attention(q, k, v, attn_mask=m) - ref).abs().max().item())

出力は「反転して渡す 0.0」「そのまま渡す 3.46」でした。Xiaomi-Lens-0 は -inf を足す加算マスクに変換して渡しているので、この罠を踏んでいません。

7. 自分のモデルで確かめる ─ 判定フローと検出関数

前節で、付け替えが出力をほとんど変えないことを確認しました。この罠は Xiaomi-Lens-0 固有ではなく、nn.MultiheadAttention を推論でマスクつきで呼ぶすべてのコードに潜んでいます。

このセクションで分かること: 自分のコードが印刷する代役を呼んでいないかの判定方法

推論と同じ inference_mode で 1 回実行し、最後の 2 次元が等しい大きなテンソルを拾う関数です。

find_square_attn.py
import torch
from torch.profiler import ProfilerActivity, profile

WATCH = ("softmax", "bmm", "baddbmm", "matmul")


def find_square_attention(run, min_len: int = 256):
    with torch.inference_mode(), profile(
        activities=[ProfilerActivity.CPU], record_shapes=True
    ) as prof:
        run()
    hits = set()
    for ev in prof.events():
        if any(w in ev.name for w in WATCH):
            for shape in ev.input_shapes:
                if len(shape) >= 3 and shape[-1] == shape[-2] >= min_len:
                    hits.add((ev.name, tuple(shape)))
    return sorted(hits)


if __name__ == "__main__":
    from peak_mem import sdpa_like_repo

    mha = torch.nn.MultiheadAttention(1024, 16, batch_first=True).eval()
    enc = torch.nn.TransformerEncoderLayer(1024, 16, batch_first=True).eval()
    x, pad = torch.randn(1, 1024, 1024), torch.zeros(1, 1024, dtype=torch.bool)
    print("MHA ", find_square_attention(lambda: mha(x, x, x, need_weights=False, key_padding_mask=pad)))
    print("Enc ", find_square_attention(lambda: enc(x, src_key_padding_mask=pad)))
    print("SDPA", find_square_attention(lambda: sdpa_like_repo(mha, x, pad)))
対象 検出結果
MultiheadAttention _masked_softmax (1, 16, 1024, 1024) ほか
TransformerEncoderLayer _masked_softmax (1, 16, 1024, 1024) ほか
SDPA 経路 なし

標準の Transformer 層でも同じ形が出ました。代役の癖は 1 つのモデルではなく、部品の側にあります。

この関数を作る途中で一度失敗しています。.eval() を呼び忘れた層で試して何も検出されず、「問題なし」と読みかけました。学習モードでは高速経路に入らないからです。検査は必ず本番と同じ eval と inference_mode の組み合わせで行うのが、この罠の一番の教訓でした。

8. OSS に出すときに揃えたもの

前節までで、原因の特定、修正、出力確認、横展開の道具が揃いました。見つけた修正は、メンテナが受け取れる形にしないと他の利用者に届きません。

このセクションで分かること: 小さな修正を短時間でマージしてもらうための PR の組み立て方

PR 本文は「問題 / 変更 / 計測 / 出力比較」の 4 段にしました。計測では測った範囲 (追加ピークであって全体ではない) を明記し、出力比較では配線の確認 (ビット単位一致) と経路差による数値差 (IoU) を分けて書いています。

10 月 19 日に出し、翌 20 日にメンテナがマージしました。変更は 1 ファイル、10 行追加・5 行削除です。差分の小ささより、その差分が正しいと短時間で判断できる材料を添えたことが効いたと考えています。

トラブルシューティング

症状 原因 対処
大きい画像だけ OOM 注意行列が $P^2$ で伸びている 第 7 節の関数で _masked_softmax を探す
need_weights=False なのに重い 推論の高速経路がマスクつきで SDPA を使わない SDPA へ置換する
検出関数が何も返さない 層が学習モードのまま .eval() を呼んでから検査する
SDPA に置換したら出力が壊れた 真偽値マスクの意味が逆 ~mask で反転するか加算マスクにする
置換後に推論結果が毎回揺れる dropout_p が推論でも有効 推論時は dropout_p=0.0 を渡す

用語集

用語 意味 たとえでの対応
フォールバック 本命の実装が使えないときの代替経路 主役が欠席したときの代役
高速経路 nn.MultiheadAttention が推論時に使う専用実装 本番の舞台でだけ出る代役
key_padding_mask 無視するキーを True で示すマスク 欠場者の名簿
SDPA scaled_dot_product_attention。入力に応じてカーネルを選ぶ 暗算で流す代役
実体化 中間結果をメモリ上に丸ごと置くこと 総当たり表の印刷
IoU 2 つのマスクの重なり ÷ 和集合 2 枚の塗り絵の一致度

学習ロードマップ

まとめ

  • flash-attn の代役だった nn.MultiheadAttention は、推論時にマスクを渡すと総当たり表 $H \times P \times P$ を丸ごと印刷していた
  • need_weights=False では防げない。防げない条件は、高速経路のソースでマスクの有無として書かれている
  • 実装済みだった SDPA 経路をフォールバックにするだけで、3 MP の追加メモリが 87% 減り、出力差は丸めの範囲に収まった

主役が欠席する環境は、Windows、古い GPU、CPU と、思ったより広くあります。オプション依存を外した状態で 1 回だけ推論を回し、代役の経路を測ってみてください。

参考文献

関連記事

PyTorch の高速経路の掘り下げを中心にした版を Zenn にも書いています。

関連記事

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?