並列学習のトレーニング

CoMLRL は、異なるモデルを異なる GPUs に配置できます。 parallel_training="mp" を使用する場合、各エージェントは agent_devices からの個別のデバイスで実行され、エージェントは同時に生成されます。 critics、報酬モデル、比較モデル、および参照モデルは、それぞれ独自のデバイス設定を使用します。 デフォルトの "none" は、エージェントを 1 つのデバイスに保持します。 両方のモードは 1 台のマシンで実行でき、ジョブスケジューラは必要ありません。

各モデルは割り当てられたデバイス上で完全に動作し、GPUs に分割されません。 デバイスのインデックスは、CUDA_VISIBLE_DEVICES を通じて確認できる GPUs を参照します。 複数のモデルの設定は、1 つのデバイスを各モデルに割り当てるか、すべてのモデルに 1 つのデバイスを割り当てるため、ユーザーは利用可能な GPUs と各モデルのメモリに合わせて配置できます。 デバイスの配置と並列スケジューリングは、数値的なトレーニング結果に影響を与える可能性があります。

次の構成は、反復 MARLHF の各モデルを個別の GPU に配置します: 2 つのエージェント、集中型の critic (MAAC のポリシー学習用)、報酬モデル、および現在のエージェントのコピーを比較モデルとして使用します。

from comlrl.trainers.preference import MARLHFIterConfig

config = MARLHFIterConfig(
    num_agents=2,
    num_turns=1,
    rl_algorithm="maac",
    parallel_training="mp",
    agent_devices=["cuda:0", "cuda:1"],
    critic_devices="cuda:2",
    reward_model_device="cuda:3",
    comparator_policy="current_copy",
    comparator_devices=["cuda:4", "cuda:5"],
)

デバイスの配置#