TRL documentation

メモリ使用量の削減

TRL

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

メモリ使用量の削減

トレーニングワークフローは、メモリ消費量を削減するように最適化できることが多く、TRLにはこれを達成するためのいくつかの組み込み機能が提供されています。

以下では、これらの手法を概説し、特定のセットアップに最適な構成を特定するために、異なる組み合わせの実験をお勧めします。

各メソッドには、サポートされているトレーナーの例が含まれています。特定のテクニックがご使用のトレーナーと互換性があるか不明な場合は、該当するトレーナーのドキュメントをご参照ください。

勾配チェックポイントなど、すべてのトレーナーでサポートされている追加の戦略については、transformersのパフォーマンスガイドを参照してください。

切り詰め

データセット内のシーケンス長は大きくばらつく可能性があります。データがバッチ化されると、バッチ内でもっとも長いシーケンスに合わせてパディングが行われるため、ほとんどのシーケンスが比較的短い場合でも、メモリ使用量が高くなる可能性があります。

Truncation prompt-completion

メモリ使用量を削減するには、シーケンスを適切な長さに切り詰めることが重要です。TRLのトレーナーはデフォルトでシーケンスを切り詰めますが、特定のユースケースに合わせるために、デフォルトの切り詰め長を調整する必要がある場合があります。

DPO

SFT

max_lengthの値はどのように選択すればよいですか？

max_length が小さすぎると、トークンの大部分が破棄され、学習に貢献しなくなります。大きすぎるとメモリ使用量が急増し、メモリ不足（OOM）エラーが発生する可能性があります。パッキングや padding-free を使用しない場合、多くのトークンが padding になるため、大きな max_length は非効率な学習につながることもあります。

適切な値を選択できるよう支援するため、データセット内のシーケンス長の分布を可視化するユーティリティを提供しています。

パッキング

この技術は、SFTトレーニングおよびFlashAttention（またはそのバリアント）を使用するセットアップでのみ利用可能です。

切り詰めにはいくつかの欠点があります：

情報の損失: シーケンス末尾の重要なトークンが破棄される可能性があります。
切り詰め長の選択：短すぎるとデータが失われ、長すぎると効率が低下します。

Packing は、複数のシーケンスを同じトレーニング行にグループ化し、各行を max_length まで埋めることで、これらの問題を軽減します。

Packing

TRLは、シーケンスを効率的にグループ化しながらパディングを最小限に抑えるBest-Fit Decreasing (BFD)によるビニングパッキングを実装しています。シーケンスがmax_lengthを超えた場合、オーバーフローしたトークンをどのように処理するかは、異なる戦略によって決定されます。

TRLは3つの戦略をサポートしています：

"bfd" (デフォルト): Best-Fit Decreasing packingを使用します。max_lengthを超えるシーケンスの場合、オーバーフローしたトークンは破棄されます。
"bfd_split": Best-Fit Decreasing packingを使用しますが、パッキングの前に長いシーケンスをmax_length以下のチャンクに分割します。これによりすべてのトークンが保持され、Fewer Truncations Improve Language Modelingで提案されたアプローチに従います。
"wrapped": すべてのトークンを1つのストリームに連結し、固定長のブロックに分割します。これによりパディングを最小限に抑えられますが、無関係な例が混在する可能性があります。この戦略は、文献で説明されているconcatenate-then-split前処理（例: Fewer Truncations Improve Language Modeling）に対応します。ただし、Qwen3-Coder-Next Technical Reportで議論されているように、データセットの大部分でシーケンスの連続性が断たれるという欠点があり、パフォーマンスに悪影響を及ぼします。

すべてのシーケンスがmax_lengthより短い場合、切り詰めや分割が不要なため、bfdとbfd_splitは同じように動作します。

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    packing=True,
    packing_strategy="bfd",
    max_length=512,
)

パラメータ効率の良いファインチューニングのためのPEFT

LoRAのようなParameter-Efficient Fine-Tuning (PEFT)手法は、学習時のメモリ使用量を削減する最も効果的な技術の一つです。モデルのすべてのパラメータを学習するのではなく、PEFT手法では少数のアダプタパラメータのみを学習するため、メモリ要件を大幅に削減し、限られたハードウェア上でより大きなモデルのファインチューニングを可能にします。

PEFTをTRLと組み合わせて使用する方法の詳細、各種アダプター手法、量子化オプション、高度な設定については、PEFT Integrationを参照してください。

メモリ使用量を削減するためにPEFTを使用するには:

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

peft_config = LoraConfig()

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
)

PEFTは、量子化（4ビットまたは8ビット）などの他のメモリ削減技術と組み合わせることで、さらに大きなメモリ削減を実現できます。量子化の例については、PEFT Integrationを参照してください。

パディングなし

Padding-free batchingは、メモリ使用量を削減するための代替アプローチです。この手法では、まずバッチをサンプリングし、その後、パディングを回避するために単一のシーケンスにフラット化します。異なるサンプルの一部を組み合わせることで不完全なシーケンスが生じる可能性があるpackingとは異なり、padding-free batchingでは、すべてのシーケンスが完全なまま保持されます。

Padding-free

DPO

SFT

アクティベーションのオフロード

Activation offloadingは、GPUのVRAM使用量を削減するためのメモリ効率化技術です。この技術では、フォワードパス中にアクティベーションテンソルを一時的にCPU RAMへ移動し、バックワードパスで必要になったタイミングでのみGPUへ戻します。これにより、わずかな訓練時間の増加を代償に、ピーク時のメモリ使用量を大幅に削減できます。

SFTトレーニング設定でアクティベーションのオフロードを有効にするには:

from trl import SFTConfig

training_args = SFTConfig(..., activation_offloading=True)

内部実装として、activation offloading は PyTorch の saved_tensors_hooks を実装し、フォワードパス中のアクティベーションをインターセプトします。サイズとコンテキストに基づいてオフロードするテンソルをインテリジェントに管理し、非効率になる出力テンソルのオフロードを回避します。パフォーマンス最適化のため、フラグ（デフォルトでは true）を介して CUDA ストリームを使用し、計算と CPU-GPU 間の転送をオーバーラップさせることができます。

シーケンスを倍数にパディング

この技術は、現在、SFTとRewardトレーナーでサポートされています。

有効にすると、このオプションはすべてのシーケンスが指定された値の倍数にパディングされることを保証します。
シーケンスの長さをメモリに優しい境界に整列させることで、一部のハードウェアでは計算効率を向上させることができます。

SFT

報酬

PyTorchのキャッシュアロケータ

長時間の実行、特にオンラインRL（GRPO、RLOO、Online DPO）では、GPUメモリが断片化することがあります。expandable_segments:Trueを設定すると、PyTorchのキャッシュアロケーターが既存のセグメントを拡張できるようになり、割り当てメモリと予約メモリの差が縮小されます。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_ALLOC_CONF=expandable_segments:True  # canonical name since PyTorch 2.10

vLLM のスリープモード（vllm_enable_sleep_mode=True）では、vLLM はこの設定を回避するために PYTORCH_CUDA_ALLOC_CONF のみを読み取ります。したがって、PYTORCH_ALLOC_CONF だけを設定しないでください。

オンライン手法における生成時のモデル集約の無効化

DeepSpeed ZeRO-3を使用すると、モデルの重みが複数のGPUにシャードされます。オンライン手法では、トレーニングプロセスの一部としてモデルから補完を生成します。このステップでは、生成のためにモデルの重みが一時的に単一のGPUに集約されます。非常に大規模なモデルの場合、この集約によりOOMエラーが発生する可能性があります。詳細は、このissue #2250に記載されています。

この問題が発生した場合、以下のパラメータを設定することで、生成時のモデル重みの集約を無効にできます。

GRPO

オンラインDPO

RLOO

この調整により、モデルの重みが集約されることを防ぎ、OOMエラーを回避できますが、生成速度が遅くなる可能性があります。

vLLM スリープモード

オンライン学習手法の生成バックエンドとしてvLLMを使用する場合、sleep modeを有効にすることで、最適化ステップ中にvLLMのパラメータとキャッシュをCPU RAMにオフロードし、重みの同期や生成が必要なときにGPU VRAMへ再読み込みできます。

GRPO

RLOO

vLLMの重みとキャッシュをオフロードすることで、GPUメモリ使用量を低く保つことができます。これは、大規模モデルの学習や、GPUリソースが限られている環境で特に有益です。ただし、vLLMエンジンをスリープモードから復帰させる際に、ホストとデバイス間の転送遅延が発生するため、学習速度がわずかに低下する可能性があります。

Gradient checkpointing

Gradient checkpointingは、フォワードパス中にすべての中間アクティベーションを保存せず、代わりにバックワードパス中にそれらを再計算することで、計算量をメモリと引き換えにします。

from trl import SFTConfig

training_args = SFTConfig(..., gradient_checkpointing=True)

メモリ使用量の最適化のため、すべてのトレーナーでグラデーションチェックポイントがデフォルトで有効になっています。必要に応じてgradient_checkpointing=Falseを設定することで無効化できます。

選択的アクティベーション・チェックポイント

SFTTrainerを使用すると、バックワードパスでの再計算の代わりに、フォワードパス中のアテンション出力を保存できます。これにより、レイヤーごとに1つの隠れ状態サイズのテンソルを追加するだけで、長いコンテキストにおけるチェックポイント化による速度低下の大部分を回復できます。この機能は、非再帰的チェックポイント化を強制します。

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"selective": True},
)

より多くのメモリ最適化技術については、Transformers Performance Guideを参照してください。

Update on GitHub

←トレーニングのカスタマイズトレーニングの高速化→