Entidades de entrenamiento

CoMLRL soporta el entrenamiento de LLM y facilita su entrenamiento con modelos de recompensa basados en LLM o critics. Los agentes y critics pueden ser modelos homogéneos o heterogéneos provenientes de Hugging Face o directorios locales. Los usuarios pueden especificar identificadores de modelo Hugging Face a través de agent_model/critic_model para configuraciones homogéneas, o proporcionar listas agents/critics para configuraciones heterogéneas. Las entidades entrenadas pueden guardarse y cargarse en un nuevo trainer para un entrenamiento adicional.

Agentes LLM#

Agentes homogéneos#

Los agentes homogéneos se inicializan estableciendo agent_model a un único identificador de modelo y dejando agents=None. Para trainers con un argumento constructor num_agents separado, ese argumento puede establecerse junto con el valor de configuración correspondiente; la configuración por sí sola no lo anula.

El siguiente ejemplo carga tres Qwen/Qwen2.5-1.5B agentes:

trainer = MAGRPOTrainer(
    agent_model="Qwen/Qwen2.5-1.5B",
    agents=None,
    num_agents=3,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAGRPOConfig(num_agents=3, num_turns=1, temperature=0.7, top_p=0.9, top_k=None),
)

Agentes heterogéneos#

Aunque los agentes homogéneos pueden ser instruidos para asumir diferentes roles y trabajar en diferentes partes de una tarea en paralelo, los agentes heterogéneos pueden aportar las capacidades especializadas de diferentes modelos.

Los usuarios pueden inicializar agentes heterogéneos suministrando identificadores de modelo o objetos de modelo cargados en agents, con tokenizers suministrado a través de tokenizer. Con agents suministrado, agent_model puede permanecer None. Si ambos se proporcionan, sus nombres de modelo deben coincidir y sus longitudes de lista deben coincidir con el número de agentes; de lo contrario, el entrenamiento genera un error.

El siguiente ejemplo carga los agentes Qwen/Qwen2.5-Coder-3B y Qwen/Qwen2.5-Coder-7B

trainer = MAGRPOTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAGRPOConfig(num_agents=2, num_turns=1, temperature=0.7, top_p=0.9, top_k=None),
)

LLM Critics#

Los algoritmos como el Actor-Critic independiente (IAC) y el Actor-Critic multiagente (MAAC) utilizan LLM critics. La forma en que se cargan estos critics depende del algoritmo y la configuración use_separate_critic. MARLHF trainers con rl_algorithm="maac" o "iac" aceptan los mismos critic_model y critics argumentos, o critic_model_name en su configuración; de lo contrario, utilizan el identificador del modelo del primer actor.

Modelo centralizado Critic#

Se puede cargar un modelo causal separado y centralizado critic a través de critic_model o una lista de un elemento critics. Debe ser un modelo de lenguaje compatible: trainer añade value head y codifica la entrada conjunta con actor’s tokenizer, por lo que tokenizer también debe ser compatible con critic.

El siguiente ejemplo carga los pares Qwen/Qwen2.5-Coder-3B y Qwen/Qwen2.5-Coder-1.5B actors con un modelo centralizado Qwen/Qwen2.5-Coder-7B critic:

trainer = MAACTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-1.5B"],
    tokenizer=[
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-3B"),
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-1.5B"),
    ],
    critic_model="Qwen/Qwen2.5-Coder-7B",
    critics=None,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAACConfig(num_agents=2, num_turns=1, temperature=0.7, top_p=0.9, top_k=None),
)

Agente independiente Critics#

Con agentes independientes critics, cada agente puede tener su propio critic. Con use_separate_critic=True, los usuarios pueden proporcionar un identificador critic_model para cargar una copia separada para cada agente, o una lista critics de longitud num_agents. Cada copia separada critic recibe token IDs de su actor tokenizer y debe ser compatible con él. Con use_separate_critic=False, cada actor y critic comparten un LLM backbone y se cargan juntos. En este caso, critic_model y critics permanecen sin asignar (null o None). Al igual que con los agentes, proporcionar tanto critic_model como critics requiere nombres y longitudes de modelo coincidentes; de lo contrario, el entrenamiento genera un error.

El siguiente ejemplo carga Qwen/Qwen2.5-Coder-3B y Qwen/Qwen2.5-Coder-7B junto con actors utilizando critics basados en los mismos modelos:

trainer = IACTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    tokenizer=[
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-3B"),
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B"),
    ],
    critic_model=None,
    critics=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=IACConfig(
        num_agents=2,
        num_turns=1,
        use_separate_critic=True,
        temperature=0.7,
        top_p=0.9,
        top_k=None,
    ),
)

Alternativamente, cada actor puede compartir el mismo modelo con su critic:

trainer = IACTrainer(
    agent_model=None,
    agents=["Qwen/Qwen2.5-Coder-3B", "Qwen/Qwen2.5-Coder-7B"],
    tokenizer=[
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-3B"),
        AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B"),
    ],
    critic_model=None,
    critics=None,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=IACConfig(
        num_agents=2,
        num_turns=1,
        use_separate_critic=False,
        temperature=0.7,
        top_p=0.9,
        top_k=None,
    ),
)

Trainers utiliza agents/critics internamente. Los parámetros agent_model y critic_model son atajos para configuraciones homogéneas y deben ser consistentes con cualquier lista correspondiente suministrada.

MAGRPO carga tokenizers automáticamente para cadenas de identificadores de modelo. MAAC/IAC cargan automáticamente tokenizer para agent_model, pero requieren un tokenizer explícito cuando se suministra agents, incluso si contiene cadenas. Los objetos de modelo cargados requieren una lista compartida compatible tokenizer o tokenizer que coincida con actors. Los modelos heterogéneos actors pueden tener vocabularios diferentes cuando cada uno utiliza su propia lista compatible tokenizer; también, las separadas critics deben aceptar actor token IDs (el primer actor IDs para MAAC).

Modelos de recompensa LLM#

En PbRL, MARLHF aprende un modelo de recompensa conjunta a partir de pares de preferencias y lo mantiene fijo como recompensa para el aprendizaje de la política. Su backbone se carga a través de reward_model_name en MARLHFConfig, desde un identificador de modelo Hugging Face o un directorio local de un modelo de lenguaje causal. Por defecto, utiliza el identificador del primer actor y es requerido cuando los actors son objetos de modelo sin uno. El trainer carga el backbone con su propio tokenizer, añade una cabeza escalar en el último token sin relleno y coloca el modelo en reward_model_device (por defecto "cuda:0"; None utiliza el dispositivo trainer). reward_freeze_backbone=True solo entrena la cabeza.

El siguiente ejemplo carga un modelo de recompensa Qwen/Qwen2.5-Coder-7B para dos agentes Qwen/Qwen2.5-Coder-3B:

from comlrl.trainers.preference import MARLHFConfig, MARLHFTrainer

trainer = MARLHFTrainer(
    agent_model="Qwen/Qwen2.5-Coder-3B",
    num_agents=2,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MARLHFConfig(
        num_agents=2,
        num_turns=1,
        reward_model_name="Qwen/Qwen2.5-Coder-7B",
        reward_model_device="cuda:1",
    ),
)

MARLHFIterTrainer recarga el modelo de recompensa de esta fuente en cada iteración antes del aprendizaje de preferencias. La recompensa de la tarea de reward_func sigue etiquetando las parejas de preferencias; Inicialización del entrenamiento describe cómo pasar recompensas de tarea basadas en reglas o LLM.

Guardar y reanudar el entrenamiento#

Guardar modelos#

Los usuarios pueden llamar a save_model(output_dir) después del entrenamiento para guardar explícitamente la configuración del modelo actors y Hugging Face de tokenizers en directorios locales:

trainer.train()
trainer.save_model("./saved-models/run-1")

Entrenamiento a partir de modelos guardados#

El entrenamiento puede continuar a partir de una configuración de dos actor MAGRPO guardada, suministrando los directorios de actors en el orden original de los roles. El directorio padre run-1 contiene el actors; cada directorio agent_i es un modelo cargable.

from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer

continued_trainer = MAGRPOTrainer(
    agents=[
        "./saved-models/run-1/agent_0",
        "./saved-models/run-1/agent_1",
    ],
    num_agents=2,
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAGRPOConfig(
        num_agents=2,
        num_turns=1,
        num_train_epochs=3,
        agent_learning_rate=5e-6,
        temperature=0.7,
        top_p=0.9,
        top_k=None,
    ),
)
continued_trainer.train()
continued_trainer.save_model("./saved-models/run-2")

MAGRPO carga tokenizers de los directorios actor automáticamente. Aquí, num_train_epochs=3 significa tres epochs en la nueva ejecución de entrenamiento. El conjunto de datos, la función de recompensa, la configuración de entrenamiento y cualquier formateador o procesador de recompensas externo deben suministrarse de nuevo; save_model no los almacena.

Otros trainers soportan el mismo enfoque de directorio local, sujeto a los requisitos de tokenizer y actor mencionados anteriormente. La preferencia trainers inicia la recolección de preferencias y el aprendizaje de recompensas según la configuración de la nueva ejecución.

IAC y MAAC guardan value heads por separado como value_head.pt. Cargar un directorio critic o compartido actor a través del constructor restaura el modelo de lenguaje backbone y crea un nuevo value head. Los usuarios pueden conservar los valores aprendidos restaurando explícitamente la cabeza guardada antes de llamar a train(). Los puntos de control de entrenamiento, a diferencia de las exportaciones de save_model, restauran value heads automáticamente.

Para una configuración de dos actor MAAC guardada utilizando el lineal value head predeterminado:

from pathlib import Path
import torch
from transformers import AutoTokenizer
from comlrl.trainers.actor_critic import MAACConfig, MAACTrainer

saved_dir = Path("./saved-models/maac-run-1")
actor_dirs = [str(saved_dir / f"agent_{i}") for i in range(2)]

continued_maac = MAACTrainer(
    agents=actor_dirs,
    tokenizer=[AutoTokenizer.from_pretrained(path) for path in actor_dirs],
    critic_model=str(saved_dir / "critic"),
    reward_func=reward_func,
    train_dataset=train_dataset,
    args=MAACConfig(num_agents=2, num_turns=1, num_train_epochs=3),
)
continued_maac.critics[0].value_head.load_state_dict(
    torch.load(
        saved_dir / "critic" / "value_head.pt",
        map_location="cpu",
        weights_only=True,
    )
)
continued_maac.train()
continued_maac.save_model("./saved-models/maac-run-2")

Para IAC, el mismo paso load_state_dict aplica a continued_iac.agents[i].value_head cuando use_separate_critic=False, o a continued_iac.critics[i].value_head cuando es True. Los pasos critic backbones también requieren la carga a través de critic_model o critics.

La arquitectura value-head debe coincidir con la ejecución guardada. IAC requiere IACConfig.value_head_hidden_dim para las cabezas compartidas y IACConfig.critic_value_head_hidden_dim para las cabezas separadas. Una cabeza MAAC personalizada requiere model_config={"critic_value_head_hidden_dim": ...} en el constructor trainer. Continuar con la misma configuración de aprendizaje también requiere critic_type ("v" o "q").

Estado de entrenamiento#

Los usuarios pueden habilitar los puntos de control de entrenamiento a través de train(checkpoint_dir=...) y reanudar un trainer compatible con resume_from_checkpoint:

# Save a checkpoint at each completed training boundary.
trainer.train(checkpoint_dir="./checkpoints", checkpoint_interval=1)

# Recreate the trainer with the same models, tokenizers, data, and callbacks.
resumed_trainer.train(
    resume_from_checkpoint="./checkpoints/checkpoint-1",
    checkpoint_dir="./checkpoints",
    checkpoint_interval=1,
)

Los puntos de control preservan el estado del modelo y del optimizador, el progreso del entrenamiento, los estados aleatorios y, cuando corresponda, los datos y el historial PbRL.

Los puntos de control se guardan en las etapas completadas epochs, iteraciones o PbRL, con un guardado final al completar. checkpoint_interval controla los guardados periódicos. La reanudación repite cualquier trabajo desde el último punto de control.

num_train_epochs y num_iterations cuentan la ejecución total, incluyendo el trabajo completado, y pueden incrementarse al reanudar. Los modelos, tokenizers, los datos y la configuración de entrenamiento deben permanecer compatibles. Los conjuntos de datos y callbacks de estado pueden proporcionar state_dict() y load_state_dict(); el estado del servicio externo no se restaura. Si un directorio de salida PbRL ya contiene archivos, la salida reanudada va a un subdirectorio resume-*.

Para W&B, la identidad de la ejecución guardada se puede fusionar en wandb_config existente antes de que se construya trainer. Esto preserva la configuración de registro y salida y evita la apertura de una ejecución vacía durante la construcción:

from comlrl.utils import read_checkpoint_metadata

checkpoint = "./checkpoints/checkpoint-1"
identity = read_checkpoint_metadata(checkpoint)["wandb"]
if identity and "id" not in wandb_config and "resume" not in wandb_config:
    wandb_config = {**wandb_config, **identity, "resume": "allow"}

# Pass wandb_config to the trainer constructor before resuming.

Las configuraciones explícitas de id o resume tienen prioridad. Si el punto de control no contiene una identidad de W&B, el ayudante devuelve None para wandb y la configuración existente se mantiene.

El punto de control está deshabilitado por defecto: train() sin opciones de punto de control no guarda puntos de control de entrenamiento. save_model() exporta modelos solo; no guarda el estado completo necesario para reanudar el entrenamiento.