TRL 1.13、チャンク損失をTensor Coreに戻し、長コンテキストのポストトレーニングでスループットを最大1.69倍に向上

Hugging Faceは9月10日、[TRL 1.13.0](https://github.com/huggingface/trl/releases)をリリースした。中核となる改善は新しいアルゴリズムではなく、デフォルトの`chunked_nll`損失パスに存在していた高コストなデータ型変換の修正だ。旧実装では、各チャンクで、すでにBF16であるhidden statesと`lm_head`の重みをFP32に変換していた。その結果、行列乗算がTensor CoreではなくFP32 SIMTで実行され、巨大な出力層のコピーも繰り返し作成されていた。語彙サイズ248,320、hidden size 2,048、256 tokenのチャンクを単一のH100でテストしたところ、forwardとbackwardを合わせた実行時間は23.37ミリ秒から3.86ミリ秒に短縮され、ピークメモリ使用量は5.99 GBから3.03 GBに減少した。

この約6倍という数値は、単一の中核処理だけを対象としたものだ。デプロイの観点でより重要なエンドツーエンドの結果では、GPU当たりのスループットがモデルと学習方式に応じて約1.20～1.69倍向上している。この修正は、各チャンクで生徒モデルと教師モデルの出力を計算する必要がある蒸留学習にも適用された。公式によると、Accelerateのmixed precision環境では新旧の結果がbitwise identicalになるが、それ以外の精度設定、GPU、カスタムautocastスコープについては、個別に回帰テストを行う必要がある。

新版では、実行可能な100万token学習のサンプルも提供されている。Qwen3-8Bを8基のH100上でBF16、`chunked_nll`、GPU当たりのbatch size 1という構成で動かし、1,048,576 tokenを処理する。1ステップ当たりの所要時間は約380秒で、各GPUのメモリ使用量は56.2 GBだった。これは単一ノードで完全な学習ステップを実行できることを示すものだが、一般的なモデルで100万tokenのコンテキストを容易に利用できることを意味するわけではない。このサンプルは引き続きfull attentionに限定され、packingには対応しておらず、gradient-checkpointing offloadとYaRN RoPEに依存している。

互換性の面では、`PPOTrainer`、`PPOConfig`、value-head wrapperがコードから削除された。また、`peft`の最低バージョンは0.13、DeepSpeedは0.18.6に引き上げられ、vLLM 0.19.0もサポート対象外となった。エンジニアリングチームは、性能向上と移行コストを併せて評価すべきだ。特に、旧PPOワークフローに依存しているプロジェクトや、カスタム損失関数を使用しているプロジェクトでは注意が必要になる。