🌊

PythonでGPU計算やってみた話 ― 沼にハマった記録

に公開
GPU
micromanba
tech

PythonでGPU計算やってみた話 ― 沼にハマった記録

2025年、冬

大規模な計算をGPUで速くしたくて、ここ数ヶ月ずっと格闘してた。その記録を残しておく。

そもそもなんでGPUなのか

何万回もループ回して計算をする必要があった。パラメータ変えながら何十パターンも。

CPUで1回95秒。これを100パターン回すと2時間半以上。

...待てない。

GPUは「同じ計算を大量に並列で回す」のが得意。パラメータ違いの計算をまとめて流せば、理論上は爆速になるはず。

そう思って始めたんやけど、まあ色々あった。

PyTorch、お前はダメだった

最初はPyTorchを検討した。機械学習界隈で実績あるし、テンソル操作も直感的やし。

でも断念。

理由その1：デカすぎる

インストールで約2GB。研究用のクラスタ環境で「2GB入れてください」とか言いにくい。

理由その2：メモリ管理がしんどい

長い計算（8万回ループとか）やると、途中のデータが大きすぎてGPUメモリに載らない。だから「チャンク」って呼ばれる小さい塊に分けて、少しずつGPUに送って処理する必要がある。PyTorchでこれやろうとしたら、activation checkpointingだのCPU offloadingだの、なんか大仰な話になってきた。

理由その3：機能が過剰

今回は一方向の計算だけ。機械学習みたいに「結果から逆算して学習」とかは要らん。PyTorchの売りである自動微分機能、まったく使わない。勿体ない。

PyTorchは「学習ありきの機械学習」には最適。でも「計算回すだけ」の用途には重装備すぎた。

CuPyも検討した

NumPyのAPIをほぼそのままGPUで動かせる。移植コスト低い。

import cupy as cp

x = cp.array([1, 2, 3])
y = cp.sum(x)

悪くない。でも後述の理由でNumbaにした。

結局Numba CUDAに落ち着いた

from numba import cuda
import numpy as np

@cuda.jit
def my_kernel(arr):
    idx = cuda.grid(1)
    if idx < arr.size:
        arr[idx] *= 2

なんでNumbaにしたかというと。

  • 軽い: PyTorchの約1/20（約100MB）
  • JITコンパイル: 初回実行時にCUDAカーネル（GPUで動く関数のこと。OSのカーネルとは別物）にコンパイル、以降キャッシュ
  • 直接カーネル書ける: メモリアクセスパターンを細かく制御できる
  • 余計なもんがない: 抽象化レイヤーなしで直接GPU計算

JITコンパイルって何やねん

「実行するとき」にコンパイルする方式。

普通のコンパイル。

ソースコード → [コンパイル] → 実行ファイル → [実行]

JIT。

ソースコード → [実行時にコンパイル] → 即実行

Numbaの場合、@cuda.jitってデコレータつけた関数は、初回呼び出し時にCUDAカーネル（GPUで動くネイティブコード）にコンパイルされる。2回目以降はキャッシュ使うから速い。

nvcc（NVIDIA CUDA Compiler）を直接叩いてCUDA C/C++書く必要がない。Pythonのまま書ける。これがNumba CUDAの強み。

メモリコピーの話 ― ここが沼

GPUプログラミングで一番ハマったのがメモリ周り。

CPU→GPU転送は遅い

# これが遅い
d_data = cuda.to_device(data)  # CPU → GPU
result = d_data.copy_to_host()  # GPU → CPU

GPUの計算自体は爆速でも、データ転送がボトルネックになる。

実際に計測したら、計算時間の半分がCPU-GPU間のデータ転送だった。GPU計算自体は爆速で終わるのに、データの行き来で同じくらい時間かかってた。

は？って思うやん。俺も思った。

broadcast_to の罠 ― 26秒が1秒になった話

これ、マジでハマった。

2D配列を3D配列に拡張してGPUに送る処理があった。同じデータをN個のモデルで共有するやつ。

遅いコード（Before）:

# 2D → 3D に拡張
p_ext = np.broadcast_to(p_ext, (N_models, p_ext.shape[0], p_ext.shape[1]))
p_ext = np.ascontiguousarray(p_ext)  # broadcast_to は view を返すので実体化

# GPU転送
p_ext_device = cuda.to_device(p_ext)

1回の呼び出しで約26秒。

速いコード（After）:

if p_ext.ndim == 2:
    if N_models == 1:
        p_ext = p_ext.reshape(1, p_ext.shape[0], p_ext.shape[1])  # コピーなし
    else:
        p_ext = np.tile(p_ext, (N_models, 1, 1))  # 実体コピー

# GPU転送
p_ext_device = cuda.to_device(p_ext)

1回の呼び出しで約1秒。25倍速くなった。

なんで遅かったかというと。

  1. np.broadcast_to() はメモリを共有するビューを返す（コピーなし、メモリ効率◎）
  2. でもビューは非連続（non-contiguous）なメモリレイアウト
  3. cuda.to_device() は連続配列を期待
  4. np.ascontiguousarray() で連続化するとき、内部で超非効率なコピーが発生

broadcast_to() はメモリ効率的な良い関数。でもGPU転送前には使うな。最初から np.tile() で実体コピー作った方が速い。

何が起きてたか、部屋の喩えで説明する

broadcast_to() は「散らかった部屋をN個コピー」する。
で、ascontiguousarray() は「部屋を掃除」する関数。

つまり、散らかった部屋N個を全部掃除する羽目になる。そりゃ遅い。

一方 tile() は「部屋を掃除してからN個コピー」。綺麗な部屋をコピーするだけやから速い。

教訓: パフォーマンス問題起きたら、まず arr.flags['C_CONTIGUOUS'] を確認。

print(arr.flags['C_CONTIGUOUS'])  # True=片付いてる、False=散らかってる

チャンキングの話

長い計算やると、全データがGPUメモリに載らない。だから小分けにして処理する。

24万回ループ / 1万回ずつ = 24チャンク（24回に分けて処理）

つまり24回のCPU→GPU転送が発生。ここがボトルネック。

ストライドメモリアクセスの罠

ストライドメモリアクセスっちゅうのは、メモリを飛び飛びに読むこと。1個飛ばし、2個飛ばし、みたいな。

もう一つハマったのがメモリレイアウト。

計算の途中結果を格納する配列が [A0, B0, A1, B1, A2, B2, ...] みたいに交互に並んでた（インターリーブ配置）。AとBを分離するのに。

A = y[:, 0::2]   # 0番目から2個おき → 0, 2, 4, 6, ...
B = y[:, 1::2]   # 1番目から2個おき → 1, 3, 5, 7, ...

これ、GPUにとって最悪。メモリアクセスが飛び飛びになる。

メモリ上: [A0][B0][A1][B1][A2][B2]...

Aを読む:  [A0]    [A1]    [A2]    ← 飛び飛び

GPUは連続したメモリを一気に読むのが得意。飛び飛びだとメモリ帯域の50%が無駄になる。

解決策

ループの外で変換する。

# 最初に1回だけ
y_block = convert_to_block(y)  # [A0,A1,...,A67,B0,B1,...,B67]

for step in range(80000):
    A = y_block[:, :68]   # 連続アクセス
    B = y_block[:, 68:]   # 連続アクセス
    ...

# 最後に1回だけ戻す
y_result = block_to_interleaved(y_block)

ループ毎に .contiguous() 呼んでたのが、ループの前後で2回だけになる。8万回×2 = 16万回のメモリコピーが、たった2回に。

環境構築の沼 ― mambaで救われた

GLIBC地獄

研究用クラスタって、マシンによってOSバージョンがバラバラ。

新しいマシンでビルドしたPythonが古いマシンで動かない。

/lib/x86_64-linux-gnu/libc.so.6: version `GLIBC_2.28' not found

GLIBCは後方互換あるけど前方互換ない。新しいGLIBCでビルド → 古いGLIBCでは動かない。

minicondaでも解決できなかった

minicondaで環境作っても、結局システムのGLIBC/libstdc++に依存してしまう。

micromambaで救われた

micromamba自体が静的リンクでビルドされてる。システムのGLIBC/libstdc++に依存しない。

# インストール（どのマシンでも動く）
curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj bin/micromamba

# 環境作成
./bin/micromamba create -n gpu_env python=3.11 numba

これでクラスタ全ノードで統一環境が作れた。神。

実測結果

項目 時間
GPU計算本体 9-18秒
データ転送 9-20秒
合計 18-38秒

CPUだと95秒だったから、まあ速くはなった。でも「GPU計算と同じくらいデータ転送に時間かかってる」のはなんかモヤる。

高速化率は2.8倍〜13.4倍。モデルによってバラつく。

学んだこと

  1. PyTorchは万能じゃない: 計算回すだけの用途には過剰
  2. Numba軽くて良い: 100MBでCUDAカーネル書ける
  3. メモリ転送がボトルネック: GPU計算より転送の方が遅いことがある
  4. メモリレイアウト大事: 飛び飛びアクセスはGPU殺し
  5. 環境構築は沼: micromambaに救われた

まとめ

GPUプログラミング、計算速くするより環境構築の方が大変だった説ある。

でも一回環境整えたら、パラメータスイープとか一気に回せるようになって、研究の幅は広がった。

次はデータ転送のボトルネックなんとかしたい。非同期転送とか、転送と計算のオーバーラップとか、まだ試せることはある。

まあ、それはまた別の沼の話。

2025年冬、GPU計算と格闘した記録

Discussion