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?

12GB×2 の Blackwell + WSL2 で exllamav3 の Tensor Parallel を動かすまで  「WSL では NCCL 無理」の壁と gather カーネルのレース

0
Posted at

TL;DR

  • RTX 5070 (12GB) ×2 = 合計24GB という制約で、27B級のコーダーモデルを実用速度・長コンテキストで動かしたい。
  • モデル本体は EXL3 4.15bpw、KVキャッシュは cache_mode: 2,2(K/Vとも2bit級)量子化 が決め手で、24GBに27B+128Kキャッシュが収まった。
  • さらに2枚のGPUを1モデルで使い切る Tensor Parallel (TP) を狙ったが、以下2つの壁を越える必要があった。
    1. 「Blackwell + WSL2 では NCCL / TP は動かない」という定説 → 環境変数3点で突破。
    2. TPロードは通るのに、推論リクエストで必ず illegal memory access で落ちる → 真犯人は rep_pen.cu ではなく gather.cu の自前 gather カーネルのレース。exllamav3 の自前 SHM gather を NCCL の send/recv に置換して回避。
  • 最終的に TP有効・2枚同時稼働で安定運用(本番は安定重視で MTP は見送り、20〜21 T/s 安定)。副産物として独立バグを3件発見。

対象読者は「限られたVRAMでローカルLLMを回している人」「exllamav3 / tabbyAPI でTPや投機的デコーディングを試している人」「WSL2でマルチGPUに挑んでいる人」あたり。以前書いた実容量とコンテキスト長を計算するツールの続き、「計算して収めた先に、どう速く回すか」の記録でもある。

環境: WSL2 (Ubuntu) / RTX 5070 12GB ×2 (Blackwell, sm_120) / NVIDIA driver 596.36 / CUDA 12.8 / exllamav3 v0.0.43 / tabbyAPI。P2P・NVLinkなし(PCIe接続のみ)。


背景:24GBに27Bを「収める」ところまで

デュアル RTX 5070 は各12GB、合わせて24GB。ここに27B級のモデルを、なるべく賢いまま・長いコンテキストで載せたい。ここで効いたのが2段の量子化だった。

  • モデル本体: EXL3 の 4.15bpw。GPTQ/AWQ の一般的な4bitより攻めた bpw を細かく選べるのがEXL3の強みで、27Bを両GPUに分割して載せてもまだKV用の余地が残る。
  • KVキャッシュ: tabbyAPI / exllamav3 の cache_mode: 2,2(KとVをそれぞれ2bit級に量子化)。今回のモデルは hybrid attention 構成(64層中フルアテンションは16層だけ、残りは軽い recurrent 層)なので、KVは1トークンあたり約10KB(両GPU合計) と極端に軽い。おかげで 128Kトークン分のキャッシュを積んでも余裕がある。

「収める」まではEXL3とKV量子化でクリア。問題はその先、2枚のGPUをどう使い切るかだった。

デフォルトの gpu_split(レイヤー分割 = パイプライン的な配置)だと、生成はレイヤーを順番に通るので実質1枚ずつしか動かない(片方が計算している間もう片方は待ち)。2枚を同時に回すには Tensor Parallel(各行列を2GPUで分担して計算し、都度集約)が要る。これが本題。


壁1:「Blackwell + WSL2 では NCCL / TP は動かない」を疑う

TP は GPU 間で頻繁に集約(all-reduce / gather)するので、その通信レイヤ(NCCL)が動くことが大前提。ところがWSL2 + Blackwell では NCCL が動かないというのが半ば定説になっていて、実際こちらでも最初は動かなかった。素の CUDA IPC テストが invalid resource handle で即死する。

原因は WSL の GPU-PV 環境で、従来型の cudaIpc*(legacy CUDA IPC)によるプロセス間メモリ共有が壊れていること。NCCL は同一ホスト・マルチGPUだと既定でこの共有メモリ経路を使うので、そこで死ぬ。

ここを、3つの環境変数(+1つの sysctl) で越えられた。

# 1) PyTorch を cuMem/VMM アロケータに切り替える。
#    プロセス間共有が、壊れている legacy cudaIpc ではなく pidfd_getfd 経由になる。
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# 2) pidfd_getfd を許可する(無いと "Operation not permitted")。永続化推奨。
#    /etc/sysctl.d/99-ptrace.conf に kernel.yama.ptrace_scope=0 を書いて永続化。
sudo sysctl -w kernel.yama.ptrace_scope=0

# 3) ソケット経由(ループバック)に固定。P2Pは使わない。
export NCCL_SOCKET_IFNAME=lo
export NCCL_P2P_DISABLE=1

ポイントは expandable_segments:True。これは元々は「メモリ断片化対策」として知られるフラグだが、内部的には PyTorch のアロケータを cuMem / VMM ベースに切り替える。この経路のプロセス間共有は pidfd_getfd(2) を使うため、壊れた legacy IPC を回避できる。そして pidfd_getfd を通すには ptrace_scope=0 が要る。この組み合わせに気づくまでが長かった。

この状態で最小の all_reduce テストを流すと あっさり "NCCL TEST PASSED"。「Blackwell + WSL2 では NCCL は無理」は、少なくともこの構成では誤りだった。ドライバを新しくする必要も、ネイティブLinuxに移る必要もなかった(むしろドライバは触らない方が安定した)。

余談: 後から振り返ると、同じマシンで vLLM の TP が以前から動いていたのがヒントだった。vLLM は同じ WSL 上でこの手の集約を涼しい顔でこなしていた。「動いている前例があるなら、無理なはずがない」というのが突破口になった。


壁2:ロードは通る、でも推論すると必ず落ちる

env を仕込むと、exllamav3 は TP 有効で モデルロードまで完全に成功する(両GPUにモデルが乗り、"NCCL warmup finished"、APIも起動)。ところが推論リクエストを投げた瞬間に、TPワーカーが illegal memory access で落ちる

GPU assert: an illegal memory access was encountered
  at .../generator/rep_pen.cu:133

最初はスタックトレースを素直に読んで「repetition penalty の CUDAカーネルが、TPで分割された logits を範囲外アクセスしている」と診断した。実際、rep_pen を無効(repetition_penalty=1.0)にすると落ちない場面もあり、状況証拠は揃っているように見えた。

が、これは誤診だった。


真犯人を追う(探偵編)

CUDA_LAUNCH_BLOCKING で「本当の発生地点」を出す

CUDA のカーネル起動は非同期なので、rep_pen.cu:133 は「その時点までに溜まった非同期エラーを最初に検出した場所」に過ぎない。cudaPeekAtLastError() がたまたまそこにあっただけで、実際に壊しているカーネルはもっと手前にある可能性が高い。

そこで CUDA_LAUNCH_BLOCKING=1 を付け、systemd 経由ではなく foreground で直接起動して、エラー検出を発生カーネルまで前倒しさせる。すると真の発生地点が出た。

GPU assert: an illegal memory access ...
  at .../exllamav3_ext/parallel/gather.cu:472

真犯人は rep_pen.cu ではなく parallel/gather.cupg_gather_kernel だった。これは TP で分割された lm_head の出力(フルvocab幅の logits)を、各 rank から出力デバイスに集約する 自前の共有メモリ・ストリーミング gather カーネル(host-pinned メモリを使ったプロデューサ/コンシューマ + cudaLaunchCooperativeKernel によるグリッド同期のリングバッファ実装)。

なぜ greedy だけは絶対に落ちないのか

切り分けで効いたのがこれ。greedy(argmax)サンプリングだけは何をしても落ちない。理由を追うと、greedy には専用の軽量経路 tp_dispatch_lm_head_argmax() があり、各 rank でローカルに argmax を取ってからスカラー1個だけを集約する。つまり問題の pg_gather(フルvocab幅のストリーミング gather)を一切通らない

逆に言うと——rep_pen に限らず、temperature / top_p / top_k など「フルの logits が必要な」サンプラーは全部この同じ pg_gather を経由し、全部同じバグを踏む。本バグは「TP + 非greedyサンプリング全般」の問題であって、rep_pen 固有ではなかった。最初の誤診はここを見誤っていた。

printf を足すと、バグの位置が動く

pg_gather_kernel の中身を手計算で追っても、単純ケース(リングバッファが一周しない経路)のインデックス計算は正しく見える。算術的な範囲外は見つからない。100%再現するのでレアレースには見えないが……。

決定打になったのは、printf デバッグを仕込んだ瞬間にクラッシュの検出地点が変わったこと。ロジックは1バイトも変えず printf を足しただけで、gather.cu:472 での検出が消え、代わりに NCCL の Watchdog スレッドが別デバイスで非同期に拾う形に変わる。printf を全部外すと、また元の gather.cu:472 にきれいに戻る。

printf の有無だけで挙動が変わるのは、決定的なインデックスバグでは起こらない。 printf はレジスタ割り付け・命令スケジューリング・warp の実行タイミングを変えるので、それに影響される=タイミング / メモリオーダリングに依存するレースコンディションである、というのが強い傍証になった。

疑わしいのは、プロデューサ/コンシューマ間の acquire/release フラグを使ったハンドシェイク自体か、あるいは cudaLaunchCooperativeKernel / grid.sync()WSL2 + Blackwell (sm_120) で完全にはサポートされていないか、のどちらか。ここは残念ながら本機では確定できなかった。compute-sanitizer の racecheck が Blackwell 非対応Device not supported で起動時から警告、無関係なカーネルで cudaErrorNoKernelImageForDevice を連発してログが汚染される)、cuda-gdb もバイナリ欠損で使えず、ライブでレースを捕まえる手段がこの環境には無かった。

(この過程で gather.cu の self-copy 分岐に + dst_offset の加算漏れという別の独立バグも見つけて直した。出力デバイスが物理device0でないときだけ他デバイスの列を上書きする静かなデータ破損。今回の再現ではたまたま踏まなかったが、これはこれで upstream に報告する価値がある。)


解決:自前 gather を捨てて NCCL send/recv にする

根本原因(レース)は潰しきれない。しかしここで発想を変えた。vLLM や sglang は、この vocab 方向の集約を自前カーネルではなく NCCL の標準集合通信に丸投げしている。 ならば exllamav3 の自前 SHM gather も、NCCL の point-to-point send/recv に置き換えてしまえばいい。 レースの温床である自前カーネルを実行パスから外せる。

model_tp_backend.pyTPBackendNCCL に手を入れた。

  • gather(): フルvocabの logits 集約(pg_gather / pg_gather_small の大きい方)を、dist.send / dist.recv ベースの実装に置換。集約順は active_devices の並びに依存させず、**物理device昇順の正規順(gather_devices / ldims)**でイテレートして列順を保証する(ここを間違えると静かに列がずれる)。
  • gather_small(): greedy / 投機的デコーディング(MTP)が使う軽量な argmax 集約経路も、同じパターンで send/recv に置換。これが無いと MTP+TP 併用は model_tp.pyvals.argmax(dim=-1) で別途 illegal access していたpg_gather_small_kernel 側の問題)。

結果、TP推論が非greedyでも安定し、副産物として MTP(投機的デコーディング)とTPの併用も同時に直った。レースの真因は未解明のままだが、実用上は完全に解決した。

注意点として、これで TP は事実上 tensor_parallel_backend: nccl 前提になった。native(NCCL不使用)に戻すと自前 gather のレースが再発しうる。また NCCL の send/recv は初回に P2P コミュニケータの遅延初期化が入るので、起動直後の最初の1〜2リクエストだけ遅く見える(異常ではない)。


ベンチマーク

モデル: Qwopus3.6-27B-Coder-exl3-4.15bpw(vocab 248320)、TP=2。

構成 生成速度 GPU稼働 備考
gpu_split(従来・パイプライン) 基準 実質片側ずつ 2枚あっても同時には回らない
TP のみ 24〜26 T/s 両GPU 67〜70% 同時 狙い通り2枚が同時に動く
TP + MTP 55〜66 T/s 変動 draft accept率 61〜81%、単発は最速

単発のピークは TP+MTP の 55〜66 T/s。ただし本番はこれを採用しなかった


本番構成の決定:速度ピークより「並列時の安定」

MTP(投機的デコーディング)を有効にすると単発は速いが、

  • バッチ窓で繰り返しループに陥るリスクがある、
  • 複数リクエストを並列で捌くと恩恵が薄い(MTP有効だと段階的に完了するのに対し、無効だと8並列がほぼ同時=真のバッチ並列が効く)、

という理由で、**複数ユーザー運用では TP のみ(MTP無効)**の方が総合的に速く・安全だった。MTP無効化で recurrent state の所要VRAMも1スロット約727MB→約148MBに激減し、キャッシュプールをさらに広げられた。

最終的な本番設定(実運用テストで 20.9〜21.3 T/s 安定・エラーなしを確認、これで確定):

# config.yml(抜粋)
model:
  model_name: Qwopus3.6-27B-Coder-exl3-4.15bpw
  tensor_parallel: true
  tensor_parallel_backend: nccl
  gpu_split: []          # TP時は自動
  cache_mode: 2,2        # KVキャッシュを2bit級に量子化(これが24GBに収める鍵)
  cache_size: 131072     # 128Kトークン分
  max_seq_len: 80000
  max_batch_size: 8
  chunk_size: 2048
draft_model:
  draft_mode: disabled   # MTPは並列時の安定を優先して見送り
# ~/.config/systemd/user/tabbyapi.service.d/nccl-tp.conf
[Service]
Environment=PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
Environment=NCCL_SOCKET_IFNAME=lo
Environment=NCCL_P2P_DISABLE=1
Environment=NCCL_SHM_DISABLE=0
Environment=MAX_JOBS=4
# 永続化しておく sysctl
# /etc/sysctl.d/99-ptrace.conf
kernel.yama.ptrace_scope = 0

サンプラーはモデル推奨に合わせて temperature 1.0 / top_k 20 / top_p 0.95。


ついでに見つけた独立バグ3件

TPとは別系統で、深掘り中に踏んだもの。誰かの再現の助けになるかもしれないので記録しておく。

  1. gather.cu self-copy 分岐の + dst_offset 加算漏れ(exllamav3)
    出力デバイスが物理device0でない場合のみ、他デバイスの列を静かに上書きするデータ破損。1行修正で直る。NCCL版に移行した今は実行パス外だが、native backend 利用時は要注意。

  2. tabbyAPI: prefill中に切断されたリクエストで Pydantic ValidationError
    _chat_stream_collector()generation = {} で初期化され、ループ内で初めて index を設定していた。まだ1トークンも出ていない(prefill中の)段階でクライアントが切断すると async for が一度も回らず、index キーの無い空dictが返って ChatCompletionRespChoice(index=None) で落ちる。generation = {"index": task_idx} に変えれば解消。

  3. exllamav3 recurrent checkpoint の共有ページ競合(未修正・スコープ外)
    job.py: maybe_stash_recurrent()assert page.kv_position == PAGE_SIZE が、大量の同時大容量リクエスト+オーバーサブスクライブでのキャンセル/再キューが重なると失敗する。1件の AssertionError が「その瞬間アクティブな全ジョブ」に配信され、まとめて503になるのが厄介。共有ページの ref-count / backup / revert プロトコルの競合とみられ、根が深いので今回は記録のみ。通常の数ユーザー利用ではまず踏まない。


まとめ・教訓

  • 「動かない」という定説は、前例があるなら疑う。 vLLM が同じWSLで動いていた事実が、NCCLの壁を越える確信になった。鍵はドライバでも特別なパッチでもなく、expandable_segments:True + ptrace_scope=0 + NCCL_SOCKET_IFNAME=lo という環境設定だった。
  • CUDAの非同期エラーは、まず CUDA_LAUNCH_BLOCKING=1 で本当の発生地点を出す。 スタックトレースの行番号を鵜呑みにすると誤診する(rep_pen.cu は無罪だった)。
  • 「あるサンプラーだけ落ちる/落ちない」は、共通経路を炙り出す最良のプローブ。 greedy だけ落ちない → 犯人はサンプラーではなく共通の gather、と分かった。
  • printf を足して挙動が変わったら、それはレースの強い傍証。 決定的バグなら計装で位置は動かない。
  • 自前カーネルで消耗するより、枯れた NCCL 集合通信に寄せる。 根本原因を潰せなくても、実用解には最短で届いた。おまけに MTP+TP まで直った。
  • 限られたVRAMの主役は量子化。 EXL3 4.15bpw と cache_mode: 2,2 が無ければ、そもそも24GBに27B+128Kは載らなかった。TPは「載せた後、どう速く回すか」の話。

exllamav3 本体の gather レースは未解明のままなので、upstream(turboderp/exllamav3)には「TP=2 で repetition_penalty か temperature≠1.0 か top_p<1.0 のいずれかを有効にすると必ず gather.cupg_gather_kernel で illegal memory access」「printf 追加で検出地点が変わる(レースの傍証)」という切り分けと、今回の NCCL 回避策をセットで報告するつもり。同じ 12GB×2 / WSL2 / Blackwell で戦っている人の役に立てば。


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
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?