CoMLRL 可以将不同的模型放置在不同的 GPUs 上。 使用 parallel_training="mp" 时，每个智能体在 agent_devices 上的设备上运行，并且智能体同时生成； critics、奖励模型、比较器和参考模型各自使用自己的设备设置。 默认设置 "none" 将智能体放在一个设备上。 两种模式都在一台机器上运行，不需要作业调度。
每个模型完全在分配的设备上运行，不会分散到 GPUs。 设备索引指的是通过 CUDA_VISIBLE_DEVICES 看到的 GPUs。 对于多个模型，设置可以是每个模型一个设备，也可以是所有模型一个设备，因此用户可以根据可用的 GPUs 和每个模型的内存进行匹配。 设备放置和并行调度会影响数值训练结果。
以下配置将迭代 MARLHF 中的每个模型放在各自的 GPU 上：两个智能体、一个用于 MAAC 策略学习的集中式 critic、奖励模型和当前智能体的副本作为比较器。
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"],
)设备放置#
Device Placement 参数
agent_devices: 智能体设备；在"mp"模式下需要critic_devices: 单个设备用于集中式 critic，或每个智能体分别使用 critics；在"mp"模式下，需要 MAAC 和单独的 critic IAC，而 MARLHF 默认使用智能体设备reward_model_device: MARLHF 奖励模型设备; 默认为"cuda:0"，或None用于 trainer 设备comparator_devices: 迭代 PbRL 中本地比较器模型的设备；默认使用智能体设备reference_devices: 使用参考 KL 的冻结参考模型的设备；默认使用智能体设备