redai-infra/Relax

解決する課題

Relaxは、オムニモーダル大規模言語モデル（LLM）向けに設計された高性能な強化学習（RL）ポストトレーニングフレームワークです。テキスト、視覚、音声といった異なるモダリティにわたるRLトレーニングのスケールアップという課題に対処し、ロールアウト（推論）フェーズとトレーニングフェーズの逐次的な性質によって通常発生するGPUのアイドル時間を排除します。

仕組み

Relaxは、Ray Serve上に構築されたサービス指向アーキテクチャを採用しており、異なるロール（Actor、Rollout、Criticなど）が独立したサービスとしてデプロイされます。TransferQueueと呼ばれるシステムを使用してトレーニングと推論を分離し、それらを別々のGPUクラスター上で並列に実行できるようにします。この非同期実行は、以下の3つのモードでサポートされています：

• Colocate (Sync): ActorとRolloutがGPUを共有し、リソースをタイムシェアリングします。
• Fully Async: 各ロールが独立したクラスターで動作し、ストリーミングデータ交換とDistributed Checkpoint Service (DCS) を介した非同期の重み同期を行います。
• Hybrid: ActorとRolloutは個別の配置グループを使用しますが、リファレンスモデルなどの他のコンポーネントはActorのGPU上でインプロセスで動作します。

バックエンドには、トレーニング用にMegatron-LM（TP/PP/CP/EP並列をサポート）を、高スループット推論用にSGLangを統合しています。

対象ユーザー

ポストトレーニング・マルチモーダルモデルに取り組むAI研究者やエンジニア、特にQwen3シリーズやGLM5のようなモデルのRLトレーニングをスケールアップする必要がある方、および複雑なエージェント的RLワークフローを実装しようとしている方を対象としています。

ハイライト

• オムニモーダル対応: テキスト、視覚、音声のRLのための統合フレームワークであり、Qwen3-Omniのようなモデルのエンドツーエンドのトレーニングをサポートします。
• エージェント的RL: クローズドループトレーニングのためのマルチターン対話、ロス・マスキング、マルチモーダル・コンテキストの引き継ぎをサポートします。
• 弾力的なロールアウト・スケーリング: トレーニング中にHTTP REST APIを介して推論エンジンを動的に追加または削除できます。
• 豊富なアルゴリズム・スイート: GRPO、GSPO、SAPO、CISPO、およびOn-Policy Distillationを即座に利用可能です。
• プロダクション対応の運用機能: 自動復旧のためのHealthManagerと、WandB、TensorBoard、ClearML用の集中型Metrics Serviceが含まれています。