⚡Unsloth カーネル＋パッキングで 3x 高速な LLM トレーニング

Unsloth がトレーニングスループットを向上させ、ファインチューニングのパディング無駄を排除する方法を学びます。

Unslothは現在最大で 5倍高速な （通常は3倍）学習を、新しいカスタム RoPEおよびMLPのTritonカーネルで、および新しいスマートな自動パッキングをサポートします。Unslothの新しいカーネルと機能は学習速度を向上させるだけでなく、さらに VRAM使用量を削減します（30%～90%） 精度の損失はありません。 Unsloth GitHub これは、例えば次のようなLLMをトレーニングできることを意味します Qwen3-4B を単に 3GBのVRAMだけでなく、3倍速で。

当社の自動 パディング不要の 汚染のないパッキングは、変更なしで全てのトレーニング実行に賢く有効化され、すべての高速アテンションバックエンド（FlashAttention 3、xFormers、SDPA）で動作します。 ベンチマーク トレーニング損失が非パッキング実行と一致することを示します 完全に.

• 2.3倍高速なQKロータリー埋め込み パッキングサポートを備えた融合Tritonカーネル

• SwiGLU、GeGLUのカーネルを更新し、 長いコンテキスト向けにint64インデックスを導入

• xformers、SDPA、FA3バックエンドで 2.5倍～5倍高速な汚染のないパッキング

• パディング不要で2.1倍高速、VRAMを50%削減、精度変化0%

• Unslothはまた、SFTの損失安定性が向上し、GPU利用率がより予測可能になりました。

• この新しいアップグレードは すべてのトレーニング手法で動作します 例：フルファインチューニング、事前学習など

🥁パッキング対応の融合QK RoPE Tritonカーネル

2023年12月に、Unslothのリリースの一環としてTritonで実装したRoPEカーネルを紹介しました。2024年3月に、コミュニティメンバーがRoPEカーネルを最適化してヘッドのグループごとにブロックを起動できるようにすることで、エンドツーエンドの学習を1～2%高速化しました。詳細は PR 238.

1つの問題は、各QとKに対して2つのTritonカーネルが存在することでした。これらを現在1つのTritonカーネルに統合し、 パディングフリーおよびパッキングサポートに不可欠だった可変長RoPEを有効にしました。これによりマイクロベンチマークにおけるRoPEカーネルは 長いコンテキスト長で2.3倍高速に、短いコンテキスト長では1.9倍高速になりました。

また、すべてのクローンと連続した転置操作を排除したため、 RoPEは現在完全にインプレースになりさらにGPUメモリを削減します。バックワードパスについては、次の点が確認されています： sin1 = -sin1 理由は：

Q * cos + rotate_half(Q) * sin
は以下と同等です
Q * cos + Q @ R * sin
ここでRは回転行列 [ 0,  I] です
                             [-I,  0]
dC/dY = dY * cos + dY @ R.T * sin
ここでR.Tは再び同じ [ 0, -I] ですが、
マイナスが転置されます。 [ I,  0]

🚃Tritonカーネルのint64インデックス化

で導入した500Kの長いコンテキスト学習中に、CUDAの範囲外エラーが発生することがありました。これはSwiGLU、GeGLUのMLPカーネルがTritonとCUDAのデフォルトであるint32インデックスを使用していたためです。 500K Context Training我々は単に

tl.program_id(0).to(tl.int64) とすることはできません なぜならint64インデックスによって学習がやや遅くなるためです。代わりにこれを LONG_INDEXING: tl.constexpr 変数にしてTritonコンパイラが特殊化できるようにしました。これにより短いコンテキストと長いコンテキストの両方でうまく動作します！

block_idx = tl.program_id(0)
if LONG_INDEXING:
    offsets = block_idx.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE).to(tl.int64)
    n_elements = tl.cast(n_elements, tl.int64)
scores.append(2.0)   # 動作したが2048には到達しなかった
    offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

🧮なぜパディングが必要か＆数学的な速度向上

コンピュータやGPUは異なる長さのデータセットを同時に処理できないため、0でパディングする必要があります。これが無駄を生みます。例えばデータセットに短いシーケンスSが50%、長いシーケンスLが50%あるとすると、最悪の場合パディングによるトークン使用量は $\text{batchsize} \times L$ 最長のシーケンス長が支配するためです。

複数の例を1つの長い一次元テンソルにパッキングすることで、かなりの量のパディングを排除できます。実際、以下のようなトークン使用量が得られます：

$$\text{Token Usage} = \frac{\text{batchsize}}{2}L+\frac{\text{batchsize}}{2}S$$

いくつかの数学と代数により、次式で速度向上を計算できます：

$$\text{Speedup} = \frac{\text{batchsize} \times L}{\frac{\text{batchsize}}{2}L+\frac{\text{batchsize}}{2}S} = 2 \frac{L}{L + S}$$

を仮定すると $S\rightarrow0$ 次に2倍の理論上の速度向上が得られます、なぜなら $2 \frac{L}{L + 0} = 2$

短いシーケンスの割合を50%から変え、短いシーケンスがより多いと仮定すると、例えば長いシーケンスが20%で短いシーケンスが80%なら、次のようになります $\frac{L}{0.2L + 0.8S}\rightarrow\frac{L}{0.2L}=5$ したがって学習は5倍速くなります！これはパッキングの速度向上がデータセット中の短い行の割合に依存することを意味します（短いほど速くなります）。

🎬デフォルトでパディング不要

packing = True を設定したときに得られる大きなスループット向上に加えて、 packing = True あなたの SFTConfig もしあなたが で自動的にパディング不要のバッチ処理を使用します パディングの無駄を減らしてスループットを改善し、トークン/秒のスループットを向上させます。結果として まったく同じ損失 がUnslothの以前のバージョンと同様に得られます。

例えばQwen3-8BやQwen3-32Bでは、メモリ使用量が60%減少し、2倍速くなり、同一の損失と勾配ノルム曲線が得られることを確認しています！

♠️汚染のないパッキングで2～5倍高速な学習

実際のデータセットは異なるシーケンス長を含む可能性があるため、例えばバッチサイズを32に増やすとパディングが発生し、学習が遅くなりVRAMを多く消費します。

過去には、 batch_size を大きな値（>32）にすると学習が速くなるどころか遅くなりました。これはパディングが原因でした—この問題を我々は packing = Trueで排除できるため、学習は速くなります！

複数のサンプルを1つの一次元テンソルにパックするとき、サンプル間でアテンションが漏れないように正しくマスクするためにシーケンス長のメタデータを保持します。また、位置IDをリセットするために New 3x Faster Training で説明したRoPEカーネルが必要です。

パッキングなしの4つの例はスペースを無駄にする

汚染のないパッキングは正しいアテンションパターンを生成する

短いシーケンスが50%の比率を変え、短いシーケンスがより多いと仮定すると、例えば長いシーケンスが20%で短いシーケンスが80%の場合、次のようになります $\frac{L}{0.2L + 0.8S}\rightarrow\frac{L}{0.2L}=5$ したがって学習は5倍速くなります！これはパッキングの速度向上がデータセット中の短い行の割合に依存することを意味します（短いほど速くなります）。

🏖️分析とベンチマーク

新しいカーネルとパックされたデータで学習したときの様々な改善を示すために、我々はファインチューニング実験を実行しました： すべての Qwen 3 デンスモデル（、Qwen3-8B、Llama 3 8B を使って yahma/alpaca-cleaned データセットでさまざまな トレーニング損失 スループットと効率の指標を測定しました。我々の新しい実行を、当社の最適化済みカーネル/最適化をオンにした標準的な最適化トレーニング実行（Flash Attention 3 (FA3)などのカーネル有効）と比較しました。我々は max_length = 1024 を固定し、バッチサイズを {1, 2, 4, 8, 16, 32} で変化させました。これによりバッチあたりの最大トークン数が {1024, 2048, 4096, 8192, 16K, 32K} の範囲で変化します。

上は、バッチサイズの変化に応じて新しいUnslothのトークン毎秒（tokens/s）学習スループットがどのように変化するかを示しています。これはデータセットのエポックを使ったモデル学習が 1.7～3倍速く（時には5倍以上）行えることを意味します！これらの利得はデータに短いシーケンスが多い場合や学習ランが長い場合により顕著になります。詳細は New 3x Faster Training

上は、バッチあたりのトークンのうち有効（すなわちパディングでない）トークンの平均割合を示しています。バッチサイズが大きくなると、アンパックされたケースではより多くのパディングトークンが発生しますが、パックされたケースでは最大シーケンス長に関係なく高いパッキング効率を達成しています。

なお、バッチロジックはバッチ内で観測された最大シーケンス長に合わせてトリムするため、バッチサイズが1の場合はアンパックデータはすべて有効トークン（すなわちパディングなし）です。しかし、バッチに例が増えると平均的にパディングが増え、バッチサイズ8ではほぼ50%のパディングに達します！我々のサンプルパッキング実装はその無駄を排除します。

最初のグラフ（上）は、 yahma/alpaca-cleaned を max_length = 2048での進捗をプロットしています。Unslothの新しい（パッキング+カーネル、マルーン）とUnslothの旧バージョン（灰色）を比較しています。どちらも max_steps = 500で訓練されていますが、x軸は実時間でプロットしています。パックされたケースでは、同じステップ数（かつ少し多めの実時間）でアンパックされたケースがエポックの5%未満を訓練する間に、ほぼ40%のエポックを訓練している点に注目してください。

同様に2番目のグラフ（上）は同じ実行からの損失をプロットしており、今回はx軸がトレーニングステップです。損失の尺度と傾向は一致しますが、パッキングケースではモデルが1ステップあたりより多くのトークンを見ているため損失の変動が少ないことに注意してください。

✨パッキングを有効にするには？

まずUnslothをアップデートすると、パディング不要はデフォルトで有効になります！したがってすべてのトレーニングは直ちに1.1～2倍速くなり、少なくとも30%少ないメモリ使用で、損失曲線指標に変化はありません！

pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth
pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth_zoo

我々はまた、Xformers経由のFlash Attention 3、SDPAサポート、Flash Attention 2をサポートしており、これは古いGPU（Tesla T4、RTX 2080）やH100、B200などの新しいGPUでも動作します！サンプルパッキングは アテンションバックエンドやモデルファミリーの選択に関係なく動作するので、これらの高速アテンション実装で以前得られたのと同じ速度向上をお楽しみください！

明示的なパッキングを有効にしたい場合は、次を追加してください packing = True で最大5倍速の学習を有効にできます！

注意 packing=True はトレーニング損失を変え、複数の短いシーケンスが1つにパックされるためデータセットの行数が切り詰められることになります。データセットの例数が減るのを確認するかもしれません。

異なるトレーニング損失値を得たくない場合は、単に packing=False に設定してください。そうすれば自動のパディングフリーが有効になり、既に学習を高速化します！

from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig

model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/Qwen3-14B",
)

trainer = SFTTrainer(
    model = model,
    processing_class = tokenizer,
    train_dataset = dataset,
    args = SFTConfig(
        per_device_train_batch_size = 1,
        max_length = 4096,
        …,
        packing = True, # サンプルパッキングを有効にするために必須！
    ),
)
trainer.train()

当社のすべてのノートブックは自動的に高速化されます（何かをする必要はありません）。詳細は Unsloth ノートブック

Qwen3 14B の高速化：

Google Colabcolab.research.google.com

Llama 3.1 Conversational の高速化：

Google Colabcolab.research.google.com

ありがとうございます！興味があれば、我々の 500K Context Training ブログ、 メモリ効率の良い RL ブログおよび Long Context gpt-oss ブログでカーネルや性能向上に関するその他のトピックをご覧ください！