研究チームは、4カ国に分散した14台のMac(個人のMacBookを含む)を用いて強化学習のポストトレーニングを実施しました。ロールアウト生成はこれらMac上で行われ、トレーニングは遠隔地にある1台のNVIDIA B200 GPUで実行されました。この構成により、PaperSearchQAタスクにおけるheld-out pass@1が29%から63%に向上しました。これは、コンシューマ向けMacのみをロールアウト生成に利用した初めてのRLポストトレーニングです。
システムはスター型トポロジーを採用し、全てのノードがCloudflare R2のオブジェクトストレージを介して通信します。Macは生成したロールアウトをストレージにアップロードし、トレーナーはそれを取得して学習し、新しいウェイトをストレージに公開します。全プロセスは非同期で行われ、ノードが互いに待つことはありません。ウェイト同期には独自のPULSE技術を採用し、int8量子化精度とbf16学習精度の差異を利用して、9GBのダウンロードを約90MBに圧縮、家庭用回線でも約7秒で完了します。
Macが選ばれた理由は、ポストトレーニングの計算の大部分が推論であり、Macの統合メモリが大規模モデルとKVキャッシュを効率的に保持できるためです。特にMoEアーキテクチャではメモリ要求が大きく計算要求が小さいため、Macの特性が活かされます。チームは14台の異なる世代のMac(M4 Pro等)を使用し、int8精度でMLXフレームワークを実行しました。トレーニング側はbf16精度でMegatronフレームワークを使用しています。
非同期によるオフポリシーギャップに対処するため、カスタムのMetal paged-attentionカーネルとDPPOのトークンゲートを実装し、学習の安定性を確保しました。その結果、pass@1は29%から63%に倍増しました。全スタックはオープンソースとして公開されており、再現可能です。