Paralelización del entrenamiento

CoMLRL puede colocar diferentes modelos en diferentes GPUs. Con parallel_training="mp", cada agente se ejecuta en su propio dispositivo de agent_devices, y los agentes generan simultáneamente; critics, modelos de recompensa, comparadores y modelos de referencia tienen sus propias configuraciones de dispositivo. El modo predeterminado, "none", mantiene a los agentes en un solo dispositivo. Ambos modos se ejecutan en una sola máquina y no requieren un programador de trabajos.

Cada modelo se ejecuta completamente en su dispositivo asignado y no se divide entre GPUs. Los índices de dispositivo se refieren a los GPUs visibles a través de CUDA_VISIBLE_DEVICES. Una configuración para varios modelos utiliza un dispositivo por modelo o un solo dispositivo para todos, para que los usuarios puedan adaptar la colocación a los GPUs disponibles y a la memoria de cada modelo. La colocación de dispositivos y la programación en paralelo pueden afectar a los resultados del entrenamiento numérico.

La siguiente configuración coloca cada modelo de un ciclo MARLHF en su propio entorno GPU: dos agentes, una política centralizada critic para MAAC aprendizaje de políticas, el modelo de recompensa, y copias de los agentes actuales como comparadores.

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"],
)

Colocación de dispositivos#