powered by TechFeed
表示モード
Deep Dive

GPU使用率98%でも7時間かかっていたLLM学習を1時間17分に短縮 — Fused Kernelとパッキング変更だけで約6倍速くなったDatabricksの手法

10月14日、Databricksが「Reducing Qwen training time by 6x on AI Runtime」と題した記事を公開した。この記事では、Qwen3シリーズの27BパラメータモデルをAI Runtime上でファインチューニングする際に、学習時間を7時間44分から1時間17分へと約6倍短縮した最適化手法が詳しく紹介されている。元記事のタイトルは「Qwen」とのみ記載されており、対象モデルの正確な表記については元記事を直接参照されたい。

10月14日、Databricksが「Reducing Qwen training time by 6x on AI Runtime」と題した記事を公開した。この記事では、Qwen3シリーズの27BパラメータモデルをAI Runtime上でファインチューニングする際に、学習時間を7時間44分から1時間17分へと約6倍短縮した最適化手法が詳しく紹介されている。元記事のタイトルは「Qwen」とのみ記載されており、対象モデルの正確な表記については元記事を直接参照されたい。


「GPU使用率98%」でも実は遅い——計測から始める最適化

LLMの学習コストを削減しようとするとき、多くのエンジニアがまず見るのがGPU使用率だ。しかしこの記事の出発点はそこへの疑問にある。

標準的なGPU使用率の指標は「GPUが何かをしていた時間の割合」に過ぎない。 GPUがアクティブな間、どれだけ効率的に計算できていたかは別の話だ。ベースラインの学習ループはGPU使用率98%を示しながら、スループットは9,396 tok/sにとどまっていた。

実態を把握するために使われたのがtorchのトレース機構だ。ステップごとの詳細を取得すると、問題が2つ浮かび上がった。

  • CPUがGPUより遅い:GPUが1ステップを処理するよりも、CPUがそのステップを準備する方に時間がかかっている。本来CPUはGPUより先行して動くはずで、これはGPUがストールしていることを意味する。
  • オペレーション数が多すぎる:1ステップあたりCPU側で約170万回、GPU側で約27万回のオペレーションが発生。CPUがGPUへのオペレーション投入に追いつかず、GPUをアイドルにさせてしまっている。

最大の改善点:Fused Kernelの導入で3.7倍高速化

記事の中で最も効果が大きかった施策がFused Kernelの適用だ。

今回対象となったモデルは、元記事によればトランスフォーマー層に加えて「gated delta network層」を持つハイブリッドアーキテクチャを採用している。このアーキテクチャはQwenシリーズ全般ではなく、この特定モデルに固有の構造として紹介されている点に注意が必要だ。この非標準層が大量の小さなGPUオペレーションを生み出していた。例えばバッチ行列積(aten::bmm)だけで1ステップに31,105回も呼ばれている。

Fused Kernelとは、複数の小さなGPUカーネルを1つにまとめる手法だ。GPU上のHBM(高帯域幅メモリ)とオンダイメモリ間の転送コストを削減し、GPUコアのスケジューリングも効率化できる。

HuggingFaceのtransformersライブラリのQwen実装には対応カーネルが明示されており、以下の2つをインストールするだけで適用できる。

インストール後、GPUオペレーション数は約27万回から約1.8万回(約15分の1)に激減。学習時間は7時間44分から2時間5分へ、3.7倍の改善を達成した。

# pip install でシンプルに導入可能
pip install flash-linear-attention causal-conv1d

注意点として、学習には順伝播・逆伝播両方のカーネルが必要だが、一部のカーネルは推論(順伝播)専用にしか実装されていない。導入前に確認が必要だ。

また、Hugging FaceのKernel Hubを使う方法もあるが、記事ではDockerコンテナにビルド時点でカーネルを組み込む方法を推奨している。Kernel Hubは使っているtorchとCUDAのバージョンに対応したカーネルが存在しない場合、警告なしに低速なフォールバックに切り替わるためだ。


追加の改善:パディング削減でさらに36%短縮

Fused Kernel導入後の2時間5分をさらに削れたのが、バッチのパッキング方式の変更だ。

デフォルトでは2Dの矩形テンソル(最大8,192トークン)を使用しており、短いサンプルが多いバッチでは大量のパディングが無駄になる。そこで複数のサンプルを1次元に詰め込む「1D packed tensor」方式に切り替えた。

ただし、ファインチューニングでこれを行う際はアテンションの漏れに注意が必要だ。プリトレーニングと異なり、ファインチューニングでは「あるサンプルのトークンが、同じバッチ内の別サンプルのトークンをアテンションしてはならない」。これを防ぐために、正しいposition_idsをモデルに渡し、Fused Attentionカーネルが適切なマスクを適用できるようにする実装が必要になる(実装例はこちら)。

この変更によってバッチあたりのシーケンス数が増え、総バッチ数が約36%削減。加えてflash-linear-attentionのJITコンパイルが入力形状ごとに走っていた余分な時間も節約できた。結果として学習時間は2時間5分から1時間17分へ短縮された。


最終結果のまとめ

施策 学習時間 改善率
ベースライン 7時間44分 —
Fused Kernel導入 2時間5分 3.7x
1D Packing追加 1時間17分 さらに約36%短縮
合計 1時間17分 約6x

記事では再現可能なサンプルコードもGitHubで公開されており、AI Runtime上ですぐに試せる。

「GPU使用率が高いから問題ない」という判断がいかに危険かを数値で示しつつ、torchのトレース→ボトルネック特定→既存カーネルの活用という再現性のあるアプローチは、今回のモデルに限らず他のアーキテクチャにも応用できる考え方だ。

詳細はReducing Qwen training time by 6x on AI Runtimeを参照していただきたい。