Hugging Faceの言語モデル学習ツールがAmazonの学習専用チップに対応し、非同期強化学習で繰り返される入力の重複計算を削減した。学習環境を広げたい開発者と、同じ質問から複数の回答を生成して学習するチームが、それぞれ注目すべき変更だ。

1. Trainiumで既存の学習コードを使う

2026年10月15日に公開されたTRL v1.16.0は、AWS Trainiumへの対応と利用ガイドを追加した。リリース説明によると、前バージョンの統合言語モデル出力層を基盤に、残っていた対応の不足を補った。

利用ガイドでは、torch-neuronxをインストールした後、同じ環境にTRLをインストールすれば、既存のトレーナーのコードを変更せずに実行できると説明している。以下は、教師ありファインチューニング、つまり回答例に沿ってモデルを学習させるドキュメントのクイックスタートコードだ。

from datasets import load_dataset
from trl import SFTTrainer

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

2. 長さをそろえて再コンパイルを減らす

Neuronは入力形状が変わるたびに計算グラフをコンパイルし、1回に約1分かかる。今回のバージョンはpad_to_multiple_ofのデフォルト値を自動選択に変更した。Neuronでは入力長を128の倍数にパディングして形状の種類を減らし、GPUでは従来の動作を維持する。

Trainium2でQwen3-0.6BをNeuronコア1個、バッチサイズ4で教師ありファインチューニングした結果、20ステップにデフォルト設定では130秒、長さの調整を無効にすると260秒かかった。これは当該例の比較であり、他のモデルの学習速度まで保証する数値ではない。

モデルのコンパイルは任意だ。torch_compile=Trueを使うには、Dynamoの自動動的形状機能を無効にする必要がある。ドキュメントの同じモデル・バッチサイズ・コア数では、初回コンパイルは約2分で、その後のステップ当たりの時間は、コンパイルしない場合の1.71秒から1.15秒に短縮された。

3. 同じ先頭部分は1回だけ計算する

AsyncGRPOのツリーアテンションは、小さな学習バッチの入力を接頭辞木にまとめる。同じ質問から分岐した回答を木の枝のようにつなぎ、共有する先頭部分を回答ごとに繰り返さず、1回だけ計算する方式だ。

内部では分岐のない部分を区間にまとめ、各区間は自分につながる経路のキーと値を集めて参照する。アテンション以外の層は固有のトークンごとに1回実行し、共有経路ではキーと値だけが重複する。

リリースのQwen3-1.7Bの測定は、生成結果当たり32Kトークンの順伝播・逆伝播の計算を合計したものだ。共有プロンプト、パッキング比率3.37の条件で、従来のシーケンスパッキングは8179ms、ツリーパッキングは2755msで、提示された速度倍率は2.97倍だ。最大メモリは35.2GiBから13.6GiBに減った。測定ハードウェアは示されていない。

個別のパッキング選択オプションはない。共有部分がなければ従来のシーケンスパッキングと同じ形になり、リリースはこの場合の追加コストを約1%と説明している。

4. オンライン生成ではキャッシュとデータ型を確認する

この環境では、学習中に回答を生成するオンライントレーナーはvLLMを使えないため、Transformersで生成する。デフォルトの動的キャッシュはトークンごとに入力形状を変えるため、ドキュメントは生成中の形状を固定する静的キャッシュの使用を案内している。

現在、静的キャッシュにはbf16モデルが必要だ。デフォルトのfp32モデルでは、キーと値のデータ型が異なる問題が報告された。生成前にプロンプト長もそろえるが、新しいキャッシュサイズが現れると再コンパイルが必要になる。

データ並列実行でもコア配置を確認する必要がある。Trainium2チップを1個搭載するtrn2.3xlargeのデフォルト設定は、物理コア8個を論理コア4個にまとめる。ドキュメントは、この構成でAccelerateを使い4個のプロセスを実行した教師ありファインチューニングが、コア1個より3.6倍速いと説明している。一部のコアだけを選ぶ際は、ドキュメントで許可された組み合わせに従う必要がある。

5. 削減される計算と準備コストを分けて考える

2つの変更は異なるコストを削減する。ツリー方式は共有入力の重複演算とメモリを減らし、Trainiumの長さ調整は余分なパディングトークンを処理する代わりに、再コンパイルする形状の種類を減らす。モデルのコンパイルを選ぶと、開始時の待ち時間を受け入れて、その後の反復実行時間を短縮する仕組みだ。

ツリーアテンションの性能測定とTrainiumの利用例は別々に示された。2つの機能を組み合わせた性能数値は示されていないため、それぞれの改善幅を足し合わせて解釈してはならない。

今回の更新は学習ハードウェアの選択肢と共有入力の処理効率を広げるが、導入判断では入力の共有度とコンパイル準備コストを個別に検討する必要がある。