テンセントは最近、Megatron-CoreとSGLang/vLLMを基盤として開発した大規模モデルトレーニングライブラリ「WeChat-YATT（Yet Another Transformer Trainer）」をリリースしました。このトレーニングライブラリの内部プロジェクトコードネームはgCoreです。このツールは強化学習およびマルチモーダルモデルのトレーニングに特化しており、開発者に対して拡張性が高く、簡潔で効率的かつ信頼性の高い大規模モデルトレーニングソリューションを提供することを目指しています。

WeChat-YATTはカスタマイズされた並列計算戦略を通じて、大規模なモデルや長系列の入力、大容量のデータセットなどの複雑なシナリオを効果的に処理することができ、WeChat内の複数の実際的な業務シナリオにおける重要な課題を成功裏に解決し、大規模モデルトレーニングの効率を顕著に向上させました。このツールは研究者や開発者に対して柔軟で拡張可能な技術的解決策を提供し、マルチモーダルおよび強化学習分野の革新と発展に貢献することが期待されています。

image.png

WeChat-YATTは大規模モデルの分散トレーニング中に直面する2つの主要なテクノロジー的な課題を重点的に解決しました。

まず、マルチモーダル環境での拡張性の制約問題です。画像や動画などのマルチモーダルデータの規模が増加するにつれて、従来のアーキテクチャでは単一のコントローラーによるデータ管理が通信やメモリのボトルネックになることが多く、システムのスループットが制限され、トレーニングプロセスが異常終了する場合もあります。WeChat-YATTは並列コントローラー（Parallel Controller）の並列管理メカニズムを取り入れることで、システムの負荷を効果的に分散させ、システムの拡張性と安定性を大幅に向上させ、マルチモーダルで大規模なデータ量を持つ複雑なシナリオに対応できるようになりました。

次に、動的サンプリングと生成型報酬計算における効率の欠点です。頻繁な動的サンプリングや生成型報酬計算が必要なトレーニングプロセスでは、モデルの頻繁な切り替えや「長尾」タスクにより多くの追加コストが生じ、GPUの計算能力が十分に活用されず、全体的なトレーニング効率に大きな影響を与えます。WeChat-YATTは部分共有戦略と非同期インタラクションメカニズムを通じて、モデルの切り替え損失や長尾タスクの影響を大幅に軽減し、トレーニングプロセスにおける高スループットと高リソース利用率を実現し、大規模なRLHFタスクの効率的な反復をより良くサポートします。

異なるビジネスシナリオのニーズに対応するために、WeChat-YATTは「全員共存」と「部分共存」の2種類のリソース配置モードをサポートし、クラスタリソースの利用効率を最大限に高めています。

全員共存モードでは、シリアルスケジューリングメカニズムが採用されており、Actor Rollouts、GenRM（生成型報酬モデル）とTrainが順番に実行されます。各役割がタスクを完了した後、計算リソースを解放し、次のタスクに必要なモデルを読み込みます。この戦略はほとんどの通常のトレーニングシナリオに適しています。注目すべきは、各段階において関連コンポーネントがすべてのGPUリソースを独占できることであり、これによりリソースの空き「バブル」時間は大幅に短縮され、全体的なトレーニングスループットと効率が顕著に向上します。

部分共存モードでは、Actor RolloutsとGenRMが独立して配置され、非同期方式で効率的にやり取りを行います。ActorトレーニングフェーズではすべてのGPUリソースを占有し、Rollouts生成フェーズではGPUリソースを解放し、Actor RolloutsとGenRMの2つのコンポーネントを協調して動作させます。システムは動的な負荷評価を通じてリソースの割当てとバランスを実施し、Rolloutsが生成された後、これらの2つはリソースを解放し、Actorが再びGPUにロードされて次のトレーニングフェーズに入ります。部分共存モードは、RolloutsとGenRMが頻繁にやり取りし、動的サンプリングを行う複雑なタスクシナリオに特に適しています。

WeChat-YATTにはいくつかの技術的な特徴があります。メモリ使用効率に関しては、プロジェクトは並列コントローラーアーキテクチャを使用しており、単一ノードのメモリ消費を効果的に低下させ、マルチモーダル環境での大規模モデルトレーニングに適しており、システムの拡張性と安定性を向上させています。GenRMサポートに関しては、生成型報酬モデルのシナリオに対応して異なるリソース配置戦略を実装しており、ユーザーは具体的なシナリオに応じて最適なトレーニング方法を選択できます。

知能チェックポイント戦略はもう一つの大きな特徴です。WeChat-YATTは非同期チェックポイント保存をサポートしており、WeChatのビジネスシナリオの特性に基づいてスケジューリングプロセスに合わせて自動的にチェックポイントを保存することで、トレーニングの安全性と高可用性をさらに確保しています。また、システムはトレーニング中に各データ並列グループ間での負荷バランスを実現し、リソースの空き時間を効果的に減少させ、全体的なトレーニングスループットを顕著に向上させています。

このトレーニングライブラリのリリースは、テンセントが大規模モデル技術インフラストラクチャの構築において重要な進展を遂げたことを示しており、業界に対して複雑なマルチモーダルトレーニングシナリオを処理するための有効なソリューションを提供しています。