概要
GoogleのLLMトレーニングフレームワーク「MaxText」が、単一ホストTPU環境でのSFTおよびGRPO/GSPOを用いた強化学習に対応し、モデルの最適化プロセスを大幅に簡素化しました。
詳細内容
Googleは、JAXベースのオープンソースLLMトレーニングフレームワーク「MaxText」において、単一ホストTPU（v5p-8やv6e-8など）でのポストトレーニング機能を新たにサポートしたことを発表しました。主な更新点は、教師あり微調整（SFT）と強化学習（RL）の導入です。SFTでは、Hugging Faceのデータセットやチェックポイントとのネイティブな統合、およびTunixライブラリによる効率的な実行が可能になりました。RLについては、メモリ効率の高いGRPO（Group Relative Policy Optimization）と、数学やコーディングの推論性能を高めるGSPO（Group Sequence Policy Optimization）の2つのアルゴリズムをサポート。vLLMを用いた高スループットな推論ループをトレーニング中に実行できることが特徴です。これにより、開発者は小規模なリソースから始めて、将来的にマルチホスト構成へシームレスにスケールアップできる道筋が提供されます。