AI RuntimeでQwenのトレーニング時間を6分の1に短縮

数十億パラメータ規模の基盤モデルのトレーニングやファインチューニングには、多くの場合数百から数千時間のGPUトレーニングに及ぶ膨大な計算リソースが必要です。しかし、LLMのトレーニングパフォーマンスのチューニングは容易ではなく、GPU使用率などのハイレベルなメトリクスは実際の利用効率を示していない可能性があります。本ブログ記事では、抽象化のレイヤーを排除し、AI Runtime上で最新のQwen 3.8 27Bモデルのファインチューニングを6倍高速化する方法を紹介します。体系的なパフォーマンス最適化により、GPU使用率がすでに98%に達している状態からでも、固定データセットのファインチューニング時間を7時間44分から1時間17分に短縮できました。

重要な指標の計測方法、実装の詳細、各最適化における設計のトレードオフなど、これらの改善の背景にある技術について詳しく解説します。

すぐに始めたい場合は、AI Runtimeのサンプルコードに直接アクセスしてお試しください！

モデルのセットアップ

まず、transformersライブラリとSFTトレーニングデータセットHuggingFaceH4/ultrachat_200kを使用した、Qwen 3.8 27Bの標準的なトレーニングループから始めます。

• torchの最新モデルシャーディング技術であるFSDP2を適用します。これはモデルの重みとオプティマイザの状態をGPUに収めるために必要です。単純な分散データ並列（DDP）アプローチでは、このサイズのモデル重みを単一のGPUに収めることができません。
• また、アクティベーションチェックポインティングも有効化します。これは逆伝播時のメモリ容量と計算量をトレードオフにすることで、GPUあたりのバッチサイズを大きくできる技術です。
• データローダーは、バッチあたりの最大トークン数を8192として2Dバッチをパッキングし、バッチ例の最大サイズに合わせてパディングを行います。

初期のトレーニングベンチマーク

全データセットに対するモデルのベースライントレーニングには7時間44分かかり、トレーニング中のGPUの平均使用率は98%を示しています。結果は9,396 tok/sのトレーニング速度です。

このGPU使用率は高く見えますが、先述の通り、この数値は誤解を招く可能性があります。標準的なGPU使用率メトリクスは、GPUが何らかの処理を行っていた時間の割合しか測定しないため、アクティブな間にモデルがGPUをどれだけ効率的に使用していたかはわかりません。これを理解するには、より詳しく調査する必要があります。

その方法の1つがtorchのトレース機構です。以下のようにトレーニングループを調整できます（素早く結果を得るためにステップ数も制限します）：

サンプルステップでは、次のようなデータが得られます。

表1：CPU - 1,714,382回の呼び出しで14.4秒。合計時間上位の操作を表示。

表2：GPU - 269,448回の操作で10.3秒。合計時間上位の操作を表示

改善の領域

これらの表から、現在のモデルにおけるいくつかのハイレベルな課題が見えてきます。

• CPUが各トレーニングステップをセットアップする時間が、GPUがそれを実行する時間よりも長くかかっている。 一般的にGPU実行のCPU側は非同期であり、CPUはGPUより十分先行して実行されるはずなので、これは予期せぬ状態です。主な影響としてGPUのストールが発生しており、ここに改善の余地があります。
• 操作数が多すぎる！ 上記の表では数十万の小さな操作が発生しており、問題はCPUがGPUを常に稼働させるのに十分な速度でGPU操作をキューに投入できていないことです。重要なアイデアは、CPUがGPUにキュー投入する操作の数を減らすことです。これには一般的な手法がいくつかあり、今回の状況に適したものがあります。
  • torchのCUDAグラフは、検出されたCUDAカーネルを記録し、単一のCPU呼び出しで再実行します。強力ではあるものの、多くの制約があり、FSDP2の動的な性質と衝突します。そのため、現時点では採用できません。
  • Torch.compileはコードをトレースし、Tritonを使用して可能な場合にカーネルを自動的に融合（フューズ）します。この例では手動で融合したカーネルを使用するため、LLMのケースではtorch.compileによる改善の余地は残されていません。
  • 手動フューズドカーネル（融合カーネル）は、torchのコンパイラが単独では適用できない特定操作用の専門的な融合技術を使用します。幸いなことに、コミュニティには確立された構築済みカーネルが多数存在します。例として「Flash Attention」があり、今回のケースでもパフォーマンスがわずかに向上します。

また、バッチ内の過剰なパディングに関する課題もあります。入力テンソルは長方形である必要があり、各行はそのバッチの最大トークン数までパディングされます。これを最小限に抑えるためのLLMのテクニックも紹介します。

以降のセクションでは、各技術トピックの概要、推奨されるベストプラクティス、適用時に考慮すべき主なトレードオフについて詳しく掘り下げていきます。

フューズドカーネル（融合カーネル）

特定のカーネルが非常に多く実行されていることにお気づきかもしれません。特にQwenでは、バッチ行列積カーネルが1ステップで3万回以上実行されています。

このケースでは、まずフューズドカーネルについて検討します。フューズドカーネルは、複数の小さなGPUカーネルを結合して1つのユニットにします。これにより、GPUの高帯域幅メモリとさらに高帯域なダイ上メモリとの間のコストの高いデータ転送を回避し、パフォーマンスの向上に貢献します。また、GPUコアのより効率的なスケジューリングも可能になります。

これほど多くの追加カーネルがあると、torch.compile は動作する可能性はありますが、実行に非常に時間がかかります。そのため、手動で融合したカーネルを使って純粋な演算（ops）数を減らすことが良い最初のステップとなります。また、Flash Attention のようなよく知られた融合演算は、torch.compile 単体では自動検出されません。

transformers ライブラリはデフォルトで、標準的なアテンション層に torch の SDPA を使用します。しかし、ここで見られるこれらの小さな操作の多くは、非標準の gated delta network 層に由来していることがわかります。幸いなことに、これらの層にはレイヤーごとにこれらの操作の多くを融合するために書かれた 2 つのカーネル、causal-conv1d と flash-linear-attention が存在します。これをどのように発見するのでしょうか？Qwen の transformers ライブラリ実装 は、サポートしているカーネルをあらかじめ示してくれています：

• flash-linear-attention
• causal-conv1d

これらのカーネルを取得するには、主に 2 つの方法があります：

• カーネルを自分でインストールする。これには以下の方法があります：
  • PyPI からの簡単な pip install（例: flash-linear-attention）
  • 該当リポジトリの git clone とコンパイル
  • コンパイルに時間がかかりすぎる場合に備え、事前コンパイル済みパッケージの別のソースを探す
• Hugging Face の Kernel Hub を使用して事前コンパイル済みカーネルを取得する。

Databricks では通常 1 つ目の方法を推奨しています。特に当社の Docker コンテナサポート と組み合わせ、ビルド時にすべてのカーネルをインストールした状態でコンテナを構築することで、ジョブの起動時間を短縮できます。

Kernel Hub は便利ですが、カーネルはインストールされている torch や CUDA のバージョンに明示的に依存することが多いため、特定の構成で利用可能なカーネルが存在しない場合があり、transformers モデルが低パフォーマンスな未最適化カーネルにサイレントにフォールバックしてしまう可能性があります。

flash-linear-attention と causal-conv1d の Python パッケージをインストールした後、エンドツーエンドの実行時間を 3.7 倍改善し、2 時間 5 分まで短縮できました。以下の表では、前の表の要素ごとの（elementwise）関数に代わって、新しいカーネル（強調表示）が配置されていることがわかります。また、実行された総操作数も 15 分の 1 に大幅に削減されました：

表 3：CPU - 3.9 秒（118,316 回呼び出し）。合計時間の長い上位操作を表示。

表 4：GPU - 3.9 秒（18,060 回呼び出し）。合計時間の長い上位操作を表示。

注意点：トレーニング中には順伝播と逆伝播の両方のカーネルが必要ですが、一部のカーネルは推論の高速化のみを目的として実装されているため、順伝播専用となっています。

まとめ

• 一般的なモデルで利用できる確立された高最適化済みカーネルの活用を検討してください。これらのカーネルを利用することで、比較的少ない工数で大幅なパフォーマンス向上を実現できます。今回のケースでは、トレーニング速度が 3.7 倍向上 し、トレーニング時間を 7 時間 44 分から 2 時間 5 分 に短縮できました。

バッチ処理、トークン化、パディング

Qwen モデルはかなり高速になりましたが、さらに高速化できます。デフォルトのバッチ設定では、メモリ不足を起こさずに収まる最大サイズである 8,192 トークンの 2D 矩形テンソルを使用しています。このアプローチの欠点は、パディングによってテンソル内のかなりのスペースが無駄になってしまうことです。

Qwen が標準的なアテンション機構のみを使用している場合、torch の NestedTensor がしっかりとサポートされており、パディングを完全に削除することでこのコストを削減するのに役立ちます。しかし、flash-linear-attention は、より従来型の 1D パックドテンソルアプローチのみをサポートしています：

しかし、これをただ無作為に行うわけにはいきません。1D テンソルを単にモデルに渡すことは事前学習（pretraining）では許容されるかもしれませんが、ファインチューニング（finetuning）では、同一サンプル内のトークン同士のみがアテンションを向けるように保証する必要があります。これを解決するには、正確なサンプルの position_ids をモデルに渡す必要があります。モデルはこれらを融合アテンションカーネルに渡し、カーネルが正しいアテンションマスクを適用します。こちらのサンプルコードを見ると、これを機能させるために必要な入力の計算方法がわかります。

この一見シンプルな変更により、トレーニング時間は 2 時間 5 分から 1 時間 17 分へと短縮されます。

その理由は 2 つあります：

1. 主な理由は、このパッキングによってバッチあたりのシーケンス数が増加し、バッチの総数が約 36% 削減されるためです。なお、これはデータセット内のサンプルの分布や長さに依存することに注意してください。
2. 2 つ目の理由は、flash linear attention カーネルを入力形状ごとに JIT（Just-In-Time）コンパイルする必要があり、以前は形状が動的だったため、コンパイルに数分余分にかかっていたためです。

トレードオフ： このような変更により入力パイプラインはより複雑になり、モデルの品質に影響を与えるエラーが発生するリスクが高まります。しかし、今回のケースでは 36% の削減効果があるため、十分その価値があります。パディングを削減するその他の手法としては、元のバッチ処理メカニズムを維持しつつ、すべての入力を長さ順にソートしてパディングを最小限に抑える方法もありますが、この並び順によって隣接するサンプル間に予期せぬ相関が生じ、モデルにバイアスがかかる可能性があります。

まとめ： 最大トレーニング可能トークン数（この場合は 8,192）よりも短い入力を持つデータセットで LLM をトレーニングする場合、データをバッチに効率的にパッキングすることで大幅な時間を節約できます。このデータセットでは、トレーニング速度が 1.6 倍向上 し、トレーニング時間が 2 時間 5 分から 1 時間 17 分 に短縮されました。

結論

このブログでは、トレーニング効率を 6 倍に高める方法をご紹介しました。ポイントは、(1) 単なる GPU 使用率だけでなく、モデルがトレーニングのどこで時間を費やしているかという本質的な指標をまず測定すること、(2) Flash Linear Attention などの既知のライブラリを活用してカーネルを最適化・融合すること、(3) モデルへのデータのバッチ供給方法を最適化すること、の 3 点に集約されます。

これを大規模に行うには、高パフォーマンスなパスがデフォルトとなっているプラットフォームが必要です。AIRには、こうしたパフォーマンス向上の多くがあらかじめ組み込まれています。ベース環境には最新のCUDAとFlash Attentionを同梱し、GPU使用率、温度、ヘルス状態などのコアメトリクスを自動追跡し、torchのバージョン（およびプロファイラー）を最新状態に維持します。社内エージェントのGenie Codeも、トレーニングをさらに最適化するこのような機会を探すように調整されています。さらに、Databricksからのデータ読み込みはUCVolumeDatasetによって最適化されています。

次回のトレーニングワークロードで、ぜひAIRをお試しください！お試しの準備はできましたか？ご自身でこのモデルを実行する方法の詳細については、AI Runtimeクックブックをご覧ください。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事