Entités d'entraînement

CoMLRL prend en charge l’entraînement des agents LLM et facilite leur entraînement avec des modèles de récompense basés sur LLM ou critics. Les agents et critics peuvent être des modèles homogènes ou hétérogènes provenant de Hugging Face ou de répertoires locaux. Les utilisateurs peuvent spécifier les identifiants de modèle Hugging Face via agent_model/critic_model pour les configurations homogènes, ou fournir des listes agents/critics pour les configurations hétérogènes. Les entités entraînées peuvent être enregistrées et chargées dans un nouveau trainer pour un entraînement ultérieur.

Agents LLM#

Agents homogènes#

Les agents homogènes sont initialisés en définissant agent_model sur un seul identifiant de modèle et en laissant agents=None. Pour trainers avec un argument de constructeur séparé num_agents, cet argument peut être défini conjointement avec la valeur de configuration correspondante ; la configuration seule ne le remplace pas.

L’exemple suivant charge trois agents Qwen/Qwen2.5-1.5B :

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

Agents hétérogènes#

Bien que les agents homogènes puissent être sollicités pour assumer différents rôles et travailler sur différentes parties d’une tâche en parallèle, les agents hétérogènes peuvent contribuer aux capacités spécialisées de différents modèles.

Les utilisateurs peuvent initialiser des agents hétérogènes en fournissant des identifiants de modèle ou des objets de modèle chargés dans agents, avec des tokenizers fournis via tokenizer. Avec agents fourni, agent_model peut rester None. Si les deux sont fournis, leurs noms de modèle doivent correspondre et leurs longueurs de liste doivent correspondre au nombre d’agents ; sinon, l’entraînement génère une erreur.

L’exemple suivant charge les agents Qwen/Qwen2.5-Coder-3B et 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#

Des algorithmes tels que Independent Actor-Critic (IAC) et Multi-Agent Actor-Critic (MAAC) utilisent LLM critics. La façon dont ces critics sont chargés dépend de l’algorithme et de la configuration use_separate_critic. MARLHF trainers avec rl_algorithm="maac" ou "iac" acceptent les mêmes arguments critic_model et critics, ou critic_model_name dans leur configuration ; sinon, ils utilisent l’identifiant de modèle du premier actor.

Centralisé Critic#

Une instance centralisée séparée critic peut être chargée via critic_model ou une liste d’un élément critics liste. Elle doit être un modèle de langage causal pris en charge : le trainer ajoute un value head et encode l’entrée conjointe avec le premier actor’s tokenizer, afin que tokenizer doit également être compatible avec le critic.

L’exemple suivant associe les agents Qwen/Qwen2.5-Coder-3B et Qwen/Qwen2.5-Coder-1.5B actors avec un modèle centralisé 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),
)

Independent Critics#

Avec une indépendance critics, chaque agent peut avoir son propre critic. Avec use_separate_critic=True, les utilisateurs peuvent fournir un critic_model identifiant pour charger une copie séparée par agent, ou une critics liste de longueur num_agents. Chaque copie séparée critic reçoit token IDs de son actor’s tokenizer et doit être compatible avec elle. Avec use_separate_critic=False, chaque actor et critic partagent un LLM backbone et sont chargés ensemble. Dans ce cas, critic_model et critics restent non définis (null ou None). Comme pour les agents, la fourniture de critic_model et critics nécessite des noms et des longueurs de modèles correspondants ; sinon, l’entraînement génère une erreur.

L’exemple suivant charge Qwen/Qwen2.5-Coder-3B et Qwen/Qwen2.5-Coder-7B ainsi que actors en utilisant des critics distincts, basés sur les mêmes modèles:

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

Alternativement, chaque actor peut partager le même modèle avec son 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 utilise agents/critics en interne. Les paramètres agent_model et critic_model sont des raccourcis pour les configurations homogènes et doivent être cohérents avec toutes les listes correspondantes.

MAGRPO charge tokenizers automatiquement pour les chaînes d’identifiants de modèle. MAAC/IAC chargent automatiquement tokenizer pour agent_model, mais nécessitent une spécification explicite de tokenizer lorsque agents est fourni, même s’il contient des chaînes. Les objets de modèle chargés nécessitent une liste partagée compatible tokenizer ou tokenizer correspondant à actors. Les modèles hétérogènes actors peuvent avoir des vocabulaires différents lorsque chacun utilise sa propre liste compatible tokenizer; les séparations critics doivent également accepter les actor de token IDs (le premier actor de IDs pour MAAC).

Modèles de récompense LLM#

Dans PbRL, MARLHF apprend un modèle de récompense conjoint à partir de paires de préférences et le maintient fixe pour l’apprentissage de la politique. Son backbone est chargé via reward_model_name dans MARLHFConfig, à partir d’un identifiant de modèle Hugging Face ou d’un répertoire local d’un modèle de langage causal. Il utilise par défaut l’identifiant du premier actor et est requis lorsque les actors sont des objets de modèle sans celui-ci. Le trainer charge le backbone avec son propre tokenizer, ajoute une tête scalaire sur le dernier élément non-remplit token et place le modèle sur reward_model_device (par défaut "cuda:0"; None utilise le trainer). reward_freeze_backbone=True entraîne uniquement la tête.

L’exemple suivant charge un modèle de récompense Qwen/Qwen2.5-Coder-7B pour deux agents 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 recharge le modèle de récompense à partir de cette source à chaque itération avant l’apprentissage des préférences. La récompense de tâche de reward_func continue de marquer les paires de préférences ; Initialisation de l’entraînement décrit comment passer les récompenses de tâche basées sur des règles ou basées sur LLM.

Sauvegarde et reprise de l’entraînement#

Sauvegarde des modèles#

Les utilisateurs peuvent appeler save_model(output_dir) après l’entraînement pour enregistrer explicitement les configurations, les poids et tokenizers du modèle actors et Hugging Face dans des répertoires locaux:

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

Entraînement à partir de modèles sauvegardés#

L’entraînement peut se poursuivre à partir d’un ensemble de deux actor et MAGRPO en fournissant les répertoires de actors dans l’ordre de rôle original. Le répertoire parent run-1 contient le actors ; chaque répertoire agent_i est un modèle chargeable.

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 charge tokenizers à partir des répertoires actor automatiquement. Ici, num_train_epochs=3 signifie trois epochs dans la nouvelle exécution d’entraînement. L’ensemble du jeu de données, la fonction de récompense, la configuration d’entraînement et tout formatteur personnalisé, processeur de récompense ou transition externe doivent être fournis à nouveau ; save_model ne les stocke pas.

D’autres trainers supportent la même approche de répertoire local, sous réserve des exigences tokenizer et actor ci-dessus. L’apprentissage des préférences trainers commence à collecter les préférences et à apprendre les récompenses conformément à la configuration de la nouvelle exécution.

IAC et MAAC sauvegardent value heads séparément sous forme de value_head.pt. Le chargement d’un répertoire critic ou partagé actor via le constructeur restaure le modèle de langage backbone et crée un nouveau value head. Les utilisateurs peuvent conserver les valeurs apprises en restaurant explicitement la tête sauvegardée avant d’appeler train(). Les points de contrôle d’entraînement, contrairement aux exportations de save_model, restaurent automatiquement value heads.

Pour un ensemble de deux actor et MAAC sauvegardés en utilisant la configuration linéaire par défaut de value head :

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

Pour IAC, la même étape load_state_dict s’applique à continued_iac.agents[i].value_head lorsque use_separate_critic=False, ou à continued_iac.critics[i].value_head lorsque c’est True. Les étapes critic backbones nécessitent également de charger via critic_model ou critics.

L’architecture value-head doit correspondre à l’exécution enregistrée. IAC nécessite IACConfig.value_head_hidden_dim pour les têtes partagées et IACConfig.critic_value_head_hidden_dim pour les têtes séparées. Une tête MAAC personnalisée nécessite model_config={"critic_value_head_hidden_dim": ...} dans le constructeur trainer. La poursuite de la même configuration d’apprentissage nécessite également critic_type ("v" ou "q").

État d’entraînement#

Les utilisateurs peuvent activer les points de contrôle d’entraînement via train(checkpoint_dir=...) et reprendre une exécution trainer compatible avec 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,
)

Les points de contrôle conservent l’état du modèle et de l’optimiseur, la progression de l’entraînement, les états aléatoires et, le cas échéant, les données et l’historique PbRL.

Les points de contrôle sont sauvegardés lors de la complétion de epochs, des itérations ou des étapes PbRL, avec une sauvegarde finale à la complétion. checkpoint_interval contrôle les sauvegardes périodiques. La reprise reprend tout le travail depuis le dernier point de contrôle.

num_train_epochs et num_iterations comptent l’exécution totale, y compris le travail terminé, et peuvent être augmentées lors de la reprise. Les modèles, tokenizers, les données et les paramètres d’entraînement doivent rester compatibles. Les ensembles de données et les appels de rappel d’état peuvent fournir state_dict() et load_state_dict(); l’état du service externe n’est pas restauré. Si un répertoire de sortie PbRL contient déjà des fichiers, la sortie reprise se place dans un sous-répertoire resume-*.

Pour W&B, l’identité de l’exécution enregistrée peut être fusionnée dans l’existant wandb_config avant la construction de trainer. Cela préserve les paramètres de journalisation et de sortie et évite de démarrer une exécution vide lors de la construction:

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.

Les paramètres explicites id ou resume ont priorité. Si le point de contrôle ne contient aucune identité W&B, l’aide renvoie None pour wandb et la configuration existante est conservée.

Le point de contrôle est désactivé par défaut : train() sans options de point de contrôle ne sauvegarde aucun point de contrôle d’entraînement. save_model() exporte uniquement les modèles ; il ne sauvegarde pas l’état complet nécessaire pour reprendre l’entraînement.