50万トークンのコンテキスト長ファインチューニング

Unsloth で >500K トークンのコンテキストウィンドウファインチューニングを有効にする方法を学びます。

Unslothでは、新しいアルゴリズムを導入しており、長文コンテキストのトレーニング限界を押し上げます。 あらゆるLLMおよびVLMのトレーニングでは、gpt-oss-20bのようなモデルが今や 500K以上のコンテキスト長 を単一の80GB H100 GPUで達成でき、以前の80Kと比べても精度の低下はありません。

到達可能なのは >750Kのコンテキストウィンドウ B200 192GB GPU上で。

我々の 80GB A100 Colabノートブック.

で500Kコンテキストのgpt-oss-20bファインチューニングを試してください。

• Unslothはメモリ使用パターン、速度、コンテキスト長の扱いを大幅に改善しました： を VRAM使用量を60%低減 3.2倍長いコンテキスト Unslothの新しい 融合およびチャンク化されたクロスエントロピー

• 損失により、速度や精度の劣化なしに Unslothの

• Gradient Checkpointing における強化されたアクティベーションオフローディングと共に、

Stas Bekman（Snowflake）とのコラボレーションでTiled MLP

によりコンテキストを2倍多く可能にします。

Unslothのアルゴリズムにより、gpt-oss-20b QLoRA（4bit）はH100上で290Kコンテキストが精度損失なしに可能になり、Tiled MLPを有効にすると500K以上に到達し、合計で > 6.4倍の長いコンテキスト長を実現します。📐 Unsloth損失の再構成：チャンク＆フューズ 新しい融合損失実装は 動的シーケンスチャンク化 を追加します：言語モデルヘッドのロジットやクロスエントロピーをシーケンス全体で一度に計算する代わりに、フラット化されたシーケンス次元に沿って扱いやすいスライスで処理します。これによりピークメモリが数GBから小さなチャンクサイズに切り詰められます。各チャンクは依然として

torch.func.grad_and_value を介した完全に融合された順伝播＋逆伝播を実行し、必要に応じてfloat32にアップキャストすることで混合精度の精度を保持します。 これらの変更はトレーニング速度や精度を劣化させません。

• 主要な革新は、

• チャンクサイズが実行時に自動的に選択される

ことです。 利用可能なVRAMに基づいてチャンクサイズが選ばれます。 より多くの空きVRAMがあれば、より大きなチャンクが使用されてより高速な実行が可能になります

より少ないVRAMの場合は、メモリ爆発を避けるためにチャンク数を増やします。 これにより 手動チューニングが不要になり 古いGPUや新しいGPU、ワークロードや異なるシーケンス長にわたってアルゴリズムの堅牢性を保ちます。自動チューニングにより、

より小さなコンテキストはより多くのVRAMを使用します

（チャンクが少ない） 不要なオーバーヘッドを避けるためです 。上記の図では、現実的なVRAM階層を反映するために損失チャンク数を調整しています。80GBのVRAMでは、これは>3.2×長いコンテキストをもたらします。 🏁 UnslothのGradient Checkpointing拡張我々の 0.1% Unsloth Gradient Checkpointing

アルゴリズムは、
2024年4月に導入され
    すぐに人気となり、現在では業界標準となり、ほとんどのトレーニングパッケージに統合されています。これはアクティベーションをCPU RAMにオフロードし、コンテキスト長を10倍にすることを可能にしました。我々の新しい拡張はCUDAストリームやその他の工夫を用い、精度に影響を与えずに最大で
    トレーニングのオーバーヘッドを追加します。以前は1〜3%のトレーニングオーバーヘッドを追加していました。
    # オリジナルのUnslothバージョン（2024年4月公開） - LGPLv3 ライセンス
        class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
        @staticmethod
        @torch_amp_custom_fwd
            def forward(ctx, forward_function, hidden_states, *args):
        ctx.device = hidden_states.device
        saved_hidden_states = hidden_states.to("cpu", non_blocking = True)
        with torch.no_grad():

    すぐに人気となり、現在では業界標準となり、ほとんどのトレーニングパッケージに統合されています。これはアクティベーションをCPU RAMにオフロードし、コンテキスト長を10倍にすることを可能にしました。我々の新しい拡張はCUDAストリームやその他の工夫を用い、精度に影響を与えずに最大で
    output = forward_function(hidden_states, *args)
    ctx.save_for_backward(saved_hidden_states)
        ctx.forward_function, ctx.args = forward_function, args
        return output
        @torch_amp_custom_bwd
        def backward(ctx, dY):
            (hidden_states,) = ctx.saved_tensors
        hidden_states = hidden_states.to(ctx.device, non_blocking = True).detach()
        hidden_states.requires_grad_(True)

with torch.enable_grad():

(output,) = ctx.forward_function(hidden_states, *ctx.args)

torch.autograd.backward(output, dY)

return (None, hidden_states.grad,) + (None,)*len(ctx.args) アクティベーションを生成され次第すぐにオフロードすることで、ピークアクティベーションフットプリントを最小化し、GPUメモリを必要なときに正確に解放します。これにより、単一のデコーダ層のアクティベーションが2GBを超えるような長文コンテキストや大きなバッチのトレーニングにおけるメモリ圧力が大幅に軽減されます。 したがって、Unslothの新しいアルゴリズムとGradient Checkpointingはほとんどの改善（3.2倍）に寄与しており、単一のH100で290kコンテキスト長のQLoRA GPT-OSSファインチューニングを可能にします。 🔓 Tiled MLP：500K+の解放 Stas Bekmanの協力により、

我々はSnowflakeのArctic Long Sequence Trainingの

論文

とブログ記事からTiled MLPを統合しました。 TiledMLPはアクティベーションメモリを削減し、重いMLP投影の前にシーケンス次元に沿って隠れ状態をタイル化することではるかに長いシーケンス長を可能にします。また、いくつかの利便性向上も導入しています： タイル化された順伝播の再計算でRNG状態を保持するため、ドロップアウトやその他の確率的演算が順伝播と逆伝播の再生間で一貫します。これによりネストされたチェックポイント計算が安定し、数値的に同一に保たれます。

実装は名前や型が

mlp とされるモジュールを自動でパッチするので、MLPモジュールを持つほとんどのモデルがTiled MLPでそのままサポートされます。

留意すべきトレードオフ

TiledMLPは追加の順伝播を行う代わりにVRAMを節約します。これはチェックポイント化されたトランスフォーマーブロック内にあり、それ自体がチェックポイントスタイルで記述されているため、実質的にネストされたチェックポイントになります：1つの

MLPは今やステップごとに約3回の順伝播と1回の逆伝播を行います 。その代わりに、ほとんどの中間MLPアクティベーションをVRAMから除去しつつ、非常に長いシーケンスをサポートできます。

図は、長文コンテキストのトレーニングステップ中の単一デコーダ層の順伝播と逆伝播におけるアクティブメモリのタイムラインを、Tiled MLPなし（左）とあり（右）で比較しています。Tiled MLPがない場合、ピークVRAMはMLPの逆伝播中に発生します；Tiled MLPを使うとピークは融合損失計算に移動します。約40%低いVRAM使用量が観測され、融合損失が利用可能VRAMに基づいて自動的にチャンク化するため、より小さいGPUではTiled MLPのピークはさらに小さくなります。

model, tokenizer = FastLanguageModel.from_pretrained(
    ...,
    クロスエントロピー損失が新たなボトルネックでないことを示すため、チャンクサイズを動的に選ぶ代わりに固定し、チャンク数を倍にしました。これにより損失に関連するメモリスパイクが大幅に減少します。最大メモリはどちらの場合も逆伝播中に発生し、全体のタイミングは類似していますが、Tiled MLPは小さなオーバーヘッドを追加します：1つの大きなGEMMが多くの逐次的な行列積に分かれ、前述の追加の順伝播が発生します。
)

総じて、このトレードオフは価値があります：Tiled MLPなしでは長文コンテキストのトレーニングはおおよそ2×のメモリを必要とすることがありますが、 Tiled MLPを使うと単一GPUで同じコンテキスト長に対してステップ時間の増加は約1.3×にとどまります。 を UnslothでTiled MLPを有効にする： unsloth_tiled_mlp = True, 設定するだけで、unsloth_tiled_mlp = True

from_pretrained とするとTiled MLPが有効になります。Arctic論文と同じロジックに従い、我々は num_shards = ceil(seq_len/hidden_size)

を選びます。 Tiled MLPを使うと単一GPUで同じコンテキスト長に対してステップ時間の増加は約1.3×にとどまります。各タイルは、スループットとメモリ節約のバランスを取るためにモデルの隠れ次元と同じサイズのシーケンス長で動作します。

---

我々はまた、Tiled MLPが実質的に3回の順伝播と1回の逆伝播を行うことをStas Bekmanと議論し、通常のGradient Checkpointingは2回の順伝播と1回の逆伝播であるのに対して、

我々の 80GB A100 Colabノートブック.

DeepSpeed がDeepSpeed内のTiled MLP向けのドキュメント更新を提供しました。 次回ファインチューニングがメモリ不足になったら、ぜひこれをオンにしてみてください：