diff --git a/README.md b/README.md index c74e4d5..a79c8aa 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,16 @@ length. The examples leave W&B logging disabled by default; enable it explicitly python examples/story-len-ratio.py --wandb ``` +Two more examples cover other trainers and tasks. `leetcode-func-print.py` trains a +function-writing agent and a test-writing agent with MAGRPO, rewarding runnable code +and correct test prints. `tldr-len-ratio.py` trains two IAC agents to summarize Reddit +posts, the second TL;DR about 2–3 times longer than the first. + +```bash +python examples/leetcode-func-print.py +python examples/tldr-len-ratio.py +``` + For use in your own Python programs, CoMLRL is also available via PyPI or conda-forge: ```bash @@ -55,10 +65,12 @@ conda install -c conda-forge comlrl - **_Multi-Agent REINFORCE_:** Critic-free policy gradient methods, including [MAREINFORCE](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/mareinforce.py), [MAGRPO](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/magrpo.py), [MARLOO](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/marloo.py), [MAREMAX](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/maremax.py). - Aligned individual response joint with `joint_mode='aligned'`. - Cartesian-product joint responses with `joint_mode='cross'`. + - Decentralized actors or a centralized actor with `collaboration_mode='centralized'`. - **_Multi-Agent Actor-Critic:_** Actor-Critic methods, including [IAC](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/actor_critic/iac.py) and [MAAC](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/actor_critic/maac.py). - Independent actor-critic (separate critic or value-head over LLM backbone). - Centralized critic over joint prompts with separate actors. - - **_Multi-Agent Preference Learning:_** [MADPO](comlrl/trainers/preference/madpo.py), [MARLHF](comlrl/trainers/preference/marlhf.py), and their [iterative variants](comlrl/trainers/preference/iterative.py). + - Decentralized actors or a centralized actor with `collaboration_mode='centralized'`. + - **_Multi-Agent Preference Learning:_** [MADPO](comlrl/trainers/pbrl/madpo.py), [MARLHF](comlrl/trainers/pbrl/marlhf.py), and their [iterative variants](comlrl/trainers/pbrl/iterative.py). - Direct joint preference optimization or learned joint rewards followed by RL. - Decentralized actors or a centralized actor generating a complete joint response. diff --git a/comlrl/trainers/__init__.py b/comlrl/trainers/__init__.py index aa58f61..f1d9db2 100644 --- a/comlrl/trainers/__init__.py +++ b/comlrl/trainers/__init__.py @@ -1,3 +1,8 @@ +"""Trainers by family: reinforce (MAGRPO and its baselines), pbrl (preference-based RL) and +actor_critic. Every trainer runs in ``collaboration_mode="decentralized"`` (one actor per +agent) or ``"centralized"`` (one joint actor answering for all agents; see ``centralized``). +""" + from .actor_critic import ( ActorCriticTrainerBase, IACConfig, @@ -5,11 +10,21 @@ MAACConfig, MAACTrainer, ) -from .preference import ( +from .centralized import ( + CentralizedCollaboration, + CentralizedComparatorAdapter, + CentralizedComparatorParseError, + TaggedCentralizedComparatorAdapter, +) +from .pbrl import ( JointRewardModel, MADPOConfig, + MADPOIterConfig, + MADPOIterTrainer, MADPOTrainer, MARLHFConfig, + MARLHFIterConfig, + MARLHFIterTrainer, MARLHFTrainer, ) from .reinforce import ( @@ -29,13 +44,21 @@ "IACTrainer", "MAACConfig", "MAACTrainer", - "MAGRPOConfig", - "MAGRPOTrainer", + "CentralizedCollaboration", + "CentralizedComparatorAdapter", + "CentralizedComparatorParseError", + "TaggedCentralizedComparatorAdapter", + "JointRewardModel", "MADPOConfig", + "MADPOIterConfig", + "MADPOIterTrainer", "MADPOTrainer", "MARLHFConfig", + "MARLHFIterConfig", + "MARLHFIterTrainer", "MARLHFTrainer", - "JointRewardModel", + "MAGRPOConfig", + "MAGRPOTrainer", "MAREINFORCEConfig", "MAREINFORCETrainer", "MAReMaxConfig", diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index f0193eb..497cc10 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -1,5 +1,4 @@ from collections import defaultdict -import numbers from typing import Any, Dict, List, Optional, Tuple from concurrent.futures import ThreadPoolExecutor, as_completed @@ -11,12 +10,27 @@ from comlrl.models.actor_critic import CausalLMWithValueHead from comlrl.utils.checkpoint import TrainingCheckpointMixin +from comlrl.utils.distributed import unwrap_model from comlrl.utils.reference_kl import reference_kl_coef, reference_kl_enabled +from ..centralized import CollaborationModeMixin -class ActorCriticTrainerBase(TrainingCheckpointMixin): + +class ActorCriticTrainerBase(CollaborationModeMixin, TrainingCheckpointMixin): """Shared training utilities for actor-critic style trainers.""" + @property + def num_agents(self) -> int: + """Actor models trained: ``args.num_agents``, or one joint actor when centralized.""" + actors = self.__dict__.get("_num_actors") + if actors is None: + actors = getattr(getattr(self, "args", None), "num_agents", 0) + return int(actors or 0) + + @num_agents.setter + def num_agents(self, value: int) -> None: + self._num_actors = int(value) + def _kl_shaped_reward( self, reward: float, data: Dict[str, Any], index: int ) -> Tuple[float, Dict[str, float]]: @@ -34,7 +48,7 @@ def _kl_shaped_reward( } def _expand_rewards(self, rewards: List[float], num_ret: int) -> List[List[float]]: - num_agents = self.args.num_agents + num_agents = self.num_agents if len(rewards) == 1: return [[rewards[0]] * num_ret for _ in range(num_agents)] if len(rewards) == num_ret: @@ -45,6 +59,21 @@ def _expand_rewards(self, rewards: List[float], num_ret: int) -> List[List[float "Reward function must return 1 value, num_generations values, or num_agents values." ) + def _sample_loss_shares(self, policy_loss, value_error, batch_size: int): + """One sample's share of a mini-batch's mean actor and weighted value losses. + + Backpropagating each share as soon as it is computed keeps one sample's + activations in memory. The summed gradients are mathematically those of the + batch means; floating-point sums may differ slightly. Dividing before weighting + keeps every share within the range of the batch total, so low-precision losses + do not overflow where the batch mean would not. + """ + actor_share = policy_loss.mean() / batch_size + value_share = value_error.mean() / batch_size * self.args.value_loss_coef + if not torch.isfinite(actor_share) or not torch.isfinite(value_share): + raise FloatingPointError("Non-finite loss detected.") + return actor_share, value_share + def _value_on_prompt_only( self, model: CausalLMWithValueHead, @@ -105,7 +134,7 @@ def _compute_sequence_stats( def _parallel_agent_mode_enabled(self) -> bool: if str(getattr(self, "parallel_training", "")).lower() != "mp": return False - num_agents = int(getattr(getattr(self, "args", None), "num_agents", 0) or 0) + num_agents = self.num_agents if num_agents <= 1: return False devices = getattr(self, "agent_devices", None) @@ -121,7 +150,7 @@ def _run_agent_tasks( agent_indices: Optional[List[int]] = None, parallel: Optional[bool] = None, ) -> List[Any]: - num_agents = int(getattr(getattr(self, "args", None), "num_agents", 0) or 0) + num_agents = self.num_agents indices = ( list(agent_indices) if agent_indices is not None @@ -178,6 +207,26 @@ def _resolve_turn_prompt( modified_item, agent_idx, external_prompts=external_prompt ) + def _generation_kwargs(self, agent_idx: int) -> Dict[str, Any]: + """Extra ``generate`` keyword arguments, such as stop strings; none by default.""" + return {} + + def _generate_in_eval_mode( + self, agent_model, generation_kwargs: Dict[str, Any] + ) -> torch.Tensor: + """Sample in eval mode, then restore the actor's mode. + + In training mode HF would apply dropout and, with gradient checkpointing + on, generate without its KV cache. + """ + model = unwrap_model(agent_model) + training = model.training + model.eval() + try: + return model.generate(**generation_kwargs) + finally: + model.train(training) + def _get_tokenizer(self, agent_idx: Optional[int] = None): tokenizers = getattr(self, "tokenizers", None) if isinstance(tokenizers, list) and tokenizers: @@ -271,9 +320,7 @@ def _summarize_rollout_metrics(self, rollouts: List[Any]) -> Dict[str, float]: return metrics def _iter_dataloader(self, dataloader, epoch: int, total_epochs: int): - dist_env = getattr(self, "dist_env", None) - is_main = bool(getattr(dist_env, "is_main", True)) - if getattr(self, "verbose", True) and is_main: + if getattr(self, "verbose", True): return enumerate( tqdm( dataloader, @@ -299,12 +346,7 @@ def _on_epoch_end( epoch_metrics: Dict[str, List[float]], ) -> None: summary = self._summarize_epoch_metrics(epoch_metrics) - dist_env = getattr(self, "dist_env", None) - if ( - summary - and getattr(self, "verbose", True) - and getattr(dist_env, "is_main", True) - ): + if summary and getattr(self, "verbose", True): print(f"Epoch {epoch + 1}/{total_epochs} metrics: {summary}") def _tag_metrics( @@ -425,9 +467,9 @@ def _process(agent_idx: int) -> Dict[str, Any]: def _run_batch(self, batch, epoch_metrics: Dict[str, List[float]]) -> None: for item in batch: rollouts = self._collect_rollouts(item) - if self.args.num_agents > 0: + if self.num_agents > 0: # Count joint-action reward evaluations (one per agent group). - self.env_step += len(rollouts) // self.args.num_agents + self.env_step += len(rollouts) // self.num_agents self._log_domain_metrics(rollouts) ready_agents: List[int] = [] @@ -522,53 +564,26 @@ def train( checkpoint_total_limit=None, ) -> None: """Train with optional resume and checkpointing at complete epoch boundaries.""" - if ( - isinstance(checkpoint_interval, bool) - or not isinstance(checkpoint_interval, numbers.Integral) - or checkpoint_interval < 1 - ): - raise ValueError("checkpoint_interval must be a positive integer.") - if ( - resume_from_checkpoint is not None - or checkpoint_dir is not None - or checkpoint_interval != 1 - or checkpoint_total_limit is not None - ): - with self._checkpoint_context( - resume_from_checkpoint=resume_from_checkpoint, - checkpoint_dir=checkpoint_dir, - checkpoint_interval=checkpoint_interval, - checkpoint_total_limit=checkpoint_total_limit, - ): + checkpointing = self._checkpoint_arguments( + resume_from_checkpoint, + checkpoint_dir, + checkpoint_interval, + checkpoint_total_limit, + ) + if checkpointing is not None: + with self._checkpoint_context(**checkpointing): return self._train_checkpoint_epochs( self._checkpoint_progress, self._checkpoint_commit ) if not getattr(self, "_checkpoint_context_active", False): self._checkpoint_save_blocked = True - - total_epochs = self.args.num_train_epochs - - for epoch in range(total_epochs): - epoch_metrics = defaultdict(list) - dataloader = self.get_train_dataloader() - it = self._iter_dataloader(dataloader, epoch, total_epochs) - for batch_idx, batch in it: - if ( - self.eval_dataset is not None - and self.args.eval_interval > 0 - and batch_idx % int(self.args.eval_interval) == 0 - ): - self.evaluate() - self._run_batch(batch, epoch_metrics) - - self._flush_buffers(epoch_metrics) - self._on_epoch_end(epoch, total_epochs, epoch_metrics) + self._run_epochs(0, None) if not getattr(self, "_checkpoint_context_active", False): self._checkpoint_progress = { "kind": "actor_critic", - "epoch": total_epochs, + "epoch": self.args.num_train_epochs, } self._checkpoint_save_blocked = False @@ -609,7 +624,6 @@ def _on_checkpoint_train_start(self, start_epoch: int, total_epochs: int) -> Non def _train_checkpoint_epochs(self, progress, commit) -> None: """Run complete RL epochs, optionally inside a parent trainer's context.""" - total_epochs = self.args.num_train_epochs start_epoch = ActorCriticTrainerBase._checkpoint_validate_progress( self, { @@ -617,9 +631,16 @@ def _train_checkpoint_epochs(self, progress, commit) -> None: "runtime": {"rollout_buffers": self.rollout_buffers}, }, ) - self._on_checkpoint_train_start(start_epoch, total_epochs) + self._on_checkpoint_train_start(start_epoch, self.args.num_train_epochs) self._checkpoint_init_wandb() + self._run_epochs(start_epoch, commit) + def _run_epochs(self, start_epoch, commit) -> None: + """Train epochs ``start_epoch`` onward; ``commit`` (if given) follows each epoch.""" + # from_pretrained returns eval mode, in which HF skips gradient checkpointing. + for model in [*self.agents, *self.critics]: + model.train() + total_epochs = self.args.num_train_epochs for epoch in range(start_epoch, total_epochs): epoch_metrics = defaultdict(list) dataloader = self.get_train_dataloader() @@ -635,7 +656,8 @@ def _train_checkpoint_epochs(self, progress, commit) -> None: self._flush_buffers(epoch_metrics) self._on_epoch_end(epoch, total_epochs, epoch_metrics) - commit( - {"kind": "actor_critic", "epoch": epoch + 1}, - force=epoch + 1 == total_epochs, - ) + if commit is not None: + commit( + {"kind": "actor_critic", "epoch": epoch + 1}, + force=epoch + 1 == total_epochs, + ) diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index 552ea1d..e6f24b6 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -11,22 +11,23 @@ from comlrl.models.actor_critic import CausalLMWithValueHead from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training -from comlrl.utils.distributed import local_context, unwrap_model +from comlrl.utils.checkpoint import register_checkpoint_record +from comlrl.utils.distributed import unwrap_model from comlrl.utils.formatters import build_formatters from comlrl.utils.model_loading import resolve_model_kwargs, resolve_model_sources from comlrl.utils.reference_kl import ( - _sequence_kl_values, clone_reference_models, load_reference_models_from_sources, reference_kl_coef, reference_kl_enabled, - reference_kl_for_sequence, + reference_kl_values, resolve_reference_devices, validate_reference_kl_config, ) from comlrl.utils.reward_utils import call_reward_function, normalize_reward_lengths from comlrl.utils.tokenizer_utils import apply_tokenizer_specials, resolve_tokenizers from comlrl.utils.wandb_utils import verbose_from_wandb_config, wandb_init_kwargs +from ..centralized import validate_collaboration_mode from .ac_base import ActorCriticTrainerBase import wandb @@ -58,6 +59,8 @@ class IACConfig: value_head_hidden_dim: Optional[int] = None num_agents: int = 2 num_turns: int = 2 + # "centralized" trains one joint actor that answers for all num_agents roles. + collaboration_mode: str = "decentralized" parallel_training: str = "none" agent_devices: Optional[Union[str, Sequence[str]]] = None critic_devices: Optional[Union[str, Sequence[str]]] = None @@ -72,6 +75,9 @@ class IACConfig: reference_kl_enabled: bool = False reference_kl_coef: float = 0.1 reference_devices: Optional[Union[str, Sequence[str]]] = None + # Mathematically the same update, holding one sample's activations at a time + # (last, so positional arguments keep their meaning). + per_sample_backward: bool = False def __post_init__(self) -> None: if self.rollout_buffer_size < 1: @@ -84,6 +90,9 @@ def __post_init__(self) -> None: raise ValueError("num_agents must be >= 1.") if self.num_turns < 1: raise ValueError("num_turns must be >= 1.") + validate_collaboration_mode(self.collaboration_mode) + if self.collaboration_mode == "centralized" and self.num_turns != 1: + raise ValueError("Centralized IAC currently supports num_turns=1 only.") if self.num_turns > 1 and self.num_generations != 1: raise ValueError("Multi-turn IAC currently supports num_generations == 1.") critic_type = (self.critic_type or "v").lower() @@ -113,9 +122,12 @@ def __post_init__(self) -> None: "use_separate_critic=True." ) self.parallel_training = mode - validate_reference_kl_config(self, self.num_agents) + validate_reference_kl_config( + self, 1 if self.collaboration_mode == "centralized" else self.num_agents + ) +@register_checkpoint_record @dataclass class RolloutSample: agent_idx: int @@ -162,8 +174,18 @@ def __init__( critics: Optional[ Sequence[Union[str, PreTrainedModel, CausalLMWithValueHead]] ] = None, + *, + centralized_adapter: Optional[Any] = None, ) -> None: self.args = args if args is not None else IACConfig() + reward_func, formatters, self.num_agents = self._setup_collaboration( + centralized_adapter, + agents=agents, + tokenizer=tokenizer, + reward_func=reward_func, + formatters=formatters, + num_agents=self.args.num_agents, + ) if reward_func is None or not callable(reward_func): raise ValueError("reward_func must be a callable.") if agent_model is None and agents is None: @@ -176,7 +198,7 @@ def __init__( ) if ( agents is None - and self.args.num_agents > 1 + and self.num_agents > 1 and isinstance(agent_model, PreTrainedModel) ): raise ValueError( @@ -190,7 +212,7 @@ def __init__( self.reward_processor = reward_processor or (lambda x: x) self.train_dataset = train_dataset self.eval_dataset = eval_dataset - self.metrics_callback = metrics_callback + self.metrics_callback = self._role_metrics_callback(metrics_callback) self.model_config = model_config or {} self.critic_type = (self.args.critic_type or "v").lower() self.parallel_training = normalize_parallel_training( @@ -198,7 +220,7 @@ def __init__( ) if self.parallel_training == "mp": self.agent_devices, self.critic_devices = DeviceScheduler.assign_devices( - self.args.num_agents, + self.num_agents, getattr(self.args, "agent_devices", None), getattr(self.args, "critic_devices", None), use_separate_critic=self.args.use_separate_critic, @@ -208,10 +230,9 @@ def __init__( getattr(self.args, "agent_devices", None), getattr(self.args, "critic_devices", None), ) - self.agent_devices = [single_device] * self.args.num_agents - self.critic_devices = [single_device] * self.args.num_agents + self.agent_devices = [single_device] * self.num_agents + self.critic_devices = [single_device] * self.num_agents self.device = self.agent_devices[0] - self.dist_env = local_context(self.device) self.agents: List[CausalLMWithValueHead] = [] self.critics: List[CausalLMWithValueHead] = [] @@ -221,9 +242,9 @@ def __init__( self.tokenizers = tokenizers self.tokenizer = tokenizers[0] if tokenizers else None else: - self.tokenizers = [tokenizers] * self.args.num_agents + self.tokenizers = [tokenizers] * self.num_agents self.tokenizer = tokenizers - self.formatters = build_formatters(formatters, self.args.num_agents) + self.formatters = build_formatters(formatters, self.num_agents) try: self._reward_signature = inspect.signature(reward_func) except (TypeError, ValueError): @@ -234,7 +255,7 @@ def __init__( kind="agents", model=agent_model, models=agents, - expected_count=self.args.num_agents, + expected_count=self.num_agents, model_label="agent_model", ) actor_model_kwargs = resolve_model_kwargs( @@ -280,7 +301,7 @@ def __init__( kind="critics", model=critic_model, models=critics, - expected_count=self.args.num_agents, + expected_count=self.num_agents, model_label="critic_model", ) for idx, critic_source in enumerate(critic_sources): @@ -322,7 +343,7 @@ def __init__( self.reference_devices = resolve_reference_devices( self.args, self.agent_devices, - self.args.num_agents, + self.num_agents, ) if actor_sources and all(isinstance(src, str) for src in actor_sources): self.reference_models = load_reference_models_from_sources( @@ -368,11 +389,11 @@ def __init__( self.critic_optimizers.append(optimizer) self.env_step = 0 - self.rollout_buffers = [[] for _ in range(self.args.num_agents)] + self.rollout_buffers = [[] for _ in range(self.num_agents)] self.wandb_config = wandb_config self.wandb_initialized = False self._last_train_log_step = -1 - if wandb_config is not None and self.dist_env.is_main: + if wandb_config is not None: self._init_wandb() self.verbose = verbose_from_wandb_config(self.wandb_config) @@ -382,8 +403,6 @@ def _critic_device(self, agent_idx: int) -> torch.device: return self.agent_devices[agent_idx] def _init_wandb(self) -> None: - if not self.dist_env.is_main: - return if self.wandb_initialized: return if self.wandb_config is None: @@ -391,7 +410,7 @@ def _init_wandb(self) -> None: config_dict: Dict[str, Any] = { "algorithm": self.algorithm_name, - "num_agents": self.args.num_agents, + "num_agents": self.num_agents, "num_turns": self.args.num_turns, "agent_learning_rate": self.args.agent_learning_rate, "critic_learning_rate": self.args.critic_learning_rate, @@ -406,6 +425,11 @@ def _init_wandb(self) -> None: "reference_devices": getattr(self.args, "reference_devices", None), } + if self.is_centralized: + config_dict["collaboration_mode"] = "centralized" + config_dict["num_roles"] = self.centralized_collaboration.num_roles + config_dict["num_actor_models"] = len(self.agents) + previous_run = wandb.run self._checkpoint_record_wandb_run( wandb.init(**wandb_init_kwargs(self.wandb_config, config_dict)), @@ -442,9 +466,9 @@ def _generate_rollout( } if self.args.top_k is not None: generation_kwargs["top_k"] = self.args.top_k + generation_kwargs.update(self._generation_kwargs(agent_idx)) - generation_model = unwrap_model(agent_model) - sequences = generation_model.generate(**generation_kwargs) + sequences = self._generate_in_eval_mode(agent_model, generation_kwargs) if sequences.size(1) <= prompt_len: raise RuntimeError("Model produced an empty completion during rollout.") @@ -521,31 +545,7 @@ def _generate_rollout( "completions": completion_texts, } - def _reference_kl_values( - self, - agent_idx: int, - policy_model: CausalLMWithValueHead, - sequences: torch.Tensor, - attention_mask: torch.Tensor, - prompt_len: int, - response_lens: Sequence[int], - ) -> List[float]: - if not reference_kl_enabled(self.args): - return [] - if not getattr(self, "reference_models", None): - raise RuntimeError( - "Reference KL is enabled but reference models are missing." - ) - reference_model = self.reference_models[agent_idx] - return _sequence_kl_values( - policy_model, - reference_model, - sequences, - attention_mask, - prompt_len, - response_lens, - sequence_kl=reference_kl_for_sequence, - ) + _reference_kl_values = reference_kl_values def _collect_rollouts(self, item: Dict[str, Any]) -> List[RolloutSample]: num_turns = max(1, int(getattr(self.args, "num_turns", 1))) @@ -559,7 +559,7 @@ def _collect_rollouts(self, item: Dict[str, Any]) -> List[RolloutSample]: ) turn_prompts = [ self._resolve_turn_prompt(item, agent_idx) - for agent_idx in range(self.args.num_agents) + for agent_idx in range(self.num_agents) ] def _generate_agent(agent_idx: int) -> Dict[str, Any]: @@ -578,13 +578,13 @@ def _generate_agent(agent_idx: int) -> Dict[str, Any]: self.reward_func, prompts, completions_per_agent, - num_agents=self.args.num_agents, + num_agents=self.num_agents, batch_items=[item], signature=self._reward_signature, ) rewards = normalize_reward_lengths( [float(self.reward_processor(r)) for r in rewards], - num_agents=self.args.num_agents, + num_agents=self.num_agents, num_generations=num_ret, algorithm="IAC", ) @@ -645,11 +645,11 @@ def _collect_rollouts_multi_turn( if self.external_transition is None: raise ValueError("external_transition is required for multi-turn IAC.") - prompt_history = [[] for _ in range(self.args.num_agents)] - response_history = [[] for _ in range(self.args.num_agents)] - previous_completions: List[Optional[str]] = [None] * self.args.num_agents + prompt_history = [[] for _ in range(self.num_agents)] + response_history = [[] for _ in range(self.num_agents)] + previous_completions: List[Optional[str]] = [None] * self.num_agents per_agent_samples: List[List[RolloutSample]] = [ - [] for _ in range(self.args.num_agents) + [] for _ in range(self.num_agents) ] rollouts: List[RolloutSample] = [] value_discount = float(getattr(self.args, "discount", 1.0)) @@ -658,19 +658,19 @@ def _collect_rollouts_multi_turn( if turn_idx == 0: turn_prompts = [ self._resolve_turn_prompt(item, agent_idx) - for agent_idx in range(self.args.num_agents) + for agent_idx in range(self.num_agents) ] else: transition_result = self.external_transition( prompt=item.get("prompt", ""), agent_completions=previous_completions, - num_agents=self.args.num_agents, + num_agents=self.num_agents, prompt_history_per_agent=prompt_history, response_history_per_agent=response_history, ) if ( not isinstance(transition_result, (list, tuple)) - or len(transition_result) != self.args.num_agents + or len(transition_result) != self.num_agents ): raise ValueError( "External transition must return per-agent prompts" @@ -680,7 +680,7 @@ def _collect_rollouts_multi_turn( self._resolve_turn_prompt( item, agent_idx, external_prompt=external_prompts[agent_idx] ) - for agent_idx in range(self.args.num_agents) + for agent_idx in range(self.num_agents) ] completions_per_agent: List[List[str]] = [] @@ -700,13 +700,13 @@ def _generate_agent_turn(agent_idx: int) -> Dict[str, Any]: self.reward_func, turn_prompts, completions_per_agent, - num_agents=self.args.num_agents, + num_agents=self.num_agents, batch_items=[item], signature=self._reward_signature, ) rewards = normalize_reward_lengths( [float(self.reward_processor(r)) for r in rewards], - num_agents=self.args.num_agents, + num_agents=self.num_agents, num_generations=1, algorithm="IAC", ) @@ -759,7 +759,7 @@ def _generate_agent_turn(agent_idx: int) -> Dict[str, Any]: if mean_reward > float(term_threshold): break - for agent_idx in range(self.args.num_agents): + for agent_idx in range(self.num_agents): traj = per_agent_samples[agent_idx] for t, sample in enumerate(traj): r = float(sample.returns.view(-1)[0].item()) @@ -771,7 +771,7 @@ def _generate_agent_turn(agent_idx: int) -> Dict[str, Any]: sample.metadata["adv_target"] = torch.tensor([target]).cpu() sample.metadata["value_target"] = torch.tensor([target]).cpu() - for agent_idx in range(self.args.num_agents): + for agent_idx in range(self.num_agents): future = 0.0 for sample in reversed(per_agent_samples[agent_idx]): immediate = float(sample.returns.view(-1)[0].item()) @@ -878,6 +878,13 @@ def _ac_step(self, agent_idx: int, batch: List[RolloutSample]) -> None: value_losses: List[torch.Tensor] = [] agent_device = self.agent_devices[agent_idx] critic_device = self._critic_device(agent_idx) + per_sample = getattr(self.args, "per_sample_backward", False) + if per_sample: + if self.args.use_separate_critic and critic_optimizer is None: + raise RuntimeError("Critic optimizer missing.") + agent_optimizer.zero_grad() + if critic_optimizer is not None: + critic_optimizer.zero_grad() for sample in batch: sequences = sample.full_input_ids.to(agent_device).unsqueeze(0) @@ -960,9 +967,24 @@ def _ac_step(self, agent_idx: int, batch: List[RolloutSample]) -> None: else: value_error = (value_target - value) ** 2 + if per_sample: + actor_share, value_share = self._sample_loss_shares( + policy_loss, value_error, len(batch) + ) + if self.args.use_separate_critic: + actor_share.backward() + value_share.backward() + else: + (actor_share + value_share).backward() + continue actor_losses.append(policy_loss) value_losses.append(value_error) + if per_sample: + agent_optimizer.step() + if self.args.use_separate_critic: + critic_optimizer.step() + return actor_loss = torch.stack(actor_losses).mean() value_loss = torch.stack(value_losses).mean() if not torch.isfinite(actor_loss) or not torch.isfinite(value_loss): @@ -1040,10 +1062,8 @@ def _update( return averaged def save_model(self, output_dir: str) -> None: - if not self.dist_env.is_main: - return os.makedirs(output_dir, exist_ok=True) - if self.args.num_agents == 1: + if self.num_agents == 1: actor = unwrap_model(self.agents[0]) actor.model.save_pretrained(output_dir) if actor.value_head is not None: @@ -1085,7 +1105,7 @@ def save_model(self, output_dir: str) -> None: ) if self.tokenizers: - if self.args.num_agents == 1: + if self.num_agents == 1: self.tokenizers[0].save_pretrained(output_dir) else: for idx, tok in enumerate(self.tokenizers): diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index 27a2872..5fbfd46 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -11,22 +11,22 @@ from comlrl.models.actor_critic import CausalLMWithValueHead from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training -from comlrl.utils.distributed import local_context, unwrap_model +from comlrl.utils.distributed import unwrap_model from comlrl.utils.formatters import build_formatters from comlrl.utils.model_loading import resolve_model_kwargs, resolve_model_sources from comlrl.utils.reference_kl import ( - _sequence_kl_values, clone_reference_models, load_reference_models_from_sources, reference_kl_coef, reference_kl_enabled, - reference_kl_for_sequence, + reference_kl_values, resolve_reference_devices, validate_reference_kl_config, ) from comlrl.utils.reward_utils import call_reward_function, normalize_reward_lengths from comlrl.utils.tokenizer_utils import apply_tokenizer_specials, resolve_tokenizers from comlrl.utils.wandb_utils import verbose_from_wandb_config, wandb_init_kwargs +from ..centralized import validate_collaboration_mode from .ac_base import ActorCriticTrainerBase from .iac import RolloutSample @@ -53,6 +53,8 @@ class MAACConfig: num_agents: int = 2 num_generations: int = 1 num_turns: int = 2 + # "centralized" trains one joint actor that answers for all num_agents roles. + collaboration_mode: str = "decentralized" parallel_training: str = "none" agent_devices: Optional[Union[str, Sequence[str]]] = None critic_devices: Optional[Union[str, Sequence[str]]] = None @@ -67,6 +69,9 @@ class MAACConfig: reference_kl_enabled: bool = False reference_kl_coef: float = 0.1 reference_devices: Optional[Union[str, Sequence[str]]] = None + # Mathematically the same update, holding one sample's activations at a time + # (last, so positional arguments keep their meaning). + per_sample_backward: bool = False def __post_init__(self) -> None: if self.rollout_buffer_size < 1: @@ -81,6 +86,9 @@ def __post_init__(self) -> None: raise ValueError("num_generations must be >= 1.") if self.num_turns < 1: raise ValueError("num_turns must be >= 1.") + validate_collaboration_mode(self.collaboration_mode) + if self.collaboration_mode == "centralized" and self.num_turns != 1: + raise ValueError("Centralized MAAC currently supports num_turns=1 only.") if self.num_turns > 1 and self.num_generations != 1: raise ValueError("Multi-turn MAAC currently supports num_generations == 1.") critic_type = (self.critic_type or "v").lower() @@ -105,7 +113,9 @@ def __post_init__(self) -> None: "parallel_training='mp' requires explicit critic_devices." ) self.parallel_training = mode - validate_reference_kl_config(self, self.num_agents) + validate_reference_kl_config( + self, 1 if self.collaboration_mode == "centralized" else self.num_agents + ) class MAACTrainer(ActorCriticTrainerBase): @@ -136,15 +146,25 @@ def __init__( critics: Optional[ Sequence[Union[str, PreTrainedModel, CausalLMWithValueHead]] ] = None, + *, + centralized_adapter: Optional[Any] = None, ) -> None: self.args = args if args is not None else MAACConfig() + reward_func, formatters, self.num_agents = self._setup_collaboration( + centralized_adapter, + agents=agents, + tokenizer=tokenizer, + reward_func=reward_func, + formatters=formatters, + num_agents=self.args.num_agents, + ) if reward_func is None or not callable(reward_func): raise ValueError("reward_func must be a callable.") if agent_model is None and agents is None: raise ValueError("Either agent_model or agents must be provided.") if ( agents is None - and self.args.num_agents > 1 + and self.num_agents > 1 and isinstance(agent_model, PreTrainedModel) ): raise ValueError( @@ -162,7 +182,7 @@ def __init__( self.reward_processor = reward_processor or (lambda x: x) self.train_dataset = train_dataset self.eval_dataset = eval_dataset - self.metrics_callback = metrics_callback + self.metrics_callback = self._role_metrics_callback(metrics_callback) self.model_config = model_config or {} self.parallel_training = normalize_parallel_training( getattr(self.args, "parallel_training", "none") @@ -170,7 +190,7 @@ def __init__( if self.parallel_training == "mp": self.agent_devices = DeviceScheduler.resolve_devices( getattr(self.args, "agent_devices", None), - self.args.num_agents, + self.num_agents, kind="agent_devices", ) self.critic_device = DeviceScheduler.assign_shared_critic_device( @@ -181,10 +201,9 @@ def __init__( getattr(self.args, "agent_devices", None), getattr(self.args, "critic_devices", None), ) - self.agent_devices = [single_device] * self.args.num_agents + self.agent_devices = [single_device] * self.num_agents self.critic_device = single_device self.device = self.agent_devices[0] - self.dist_env = local_context(self.device) self._parallel_update_enabled = False tokenizers = resolve_tokenizers(agent_model, tokenizer, agents) @@ -192,7 +211,7 @@ def __init__( self.tokenizers = tokenizers self.tokenizer = tokenizers[0] if tokenizers else None else: - self.tokenizers = [tokenizers] * self.args.num_agents + self.tokenizers = [tokenizers] * self.num_agents self.tokenizer = tokenizers self.external_transition = external_transition @@ -201,7 +220,7 @@ def __init__( kind="agents", model=agent_model, models=agents, - expected_count=self.args.num_agents, + expected_count=self.num_agents, model_label="agent_model", ) actor_model_kwargs = resolve_model_kwargs( @@ -285,7 +304,7 @@ def __init__( self.reference_devices = resolve_reference_devices( self.args, self.agent_devices, - self.args.num_agents, + self.num_agents, ) if actor_sources and all(isinstance(src, str) for src in actor_sources): self.reference_models = load_reference_models_from_sources( @@ -311,7 +330,7 @@ def __init__( self.tokenizer, [*self.agents, self.critics[0], *self.reference_models] ) - self.formatters = build_formatters(formatters, self.args.num_agents) + self.formatters = build_formatters(formatters, self.num_agents) try: self._reward_signature = inspect.signature(reward_func) except (TypeError, ValueError): @@ -330,23 +349,21 @@ def __init__( lr=self.args.critic_learning_rate, ) - self.rollout_buffers = [[] for _ in range(self.args.num_agents)] + self.rollout_buffers = [[] for _ in range(self.num_agents)] self.wandb_config = wandb_config self.wandb_initialized = False self.env_step = 0 self._last_train_log_step = -1 - if wandb_config is not None and self.dist_env.is_main: + if wandb_config is not None: self._init_wandb() self.verbose = verbose_from_wandb_config(self.wandb_config) def _init_wandb(self) -> None: - if not self.dist_env.is_main: - return if self.wandb_config is None: self.wandb_config = {} config_dict: Dict[str, Any] = { "algorithm": self.algorithm_name, - "num_agents": self.args.num_agents, + "num_agents": self.num_agents, "num_turns": self.args.num_turns, "agent_learning_rate": self.args.agent_learning_rate, "critic_learning_rate": self.args.critic_learning_rate, @@ -358,6 +375,11 @@ def _init_wandb(self) -> None: "reference_devices": getattr(self.args, "reference_devices", None), } + if self.is_centralized: + config_dict["collaboration_mode"] = "centralized" + config_dict["num_roles"] = self.centralized_collaboration.num_roles + config_dict["num_actor_models"] = len(self.agents) + previous_run = wandb.run self._checkpoint_record_wandb_run( wandb.init(**wandb_init_kwargs(self.wandb_config, config_dict)), @@ -383,6 +405,22 @@ def _build_critic_input( action_lines.append(f"[Agent {idx} action]\n{comp}") return base + "\n\n" + "\n\n".join(action_lines) + def _critic_input_for( + self, + item: Dict[str, Any], + prompts: Sequence[str], + joint_action: Optional[Sequence[str]] = None, + ) -> str: + """Critic text for one item; a joint actor's output is split into role outputs.""" + collaboration = self.centralized_collaboration + if collaboration is None: + return self._build_critic_input(prompts, joint_action) + role_prompts = [formatter(item) for formatter in collaboration.formatters] + role_actions = ( + None if joint_action is None else collaboration.split(joint_action[0], item) + ) + return self._build_critic_input(role_prompts, role_actions) + def _critic_value_from_text(self, critic_input: str) -> Dict[str, Any]: encoded = self._encode_prompt( critic_input, @@ -427,9 +465,9 @@ def _generate(self, agent_model, prompt: str, agent_idx: int) -> Dict[str, Any]: } if self.args.top_k is not None: generation_kwargs["top_k"] = self.args.top_k + generation_kwargs.update(self._generation_kwargs(agent_idx)) - generation_model = unwrap_model(agent_model) - sequences = generation_model.generate(**generation_kwargs) + sequences = self._generate_in_eval_mode(agent_model, generation_kwargs) if sequences.size(1) <= prompt_len: raise RuntimeError("Model produced an empty completion during rollout.") @@ -467,31 +505,7 @@ def _generate(self, agent_model, prompt: str, agent_idx: int) -> Dict[str, Any]: "completions": completion_texts, } - def _reference_kl_values( - self, - agent_idx: int, - policy_model: CausalLMWithValueHead, - sequences: torch.Tensor, - attention_mask: torch.Tensor, - prompt_len: int, - response_lens: Sequence[int], - ) -> List[float]: - if not reference_kl_enabled(self.args): - return [] - if not getattr(self, "reference_models", None): - raise RuntimeError( - "Reference KL is enabled but reference models are missing." - ) - reference_model = self.reference_models[agent_idx] - return _sequence_kl_values( - policy_model, - reference_model, - sequences, - attention_mask, - prompt_len, - response_lens, - sequence_kl=reference_kl_for_sequence, - ) + _reference_kl_values = reference_kl_values def _collect_rollouts(self, item: Dict[str, Any]) -> List[RolloutSample]: num_turns = max(1, int(getattr(self.args, "num_turns", 1))) @@ -505,7 +519,7 @@ def _collect_rollouts(self, item: Dict[str, Any]) -> List[RolloutSample]: ) turn_prompts = [ self._resolve_turn_prompt(item, agent_idx) - for agent_idx in range(self.args.num_agents) + for agent_idx in range(self.num_agents) ] def _generate_agent(agent_idx: int) -> Dict[str, Any]: @@ -533,13 +547,13 @@ def _generate_agent(agent_idx: int) -> Dict[str, Any]: self.reward_func, prompts, completions_per_agent, - num_agents=self.args.num_agents, + num_agents=self.num_agents, batch_items=[item], signature=self._reward_signature, ) rewards = normalize_reward_lengths( [float(self.reward_processor(r)) for r in rewards], - num_agents=self.args.num_agents, + num_agents=self.num_agents, num_generations=num_ret, algorithm="MAAC", ) @@ -552,15 +566,15 @@ def _generate_agent(agent_idx: int) -> Dict[str, Any]: critic_type = (self.args.critic_type or "v").lower() critic_values_by_i: List[Dict[str, Any]] = [] if critic_type == "v": - critic_input = self._build_critic_input(prompts) + critic_input = self._critic_input_for(item, prompts) with torch.no_grad(): critic_values_by_i = [self._critic_value_from_text(critic_input)] else: for i in range(num_ret): joint_action = [ - completions_per_agent[a][i] for a in range(self.args.num_agents) + completions_per_agent[a][i] for a in range(self.num_agents) ] - critic_input = self._build_critic_input(prompts, joint_action) + critic_input = self._critic_input_for(item, prompts, joint_action) with torch.no_grad(): critic_values_by_i.append( self._critic_value_from_text(critic_input) @@ -639,11 +653,11 @@ def _collect_rollouts_multi_turn( if self.args.num_generations != 1: raise ValueError("Multi-turn MAAC currently supports num_generations == 1.") - prompt_history = [[] for _ in range(self.args.num_agents)] - response_history = [[] for _ in range(self.args.num_agents)] - previous_completions: List[Optional[str]] = [None] * self.args.num_agents + prompt_history = [[] for _ in range(self.num_agents)] + response_history = [[] for _ in range(self.num_agents)] + previous_completions: List[Optional[str]] = [None] * self.num_agents per_agent_samples: List[List[RolloutSample]] = [ - [] for _ in range(self.args.num_agents) + [] for _ in range(self.num_agents) ] rollouts: List[RolloutSample] = [] value_discount = float(getattr(self.args, "discount", 1.0)) @@ -652,7 +666,7 @@ def _collect_rollouts_multi_turn( if turn_idx == 0: turn_prompts = [ self._resolve_turn_prompt(item, agent_idx) - for agent_idx in range(self.args.num_agents) + for agent_idx in range(self.num_agents) ] else: if self.external_transition is None: @@ -660,13 +674,13 @@ def _collect_rollouts_multi_turn( transition_result = self.external_transition( prompt=item.get("prompt", ""), agent_completions=previous_completions, - num_agents=self.args.num_agents, + num_agents=self.num_agents, prompt_history_per_agent=prompt_history, response_history_per_agent=response_history, ) if ( not isinstance(transition_result, (list, tuple)) - or len(transition_result) != self.args.num_agents + or len(transition_result) != self.num_agents ): raise ValueError( "External transition must return per-agent prompts" @@ -676,7 +690,7 @@ def _collect_rollouts_multi_turn( self._resolve_turn_prompt( item, agent_idx, external_prompt=external_prompts[agent_idx] ) - for agent_idx in range(self.args.num_agents) + for agent_idx in range(self.num_agents) ] completions_per_agent: List[List[str]] = [] @@ -707,13 +721,13 @@ def _generate_agent_turn(agent_idx: int) -> Dict[str, Any]: self.reward_func, turn_prompts, completions_per_agent, - num_agents=self.args.num_agents, + num_agents=self.num_agents, batch_items=[item], signature=self._reward_signature, ) rewards = normalize_reward_lengths( [float(self.reward_processor(r)) for r in rewards], - num_agents=self.args.num_agents, + num_agents=self.num_agents, num_generations=1, algorithm="MAAC", ) @@ -786,7 +800,7 @@ def _generate_agent_turn(agent_idx: int) -> Dict[str, Any]: if mean_reward > float(term_threshold): break - for agent_idx in range(self.args.num_agents): + for agent_idx in range(self.num_agents): traj = per_agent_samples[agent_idx] for t, sample in enumerate(traj): r = float(sample.returns.view(-1)[0].item()) @@ -798,7 +812,7 @@ def _generate_agent_turn(agent_idx: int) -> Dict[str, Any]: sample.metadata["adv_target"] = torch.tensor([target]).cpu() sample.metadata["value_target"] = torch.tensor([target]).cpu() - for agent_idx in range(self.args.num_agents): + for agent_idx in range(self.num_agents): future = 0.0 for sample in reversed(per_agent_samples[agent_idx]): immediate = float(sample.returns.view(-1)[0].item()) @@ -849,6 +863,10 @@ def _ac_step(self, agent_idx: int, batch: List[RolloutSample]) -> None: value_losses: List[torch.Tensor] = [] agent_device = self.agent_devices[agent_idx] critic_device = self.critic_device + per_sample = getattr(self.args, "per_sample_backward", False) + if per_sample: + agent_optimizer.zero_grad() + self.critic_optimizer.zero_grad() for sample in batch: sequences = sample.full_input_ids.to(agent_device).unsqueeze(0) @@ -896,9 +914,20 @@ def _ac_step(self, agent_idx: int, batch: List[RolloutSample]) -> None: policy_loss = -(logprob * policy_advantage) value_error = (returns - value) ** 2 + if per_sample: + actor_share, value_share = self._sample_loss_shares( + policy_loss, value_error, len(batch) + ) + actor_share.backward() + value_share.backward() + continue actor_losses.append(policy_loss) value_losses.append(value_error) + if per_sample: + agent_optimizer.step() + self.critic_optimizer.step() + return actor_loss = torch.stack(actor_losses).mean() value_loss = torch.stack(value_losses).mean() value_total = self.args.value_loss_coef * value_loss @@ -929,8 +958,6 @@ def _update( return metrics def save_model(self, output_dir: str) -> None: - if not self.dist_env.is_main: - return os.makedirs(output_dir, exist_ok=True) for agent_idx, actor in enumerate(self.agents): actor = unwrap_model(actor) diff --git a/comlrl/trainers/centralized/__init__.py b/comlrl/trainers/centralized/__init__.py new file mode 100644 index 0000000..bdc5760 --- /dev/null +++ b/comlrl/trainers/centralized/__init__.py @@ -0,0 +1,29 @@ +"""Pieces shared by every trainer in centralized collaboration mode (one joint actor).""" + +from .adapters import ( + CentralizedComparatorAdapter, + CentralizedComparatorParseError, + TaggedCentralizedComparatorAdapter, +) +from .collaboration import ( + CentralizedCollaboration, + joint_policy_gradient_loss, + joint_sequence_log_prob, +) +from .mode import ( + COLLABORATION_MODES, + CollaborationModeMixin, + validate_collaboration_mode, +) + +__all__ = [ + "COLLABORATION_MODES", + "CentralizedCollaboration", + "CentralizedComparatorAdapter", + "CentralizedComparatorParseError", + "CollaborationModeMixin", + "TaggedCentralizedComparatorAdapter", + "joint_policy_gradient_loss", + "joint_sequence_log_prob", + "validate_collaboration_mode", +] diff --git a/comlrl/trainers/preference/centralized.py b/comlrl/trainers/centralized/adapters.py similarity index 100% rename from comlrl/trainers/preference/centralized.py rename to comlrl/trainers/centralized/adapters.py diff --git a/comlrl/trainers/preference/collaboration.py b/comlrl/trainers/centralized/collaboration.py similarity index 99% rename from comlrl/trainers/preference/collaboration.py rename to comlrl/trainers/centralized/collaboration.py index 43d9855..feb9823 100644 --- a/comlrl/trainers/preference/collaboration.py +++ b/comlrl/trainers/centralized/collaboration.py @@ -11,7 +11,7 @@ from comlrl.utils.reward_utils import call_reward_function, resolve_reward_range from comlrl.utils.tokenizer_utils import apply_tokenizer_specials -from .centralized import CentralizedComparatorParseError +from .adapters import CentralizedComparatorParseError def joint_sequence_log_prob(agent, tokenizer, prompt_ids, completion_ids): diff --git a/comlrl/trainers/centralized/mode.py b/comlrl/trainers/centralized/mode.py new file mode 100644 index 0000000..ec447d3 --- /dev/null +++ b/comlrl/trainers/centralized/mode.py @@ -0,0 +1,88 @@ +"""Collaboration modes shared by every trainer family. + +``decentralized`` trains one actor per agent. ``centralized`` trains one joint actor that +answers for all ``num_agents`` roles: an adapter builds one prompt from the role prompts and +splits the joint response back into role outputs, which the task reward and the evaluation +logger see unchanged. Trainers opt in through :class:`CollaborationModeMixin`. +""" + +from typing import Any, Callable, Optional, Sequence, Tuple + +from .collaboration import CentralizedCollaboration + +COLLABORATION_MODES = ("decentralized", "centralized") + + +def validate_collaboration_mode(mode: str) -> None: + if mode not in COLLABORATION_MODES: + raise ValueError("collaboration_mode must be decentralized or centralized.") + + +class CollaborationModeMixin: + """Hooks a trainer calls to support ``args.collaboration_mode="centralized"``. + + * ``_setup_collaboration`` in ``__init__``: wraps the role formatters and the task + reward around one joint actor and returns the actor-side arguments. + * ``_role_completions`` before logging evaluation metrics: splits each joint + completion into role outputs. + * ``is_centralized`` wherever the joint actor needs its own objective. + """ + + centralized_collaboration: Optional[CentralizedCollaboration] = None + + @property + def is_centralized(self) -> bool: + return self.centralized_collaboration is not None + + def _setup_collaboration( + self, + adapter: Any, + *, + agents: Optional[Sequence[Any]], + tokenizer: Any, + reward_func: Optional[Callable], + formatters: Any, + num_agents: int, + ) -> Tuple[Any, Any, int]: + """Return ``(reward_func, formatters, num_actors)`` for the configured mode.""" + if getattr(self.args, "collaboration_mode", "decentralized") != "centralized": + return reward_func, formatters, num_agents + if getattr(self.args, "num_turns", 1) != 1: + raise ValueError( + "Centralized collaboration currently supports num_turns=1 only." + ) + if num_agents != self.args.num_agents: + raise ValueError("num_agents must match args.num_agents (the role count).") + if agents is not None and len(agents) != 1: + raise ValueError( + "Centralized collaboration requires exactly one actor model." + ) + if isinstance(tokenizer, (list, tuple)) and len(tokenizer) != 1: + raise ValueError("Centralized collaboration requires one actor tokenizer.") + if not callable(reward_func): + raise ValueError("reward_func must be a callable.") + for name in ("build_prompt", "parse_completion"): + if not callable(getattr(adapter, name, None)): + raise TypeError(f"centralized_adapter.{name} must be callable.") + self.centralized_collaboration = CentralizedCollaboration( + adapter, formatters, reward_func, num_agents + ) + return ( + self.centralized_collaboration, + self.centralized_collaboration.build_prompt, + 1, + ) + + def _role_metrics_callback(self, callback): + """Give a rollout-metrics callback per-role samples (no-op when decentralized).""" + if callback is None or self.centralized_collaboration is None: + return callback + return self.centralized_collaboration.wrap_metrics_callback(callback) + + def _role_completions(self, agent_completions_turns, batch_items): + """Split joint completions into per-role completions (no-op when decentralized).""" + if self.centralized_collaboration is None: + return agent_completions_turns + return self.centralized_collaboration.split_eval( + agent_completions_turns, batch_items + ) diff --git a/comlrl/trainers/preference/__init__.py b/comlrl/trainers/pbrl/__init__.py similarity index 63% rename from comlrl/trainers/preference/__init__.py rename to comlrl/trainers/pbrl/__init__.py index 7bb797b..2ecd945 100644 --- a/comlrl/trainers/preference/__init__.py +++ b/comlrl/trainers/pbrl/__init__.py @@ -1,9 +1,4 @@ from .madpo import MADPOConfig, MADPOTrainer -from .centralized import ( - CentralizedComparatorAdapter, - CentralizedComparatorParseError, - TaggedCentralizedComparatorAdapter, -) from .iterative import ( MADPOIterConfig, MADPOIterTrainer, @@ -15,9 +10,6 @@ __all__ = [ "MADPOConfig", "MADPOTrainer", - "CentralizedComparatorAdapter", - "CentralizedComparatorParseError", - "TaggedCentralizedComparatorAdapter", "MADPOIterConfig", "MADPOIterTrainer", "MARLHFConfig", diff --git a/comlrl/trainers/preference/checkpointing.py b/comlrl/trainers/pbrl/checkpoint_extras.py similarity index 89% rename from comlrl/trainers/preference/checkpointing.py rename to comlrl/trainers/pbrl/checkpoint_extras.py index dc6e6bc..b863d0e 100644 --- a/comlrl/trainers/preference/checkpointing.py +++ b/comlrl/trainers/pbrl/checkpoint_extras.py @@ -11,7 +11,9 @@ import threading import torch +from transformers import AutoTokenizer +from comlrl.utils.checkpoint import capture_trainer_state from comlrl.utils.distributed import unwrap_model from .output_paths import resolve_output_path @@ -95,7 +97,7 @@ def validate_preference_progress(trainer, state, kind): if type(progress[name]) is not int or progress[name] < 0: raise ValueError(f"Invalid preference checkpoint counter: {name}") if kind == "madpo": - from .madpo import PreferencePair + from .madpo import PreferencePair # madpo imports this module pairs = progress["preference_pairs"] if not isinstance(pairs, list) or not all( @@ -146,8 +148,6 @@ def preference_extra_state(trainer): state[name] = getattr(trainer, name) child = getattr(trainer, "_checkpoint_ac_trainer", None) if child is not None: - from comlrl.utils.checkpoint import capture_trainer_state - state["actor_critic"] = capture_trainer_state(child) elif getattr(trainer, "_checkpoint_pending_ac_state", None) is not None: state["actor_critic"] = trainer._checkpoint_pending_ac_state @@ -560,7 +560,7 @@ def _copy_restored_asset(source, destination): def load_preference_assets(trainer, checkpoint_dir, manifest): if manifest.get("iterative"): - from .iterative import PreferenceReplayShard + from .iterative import PreferenceReplayShard # iterative imports this module destinations = _preference_resume_destinations(trainer, checkpoint_dir) allocated = set() @@ -596,8 +596,6 @@ def load_preference_assets(trainer, checkpoint_dir, manifest): setattr(trainer, name, str(path)) if manifest.get("reward_tokenizer"): - from transformers import AutoTokenizer - trainer.reward_tokenizer = AutoTokenizer.from_pretrained( str(_asset_path(checkpoint_dir, manifest["reward_tokenizer"])) ) @@ -643,15 +641,8 @@ def train_iterations_checkpointed(trainer, *, reward_model, **kwargs): if progress.get("complete") and start == total_iterations: return trainer._checkpoint_init_wandb() - if not reward_model: - for index, agent in enumerate(trainer.agents): - agent.to(trainer.agent_devices[index]) - agent.train() - trainer._maybe_save_initial_policy_checkpoint() - updates_seen = progress.get("updates_seen", 0) - total_pairs = progress.get("total_pairs", 0) - - def commit(next_iteration): + + def commit(next_iteration, updates_seen, total_pairs): complete = next_iteration == total_iterations trainer._checkpoint_commit( { @@ -664,60 +655,11 @@ def commit(next_iteration): force=complete, ) - for iteration in range(start, total_iterations): - if reward_model: - trainer._reward_model_active = False - trainer._evaluating_with_task_reward = False - if trainer.args.preference_scoring_reward != "reward_model": - trainer._clear_reward_model() - trainer._reset_iteration_reward_distribution() - trainer._prepare_iteration_current_copy_comparator() - try: - pairs = trainer._build_preference_dataset(iteration_idx=iteration, **kwargs) - finally: - trainer._clear_iteration_current_copy_comparator() - if reward_model: - trainer._clear_iteration_model_comparator() - current_pair_count = len(pairs) - train_pairs = trainer._select_iteration_preference_pairs( - pairs, iteration_idx=iteration - ) - total_pairs += len(train_pairs) - trainer._log_iteration_replay( - iteration, - train_pairs=train_pairs, - current_pair_count=current_pair_count, - train_pair_count=len(train_pairs), - ) - pairs.clear() - if not train_pairs: - if trainer.verbose: - print( - f"{name}: no replay preference pairs were available " - f"for iteration {iteration + 1} " - f"(current generated {current_pair_count})." - ) - trainer._maybe_save_iteration_policy_checkpoint(iteration) - commit(iteration + 1) - continue - if reward_model: - trainer._clear_reward_model() - trainer._init_reward_model() - trainer._train_reward_model(train_pairs) - trainer._reward_model_active = True - trainer._train_online_rl(**kwargs) - else: - updates_seen = trainer._train_preference_pairs( - train_pairs, - updates_seen=updates_seen, - desc=f"{name} iteration {iteration + 1}/{total_iterations}", - ) - trainer._maybe_save_iteration_policy_checkpoint(iteration) - if reward_model and trainer.args.preference_scoring_reward != "reward_model": - trainer._clear_reward_model() - commit(iteration + 1) - - if start == total_iterations: - commit(start) - if total_pairs == 0 and trainer.verbose: - print(f"{name}: no non-tied preference pairs were generated.") + trainer._train_iterations( + kwargs, + reward_model=reward_model, + start=start, + updates_seen=progress.get("updates_seen", 0), + total_pairs=progress.get("total_pairs", 0), + commit=commit, + ) diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/pbrl/iterative.py similarity index 94% rename from comlrl/trainers/preference/iterative.py rename to comlrl/trainers/pbrl/iterative.py index 09f78b1..861d40b 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -2,7 +2,6 @@ import gc import inspect import json -import numbers import os import random import threading @@ -18,6 +17,7 @@ from transformers import AutoModelForCausalLM from comlrl.schedulers import DeviceScheduler +from comlrl.utils.checkpoint import register_checkpoint_record from comlrl.utils.distributed import unwrap_model from comlrl.utils.model_loading import ( freeze_model, @@ -33,7 +33,8 @@ MADPOTrainer, PreferencePair, ) -from .centralized import CentralizedComparatorAdapter +from ..centralized import CentralizedComparatorAdapter +from .checkpoint_extras import train_iterations_checkpointed from .marlhf import MARLHFConfig, MARLHFTrainer from .output_paths import resolve_output_path @@ -120,6 +121,14 @@ def _validate_iterative_config(args: Any) -> None: args.comparator_centralized_agent_index = int( args.comparator_centralized_agent_index ) + if args.comparator_centralized_max_new_tokens is not None: + if int(args.comparator_centralized_max_new_tokens) < 1: + raise ValueError( + "comparator_centralized_max_new_tokens must be >= 1 or null." + ) + args.comparator_centralized_max_new_tokens = int( + args.comparator_centralized_max_new_tokens + ) if args.comparator_num_candidates is not None: if int(args.comparator_num_candidates) < 1: raise ValueError("comparator_num_candidates must be >= 1 or null.") @@ -274,6 +283,7 @@ class _IterativePreferenceConfig: comparator_policy: str = "current" comparator_generation_mode: str = "decentralized" comparator_centralized_agent_index: int = 0 + comparator_centralized_max_new_tokens: Optional[int] = None comparator_model_name: Optional[str] = None comparator_agents: Optional[Sequence[str]] = None comparator_devices: Optional[Union[str, Sequence[str]]] = None @@ -328,6 +338,7 @@ def __post_init__(self) -> None: ) +@register_checkpoint_record @dataclass class PreferenceReplayShard: iteration: int @@ -346,11 +357,7 @@ class MADPOIterTrainer(MADPOTrainer): default_config_cls = MADPOIterConfig algorithm_name = "MADPOIter" - - def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress - - validate_preference_progress(self, state, "madpo_iter") + _checkpoint_progress_kind = "madpo_iter" def __init__( self, @@ -611,52 +618,69 @@ def train( checkpoint_total_limit=None, **kwargs, ): - if ( - isinstance(checkpoint_interval, bool) - or not isinstance(checkpoint_interval, numbers.Integral) - or checkpoint_interval < 1 - ): - raise ValueError("checkpoint_interval must be a positive integer.") - if ( - resume_from_checkpoint is not None - or checkpoint_dir is not None - or checkpoint_interval != 1 - or checkpoint_total_limit is not None - ): - from .checkpointing import train_iterations_checkpointed - - with self._checkpoint_context( - resume_from_checkpoint=resume_from_checkpoint, - checkpoint_dir=checkpoint_dir, - checkpoint_interval=checkpoint_interval, - checkpoint_total_limit=checkpoint_total_limit, - ): - return train_iterations_checkpointed(self, reward_model=False, **kwargs) + checkpointing = self._checkpoint_arguments( + resume_from_checkpoint, + checkpoint_dir, + checkpoint_interval, + checkpoint_total_limit, + ) + reward_model = isinstance(self, MARLHFTrainer) + if checkpointing is not None: + with self._checkpoint_context(**checkpointing): + return train_iterations_checkpointed( + self, reward_model=reward_model, **kwargs + ) - self._checkpoint_untracked_preference = True + self._checkpoint_untracked = True + name = "MARLHFIter" if reward_model else "MADPOIter" if int(self.args.num_turns) != 1: - raise ValueError("MADPOIter currently supports num_turns=1 only.") + raise ValueError(f"{name} currently supports num_turns=1 only.") if self.wandb_config is not None and not self.wandb_initialized: self._init_wandb() + self._train_iterations(kwargs, reward_model=reward_model) - for agent_idx, agent in enumerate(self.agents): - agent.to(self.agent_devices[agent_idx]) - agent.train() + def _train_iterations( + self, + generation_kwargs, + *, + reward_model, + start=0, + updates_seen=0, + total_pairs=0, + commit=None, + ): + """Refresh preferences and train, from iteration ``start`` to the last. + + MARLHFIter (``reward_model``) retrains its reward model and then runs + online RL; MADPOIter runs DPO on the selected pairs. ``commit`` receives + the next iteration and both counters after every iteration. + """ + name = "MARLHFIter" if reward_model else "MADPOIter" + total_iterations = int(self.args.num_iterations) + if not reward_model: + for agent_idx, agent in enumerate(self.agents): + agent.to(self.agent_devices[agent_idx]) + agent.train() self._maybe_save_initial_policy_checkpoint() - updates_seen = 0 - total_pairs = 0 - for iteration_idx in range(int(self.args.num_iterations)): + for iteration_idx in range(start, total_iterations): + if reward_model: + self._reward_model_active = False + self._evaluating_with_task_reward = False + if self.args.preference_scoring_reward != "reward_model": + self._clear_reward_model() self._reset_iteration_reward_distribution() self._prepare_iteration_current_copy_comparator() try: preference_pairs = self._build_preference_dataset( iteration_idx=iteration_idx, - **kwargs, + **generation_kwargs, ) finally: self._clear_iteration_current_copy_comparator() + if reward_model: + self._clear_iteration_model_comparator() current_pair_count = len(preference_pairs) train_pairs = self._select_iteration_preference_pairs( preference_pairs, @@ -673,25 +697,34 @@ def train( if not train_pairs: if self.verbose: print( - "MADPOIter: no replay preference pairs were available " + f"{name}: no replay preference pairs were available " f"for iteration {iteration_idx + 1} " f"(current generated {current_pair_count})." ) self._maybe_save_iteration_policy_checkpoint(iteration_idx) - continue - - updates_seen = self._train_preference_pairs( - train_pairs, - updates_seen=updates_seen, - desc=( - f"MADPOIter iteration {iteration_idx + 1}/" - f"{int(self.args.num_iterations)}" - ), - ) - self._maybe_save_iteration_policy_checkpoint(iteration_idx) + elif reward_model: + self._clear_reward_model() + self._init_reward_model() + self._train_reward_model(train_pairs) + self._reward_model_active = True + self._train_online_rl(**generation_kwargs) + self._maybe_save_iteration_policy_checkpoint(iteration_idx) + if self.args.preference_scoring_reward != "reward_model": + self._clear_reward_model() + else: + updates_seen = self._train_preference_pairs( + train_pairs, + updates_seen=updates_seen, + desc=f"{name} iteration {iteration_idx + 1}/{total_iterations}", + ) + self._maybe_save_iteration_policy_checkpoint(iteration_idx) + if commit is not None: + commit(iteration_idx + 1, updates_seen, total_pairs) + if commit is not None and start == total_iterations: + commit(start, updates_seen, total_pairs) if total_pairs == 0 and self.verbose: - print("MADPOIter: no non-tied preference pairs were generated.") + print(f"{name}: no non-tied preference pairs were generated.") def _select_iteration_preference_pairs( self, @@ -1333,17 +1366,6 @@ def _preference_pair_to_record(pair: PreferencePair) -> Dict[str, Any]: ), } - def _load_replay_shard( - self, - shard: Optional[PreferenceReplayShard], - ) -> List[PreferencePair]: - if shard is None: - return [] - return [ - self._preference_pair_from_record(record) - for record in self._load_replay_records(shard) - ] - @staticmethod def _load_replay_records(shard: PreferenceReplayShard) -> List[Dict[str, Any]]: with open(shard.path, "r", encoding="utf-8") as handle: @@ -1900,7 +1922,11 @@ def _generate_centralized_policy_output_for_agent( [batch_item], agent_idx=agent_idx, num_return_sequences=num_candidates, - max_new_tokens=self.args.max_new_tokens, + # One joint answer holds every agent's output. + max_new_tokens=( + self.args.comparator_centralized_max_new_tokens + or self.args.max_new_tokens + ), prompts_override=[prompt], **kwargs, ) @@ -2016,7 +2042,7 @@ def _split_centralized_comparator_outputs( batch_item: Dict[str, Any], prompt: str, ) -> List[Dict[str, Any]]: - if getattr(self, "centralized_collaboration", None) is not None: + if self.is_centralized: if not completions: raise ValueError("Centralized comparator produced no candidates.") # Preserve the exact joint text, including tags, for preference learning. @@ -2858,11 +2884,7 @@ class MARLHFIterTrainer(MADPOIterTrainer, MARLHFTrainer): default_config_cls = MARLHFIterConfig algorithm_name = "MARLHFIter" - - def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress - - validate_preference_progress(self, state, "marlhf_iter") + _checkpoint_progress_kind = "marlhf_iter" def _preference_scoring_raw_and_processed_rewards( self, @@ -2885,93 +2907,3 @@ def _preference_scoring_raw_and_processed_rewards( batch_items=batch_items, ) return list(scores), list(scores) - - def train( - self, - resume_from_checkpoint=None, - checkpoint_dir=None, - checkpoint_interval=1, - checkpoint_total_limit=None, - **kwargs, - ): - if ( - isinstance(checkpoint_interval, bool) - or not isinstance(checkpoint_interval, numbers.Integral) - or checkpoint_interval < 1 - ): - raise ValueError("checkpoint_interval must be a positive integer.") - if ( - resume_from_checkpoint is not None - or checkpoint_dir is not None - or checkpoint_interval != 1 - or checkpoint_total_limit is not None - ): - from .checkpointing import train_iterations_checkpointed - - with self._checkpoint_context( - resume_from_checkpoint=resume_from_checkpoint, - checkpoint_dir=checkpoint_dir, - checkpoint_interval=checkpoint_interval, - checkpoint_total_limit=checkpoint_total_limit, - ): - return train_iterations_checkpointed(self, reward_model=True, **kwargs) - - self._checkpoint_untracked_preference = True - if int(self.args.num_turns) != 1: - raise ValueError("MARLHFIter currently supports num_turns=1 only.") - - if self.wandb_config is not None and not self.wandb_initialized: - self._init_wandb() - - self._maybe_save_initial_policy_checkpoint() - total_pairs = 0 - for iteration_idx in range(int(self.args.num_iterations)): - self._reward_model_active = False - self._evaluating_with_task_reward = False - if self.args.preference_scoring_reward != "reward_model": - self._clear_reward_model() - self._reset_iteration_reward_distribution() - self._prepare_iteration_current_copy_comparator() - try: - preference_pairs = self._build_preference_dataset( - iteration_idx=iteration_idx, - **kwargs, - ) - finally: - self._clear_iteration_current_copy_comparator() - self._clear_iteration_model_comparator() - current_pair_count = len(preference_pairs) - train_pairs = self._select_iteration_preference_pairs( - preference_pairs, - iteration_idx=iteration_idx, - ) - total_pairs += len(train_pairs) - self._log_iteration_replay( - iteration_idx, - train_pairs=train_pairs, - current_pair_count=current_pair_count, - train_pair_count=len(train_pairs), - ) - preference_pairs.clear() - if not train_pairs: - if self.verbose: - print( - "MARLHFIter: no replay preference pairs were available " - f"for iteration {iteration_idx + 1} " - f"(current generated {current_pair_count})." - ) - self._maybe_save_iteration_policy_checkpoint(iteration_idx) - continue - - self._clear_reward_model() - self._init_reward_model() - self._train_reward_model(train_pairs) - self._reward_model_active = True - - self._train_online_rl(**kwargs) - self._maybe_save_iteration_policy_checkpoint(iteration_idx) - if self.args.preference_scoring_reward != "reward_model": - self._clear_reward_model() - - if total_pairs == 0 and self.verbose: - print("MARLHFIter: no non-tied preference pairs were generated.") diff --git a/comlrl/trainers/preference/madpo.py b/comlrl/trainers/pbrl/madpo.py similarity index 84% rename from comlrl/trainers/preference/madpo.py rename to comlrl/trainers/pbrl/madpo.py index cee83a2..054ac7c 100644 --- a/comlrl/trainers/preference/madpo.py +++ b/comlrl/trainers/pbrl/madpo.py @@ -1,4 +1,3 @@ -import numbers import random from dataclasses import dataclass from typing import Any, Dict, List, Optional, Sequence, Tuple, Union @@ -10,18 +9,31 @@ from tqdm import tqdm # type: ignore from comlrl.schedulers.device_scheduler import normalize_parallel_training +from comlrl.utils.checkpoint import register_checkpoint_record from comlrl.utils.distributed import unwrap_model from comlrl.utils.reference_kl import validate_reference_kl_config from comlrl.utils.tokenizer_utils import apply_tokenizer_specials from ..reinforce.magrpo import MAGRPOTrainer -from .centralized import ( +from ..centralized import ( CentralizedComparatorAdapter, TaggedCentralizedComparatorAdapter, + joint_sequence_log_prob, + validate_collaboration_mode, +) +from .checkpoint_extras import ( + load_preference_assets, + preference_extra_state, + prune_preference_asset_pool, + restore_preference_extra_state, + save_preference_assets, + validate_preference_assets, + validate_preference_extra_state, + validate_preference_progress, ) -from .collaboration import CentralizedCollaboration, joint_sequence_log_prob +@register_checkpoint_record @dataclass class AgentPreferenceTensors: prompt_input_ids: torch.Tensor @@ -29,6 +41,7 @@ class AgentPreferenceTensors: loser_completion_ids: torch.Tensor +@register_checkpoint_record @dataclass class PreferencePair: prompts: List[str] @@ -87,8 +100,7 @@ class MADPOConfig: environment_steps_per_pair: int = 2 def __post_init__(self) -> None: - if self.collaboration_mode not in {"decentralized", "centralized"}: - raise ValueError("collaboration_mode must be decentralized or centralized.") + validate_collaboration_mode(self.collaboration_mode) if self.collaboration_mode == "centralized": if getattr(self, "comparator_centralized_agent_index", 0) != 0: raise ValueError("Centralized collaboration has only actor index 0.") @@ -191,29 +203,7 @@ def __init__( f"centralized_comparator_adapter.{method} must be callable." ) self.centralized_comparator_adapter = adapter - self.centralized_collaboration = None - self._centralized_eval_items = [] - if config.collaboration_mode == "centralized": - if agents is not None and len(agents) != 1: - raise ValueError( - "Centralized collaboration requires exactly one actor model." - ) - if isinstance(tokenizer, (list, tuple)) and len(tokenizer) != 1: - raise ValueError( - "Centralized collaboration requires one actor tokenizer." - ) - if not callable(reward_func): - raise ValueError("reward_func must be a callable.") - if num_agents != config.num_agents: - raise ValueError( - "num_agents must match args.num_agents (the role count)." - ) - self.centralized_collaboration = CentralizedCollaboration( - adapter, formatters, reward_func, num_agents - ) - formatters = self.centralized_collaboration.build_prompt - reward_func = self.centralized_collaboration - num_agents = 1 + # In centralized mode MAGRPOTrainer wraps one joint actor with this adapter. super().__init__( agent_model=agent_model, agents=agents, @@ -230,27 +220,9 @@ def __init__( eval_logger=eval_logger, eval_aggregator=eval_aggregator, args=config, + centralized_adapter=adapter, ) - def evaluate(self, num_eval_samples: Optional[int] = None) -> Dict[str, float]: - self._centralized_eval_items = [] - try: - return super().evaluate(num_eval_samples=num_eval_samples) - finally: - self._centralized_eval_items = [] - - def _evaluate_sample(self, batch_item, *args, **kwargs): - super()._evaluate_sample(batch_item, *args, **kwargs) - if self.centralized_collaboration is not None: - self._centralized_eval_items.append(batch_item) - - def _log_eval_metrics(self, all_agent_completions_turns, *args, **kwargs): - if self.centralized_collaboration is not None: - all_agent_completions_turns = self.centralized_collaboration.split_eval( - all_agent_completions_turns, self._centralized_eval_items - ) - return super()._log_eval_metrics(all_agent_completions_turns, *args, **kwargs) - def train( self, resume_from_checkpoint=None, @@ -259,27 +231,17 @@ def train( checkpoint_total_limit=None, **kwargs, ): - if ( - isinstance(checkpoint_interval, bool) - or not isinstance(checkpoint_interval, numbers.Integral) - or checkpoint_interval < 1 - ): - raise ValueError("checkpoint_interval must be a positive integer.") - if ( - resume_from_checkpoint is not None - or checkpoint_dir is not None - or checkpoint_interval != 1 - or checkpoint_total_limit is not None - ): - with self._checkpoint_context( - resume_from_checkpoint=resume_from_checkpoint, - checkpoint_dir=checkpoint_dir, - checkpoint_interval=checkpoint_interval, - checkpoint_total_limit=checkpoint_total_limit, - ): + checkpointing = self._checkpoint_arguments( + resume_from_checkpoint, + checkpoint_dir, + checkpoint_interval, + checkpoint_total_limit, + ) + if checkpointing is not None: + with self._checkpoint_context(**checkpointing): return self._train_madpo_checkpointed(**kwargs) - self._checkpoint_untracked_preference = True + self._checkpoint_untracked = True if int(self.args.num_turns) != 1: raise ValueError("MADPO currently supports num_turns=1 only.") @@ -382,40 +344,20 @@ def commit_epoch(next_epoch, updates_seen, *, force=False): if not generated and start_epoch == int(self.args.num_train_epochs): commit_epoch(start_epoch, updates_seen) - def _checkpoint_extra_state(self): - from .checkpointing import preference_extra_state - - return preference_extra_state(self) - - def _restore_checkpoint_extra_state(self, state, checkpoint_dir): - from .checkpointing import restore_preference_extra_state - - restore_preference_extra_state(self, state, checkpoint_dir) - - def _checkpoint_save_assets(self, checkpoint_dir): - from .checkpointing import save_preference_assets - - return save_preference_assets(self, checkpoint_dir) - - def _checkpoint_load_assets(self, checkpoint_dir, assets): - from .checkpointing import load_preference_assets - - load_preference_assets(self, checkpoint_dir, assets) - - def _checkpoint_validate_assets(self, checkpoint_dir, assets): - from .checkpointing import validate_preference_assets - - validate_preference_assets(self, checkpoint_dir, assets) + # Checkpoint hooks looked up by comlrl.utils.checkpoint (see checkpoint_extras). + _checkpoint_progress_kind = "madpo" + _checkpoint_extra_state = preference_extra_state + _checkpoint_restore_extra_state = restore_preference_extra_state + _checkpoint_save_assets = save_preference_assets + _checkpoint_load_assets = load_preference_assets + _checkpoint_prune_assets = staticmethod(prune_preference_asset_pool) + _checkpoint_validate_assets = validate_preference_assets def _checkpoint_validate_extra_state(self, state, checkpoint_dir): - from .checkpointing import validate_preference_extra_state - validate_preference_extra_state(state) def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress - - validate_preference_progress(self, state, "madpo") + validate_preference_progress(self, state, self._checkpoint_progress_kind) def _train_preference_pairs( self, diff --git a/comlrl/trainers/preference/marlhf.py b/comlrl/trainers/pbrl/marlhf.py similarity index 95% rename from comlrl/trainers/preference/marlhf.py rename to comlrl/trainers/pbrl/marlhf.py index abebc32..cb7fb29 100644 --- a/comlrl/trainers/preference/marlhf.py +++ b/comlrl/trainers/pbrl/marlhf.py @@ -1,6 +1,5 @@ import copy import gc -import numbers from dataclasses import dataclass from functools import partial from typing import Any, Dict, List, Optional, Sequence, Union @@ -11,19 +10,31 @@ from torch.utils.data import DataLoader from torch.overrides import TorchFunctionMode from tqdm import tqdm # type: ignore -from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase - +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoTokenizer, + GenerationConfig, + PreTrainedTokenizerBase, +) + +from comlrl.schedulers import DeviceScheduler +from comlrl.schedulers.device_scheduler import normalize_parallel_training +from comlrl.utils.checkpoint import ( + _validate_state, + capture_rng_state, + restore_rng_state, + restore_trainer_state, +) +from comlrl.utils.model_loading import resolve_model_kwargs from comlrl.utils.tokenizer_utils import ensure_pad_token -from .collaboration import joint_policy_gradient_loss +from ..centralized import joint_policy_gradient_loss from .madpo import MADPOConfig, MADPOTrainer, PreferencePair from ..actor_critic import IACConfig, IACTrainer, MAACConfig, MAACTrainer from ..reinforce.magrpo import MAGRPOTrainer -_RL_ALGORITHM_ALIASES = { - "mareforce": "mareinforce", -} _MAGRPO_ADVANTAGE_MODES = { "magrpo": "mean", "mareinforce": "raw", @@ -34,8 +45,7 @@ def _normalize_rl_algorithm(name: Optional[str]) -> str: - algorithm = str(name or "magrpo").strip().lower() - return _RL_ALGORITHM_ALIASES.get(algorithm, algorithm) + return str(name or "magrpo").strip().lower() def _apply_magrpo_family_args(args: Any) -> None: @@ -166,6 +176,7 @@ class MARLHFTrainer(MADPOTrainer): default_config_cls = MARLHFConfig algorithm_name = "MARLHF" + _checkpoint_progress_kind = "marlhf" def __init__( self, @@ -205,27 +216,17 @@ def train( checkpoint_total_limit=None, **kwargs, ): - if ( - isinstance(checkpoint_interval, bool) - or not isinstance(checkpoint_interval, numbers.Integral) - or checkpoint_interval < 1 - ): - raise ValueError("checkpoint_interval must be a positive integer.") - if ( - resume_from_checkpoint is not None - or checkpoint_dir is not None - or checkpoint_interval != 1 - or checkpoint_total_limit is not None - ): - with self._checkpoint_context( - resume_from_checkpoint=resume_from_checkpoint, - checkpoint_dir=checkpoint_dir, - checkpoint_interval=checkpoint_interval, - checkpoint_total_limit=checkpoint_total_limit, - ): + checkpointing = self._checkpoint_arguments( + resume_from_checkpoint, + checkpoint_dir, + checkpoint_interval, + checkpoint_total_limit, + ) + if checkpointing is not None: + with self._checkpoint_context(**checkpointing): return self._train_marlhf_checkpointed(**kwargs) - self._checkpoint_untracked_preference = True + self._checkpoint_untracked = True if int(self.args.num_turns) != 1: raise ValueError("MARLHF currently supports num_turns=1 only.") @@ -240,11 +241,6 @@ def train( self._reward_model_active = True self._train_online_rl(**kwargs) - def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress - - validate_preference_progress(self, state, "marlhf") - def _checkpoint_validate_extra_state(self, state, checkpoint_dir): super()._checkpoint_validate_extra_state(state, checkpoint_dir) child_state = state.get("actor_critic") @@ -253,12 +249,6 @@ def _checkpoint_validate_extra_state(self, state, checkpoint_dir): if self.args.rl_algorithm not in _ACTOR_CRITIC_ALGORITHMS: raise ValueError("Unexpected MARLHF actor-critic checkpoint state.") - from comlrl.utils.checkpoint import ( - _validate_state, - capture_rng_state, - restore_rng_state, - ) - # AC constructors can move supplied models and attach heads. A meta # probe checks their structure without allocating another set of LLM # weights on a training device or changing the live parent. @@ -371,10 +361,8 @@ def __torch_function__(self, func, types, args=(), kwargs=None): def _checkpoint_prepare_meta_critics(self): """Construct identifier-based critics from configuration, without weights.""" - from transformers import AutoConfig, GenerationConfig from ..actor_critic.iac import IACTrainer as BuiltinIACTrainer from ..actor_critic.maac import MAACTrainer as BuiltinMAACTrainer - from comlrl.utils.model_loading import resolve_model_kwargs # Replacement factories keep their original inputs and CPU fallback. # The built-in factory accepts these same preloaded model objects. @@ -410,9 +398,6 @@ def _checkpoint_prepare_meta_critics(self): def _checkpoint_actor_critic_device_plan(self, child, args): """Resolve the online child's real placement, not its probe devices.""" - from comlrl.schedulers import DeviceScheduler - from comlrl.schedulers.device_scheduler import normalize_parallel_training - if normalize_parallel_training(args.parallel_training) == "mp": if self.args.rl_algorithm == "maac": agent_devices = DeviceScheduler.resolve_devices( @@ -595,8 +580,6 @@ def _train_actor_critic_rl( and getattr(self, "_checkpoint_pending_ac_state", None) is not None ) if restoring_child: - from comlrl.utils.checkpoint import capture_rng_state, restore_rng_state - rng_before_init = capture_rng_state() ac_trainer = self._create_actor_critic_trainer() if _checkpoint_commit is None: @@ -604,8 +587,6 @@ def _train_actor_critic_rl( else: self._checkpoint_ac_trainer = ac_trainer if restoring_child: - from comlrl.utils.checkpoint import restore_trainer_state - restore_trainer_state( ac_trainer, self._checkpoint_pending_ac_state, diff --git a/comlrl/trainers/preference/output_paths.py b/comlrl/trainers/pbrl/output_paths.py similarity index 100% rename from comlrl/trainers/preference/output_paths.py rename to comlrl/trainers/pbrl/output_paths.py diff --git a/comlrl/trainers/reinforce/__init__.py b/comlrl/trainers/reinforce/__init__.py index 2ad3157..ff1cb8b 100644 --- a/comlrl/trainers/reinforce/__init__.py +++ b/comlrl/trainers/reinforce/__init__.py @@ -1,5 +1,4 @@ from .magrpo import MAGRPOConfig, MAGRPOTrainer -from .centralized_magrpo import CentralizedMAGRPOConfig, CentralizedMAGRPOTrainer from .mareinforce import MAREINFORCEConfig, MAREINFORCETrainer from .maremax import MAReMaxConfig, MAReMaxTrainer from .marloo import MARLOOConfig, MARLOOTrainer @@ -7,8 +6,6 @@ __all__ = [ "MAGRPOConfig", "MAGRPOTrainer", - "CentralizedMAGRPOConfig", - "CentralizedMAGRPOTrainer", "MAREINFORCEConfig", "MAREINFORCETrainer", "MAReMaxConfig", diff --git a/comlrl/trainers/reinforce/centralized_magrpo.py b/comlrl/trainers/reinforce/centralized_magrpo.py deleted file mode 100644 index 69458b9..0000000 --- a/comlrl/trainers/reinforce/centralized_magrpo.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Single joint-input/joint-output actor optimized with task-reward MAGRPO.""" - -from dataclasses import dataclass - -from .magrpo import MAGRPOConfig, MAGRPOTrainer - - -@dataclass -class CentralizedMAGRPOConfig(MAGRPOConfig): - num_turns: int = 1 - collaboration_mode: str = "centralized" - - def __post_init__(self): - super().__post_init__() - if self.collaboration_mode != "centralized": - raise ValueError( - "CentralizedMAGRPOConfig requires centralized collaboration." - ) - if self.num_turns != 1: - raise ValueError("Centralized MAGRPO currently supports num_turns=1 only.") - - -class CentralizedMAGRPOTrainer(MAGRPOTrainer): - """Train one joint policy; split role responses only for rewards and evaluation.""" - - default_config_cls = CentralizedMAGRPOConfig - - def __init__( - self, - *, - centralized_adapter, - num_agents=2, - agents=None, - tokenizer=None, - reward_func=None, - formatters=None, - args=None, - **kwargs, - ): - # Loaded lazily because preference trainers also use the MAGRPO base. - from ..preference.collaboration import CentralizedCollaboration - - config = args if args is not None else self.default_config_cls() - if getattr(config, "collaboration_mode", None) != "centralized": - raise ValueError( - "CentralizedMAGRPOTrainer requires centralized configuration." - ) - if config.num_turns != 1: - raise ValueError("Centralized MAGRPO currently supports num_turns=1 only.") - if num_agents != config.num_agents: - raise ValueError("num_agents must match args.num_agents (the role count).") - if agents is not None and len(agents) != 1: - raise ValueError( - "Centralized collaboration requires exactly one actor model." - ) - if isinstance(tokenizer, (list, tuple)) and len(tokenizer) != 1: - raise ValueError("Centralized collaboration requires one actor tokenizer.") - if not callable(reward_func): - raise ValueError("reward_func must be callable.") - for name in ("build_prompt", "parse_completion"): - if not callable(getattr(centralized_adapter, name, None)): - raise TypeError(f"centralized_adapter.{name} must be callable.") - self.centralized_collaboration = CentralizedCollaboration( - centralized_adapter, formatters, reward_func, num_agents - ) - self._centralized_eval_items = [] - super().__init__( - agents=agents, - tokenizer=tokenizer, - num_agents=1, - args=config, - reward_func=self.centralized_collaboration, - formatters=self.centralized_collaboration.build_prompt, - **kwargs, - ) - - def evaluate(self, num_eval_samples=None): - self._centralized_eval_items = [] - try: - return super().evaluate(num_eval_samples=num_eval_samples) - finally: - self._centralized_eval_items = [] - - def _evaluate_sample(self, batch_item, *args, **kwargs): - super()._evaluate_sample(batch_item, *args, **kwargs) - self._centralized_eval_items.append(batch_item) - - def _log_eval_metrics(self, all_agent_completions_turns, *args, **kwargs): - role_outputs = self.centralized_collaboration.split_eval( - all_agent_completions_turns, self._centralized_eval_items - ) - return super()._log_eval_metrics(role_outputs, *args, **kwargs) - - def _compute_loss_with_gradients(self, agent, completions_data, returns): - from ..preference.collaboration import joint_policy_gradient_loss - - return joint_policy_gradient_loss(self, agent, completions_data, returns) diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index 7c695ee..ee361d9 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -1,5 +1,4 @@ import inspect -import numbers import os import random from dataclasses import dataclass @@ -17,8 +16,8 @@ from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training -from comlrl.utils.checkpoint import TrainingCheckpointMixin -from comlrl.utils.distributed import local_context, unwrap_model +from comlrl.utils.checkpoint import TrainingCheckpointMixin, register_checkpoint_record +from comlrl.utils.distributed import unwrap_model from comlrl.utils.formatters import build_formatters from comlrl.utils.model_loading import ( infer_model_name, @@ -26,12 +25,11 @@ resolve_model_sources, ) from comlrl.utils.reference_kl import ( - _sequence_kl_values, clone_reference_models, load_reference_models_from_sources, reference_kl_coef, reference_kl_enabled, - reference_kl_for_sequence, + reference_kl_values, resolve_reference_devices, validate_reference_kl_config, ) @@ -42,6 +40,12 @@ ) from comlrl.utils.wandb_utils import verbose_from_wandb_config, wandb_init_kwargs +from ..centralized import ( + CollaborationModeMixin, + joint_policy_gradient_loss, + validate_collaboration_mode, +) + @dataclass class MAGRPOConfig: @@ -76,6 +80,16 @@ class MAGRPOConfig: train_batch_size: Optional[int] = None advantage_normalization: bool = True advantage_mode: str = "mean" + # Without a gradient (every advantage zero), AdamW would still move weights + # along its stored momentum; skipping keeps such updates from changing them. + skip_zero_grad_steps: bool = False + # "centralized" trains one joint actor that answers for all num_agents roles; + # role outputs are split from its response only for rewards and evaluation. + collaboration_mode: str = "decentralized" + # Centralized only: score the joint response with the decentralized loss (the + # same token range, precision and non-finite fallback) instead of the + # joint-sequence loss, so both modes share one objective implementation. + match_decentralized_loss: bool = False reference_kl_enabled: bool = False reference_kl_coef: float = 0.1 reference_devices: Optional[Union[str, Sequence[str]]] = None @@ -99,6 +113,9 @@ def __post_init__(self) -> None: raise ValueError("eval_batch_size must be >= 1.") if self.num_turns < 1: raise ValueError("num_turns must be >= 1.") + validate_collaboration_mode(self.collaboration_mode) + if self.collaboration_mode == "centralized" and self.num_turns != 1: + raise ValueError("Centralized MAGRPO currently supports num_turns=1 only.") if self.logging_steps < 1: raise ValueError("logging_steps must be >= 1.") if self.train_batch_size is None: @@ -109,14 +126,11 @@ def __post_init__(self) -> None: if mode == "mp" and self.agent_devices is None: raise ValueError("parallel_training='mp' requires explicit agent_devices.") self.parallel_training = mode - actor_count = ( - 1 - if getattr(self, "collaboration_mode", None) == "centralized" - else self.num_agents - ) + actor_count = 1 if self.collaboration_mode == "centralized" else self.num_agents validate_reference_kl_config(self, actor_count) +@register_checkpoint_record @dataclass class NodeSample: agent_idx: int @@ -128,7 +142,17 @@ class NodeSample: node_env_step: int -class MAGRPOTrainer(TrainingCheckpointMixin): +def _has_gradient(optimizer) -> bool: + """Whether any parameter of the optimizer holds a nonzero gradient.""" + return any( + bool(param.grad.ne(0).any()) + for group in optimizer.param_groups + for param in group["params"] + if param.grad is not None + ) + + +class MAGRPOTrainer(CollaborationModeMixin, TrainingCheckpointMixin): """ Multi-Agent Group Relative Policy Optimization Trainer (MAGRPO). Supports both single-turn and multi-turn training with external transitions. @@ -154,6 +178,10 @@ class MAGRPOTrainer(TrainingCheckpointMixin): eval_aggregator: Called as ``eval_aggregator(metrics, num_turns=...)`` and returns ``turn_k/...`` metrics, logged under ``eval/`` external_transition: Function that provides external transitions between turns + centralized_adapter: With ``args.collaboration_mode="centralized"``, builds one + joint prompt from the role prompts and splits the joint response back into + ``num_agents`` role outputs (``build_prompt`` / ``parse_completion``). One + actor is then trained; rewards and evaluation see the split role outputs. """ default_config_cls: Type[Any] = MAGRPOConfig @@ -179,8 +207,18 @@ def __init__( eval_logger: Optional[Callable] = None, eval_aggregator: Optional[Callable] = None, args: Optional[Any] = None, + *, + centralized_adapter: Optional[Any] = None, ): self.args = args if args is not None else self.default_config_cls() + reward_func, formatters, num_agents = self._setup_collaboration( + centralized_adapter, + agents=agents, + tokenizer=tokenizer, + reward_func=reward_func, + formatters=formatters, + num_agents=num_agents, + ) self.parallel_training = normalize_parallel_training( getattr(self.args, "parallel_training", "none") ) @@ -245,7 +283,6 @@ def __init__( ) self.agent_devices = [single_device] * self.num_agents self.device = self.agent_devices[0] - self.dist_env = local_context(self.device) model_kwargs = resolve_model_kwargs(self.model_config) if actor_sources and all(isinstance(src, str) for src in actor_sources): from transformers import AutoModelForCausalLM @@ -337,7 +374,7 @@ def __init__( self.wandb_config = wandb_config self.wandb_initialized = False - if self.wandb_config is not None and self.dist_env.is_main: + if self.wandb_config is not None: self._init_wandb() self.verbose = verbose_from_wandb_config(self.wandb_config) @@ -423,8 +460,6 @@ def _shutdown_agent_task_executors(self, *, wait: bool = True) -> None: def _init_wandb(self): """Initialize Weights & Biases for tracking with multi-turn config.""" - if not self.dist_env.is_main: - return if not self.wandb_initialized: if self.wandb_config is None: self.wandb_config = {} @@ -445,7 +480,7 @@ def _init_wandb(self): } if hasattr(self.args, "num_generations"): config_dict["num_generations"] = self.args.num_generations - if getattr(self, "centralized_collaboration", None) is not None: + if self.is_centralized: config_dict["collaboration_mode"] = "centralized" config_dict["num_roles"] = self.centralized_collaboration.num_roles config_dict["num_actor_models"] = len(self.agents) @@ -663,6 +698,9 @@ def _log_eval_metrics( extra_metrics: Optional[Dict[str, Any]] = None, ) -> Dict[str, float]: """Log evaluation metrics for any number of turns.""" + all_agent_completions_turns = self._role_completions( + all_agent_completions_turns, all_batch_items + ) eval_metrics = {} # Detailed logging (if logger is provided), standardized to modern interface @@ -709,24 +747,14 @@ def train( their logs have been processed. The configured epoch count is the total for the run, including epochs already completed in a resumed checkpoint. """ - if ( - isinstance(checkpoint_interval, bool) - or not isinstance(checkpoint_interval, numbers.Integral) - or checkpoint_interval < 1 - ): - raise ValueError("checkpoint_interval must be a positive integer.") - if ( - resume_from_checkpoint is not None - or checkpoint_dir is not None - or checkpoint_interval != 1 - or checkpoint_total_limit is not None - ): - with self._checkpoint_context( - resume_from_checkpoint=resume_from_checkpoint, - checkpoint_dir=checkpoint_dir, - checkpoint_interval=checkpoint_interval, - checkpoint_total_limit=checkpoint_total_limit, - ): + checkpointing = self._checkpoint_arguments( + resume_from_checkpoint, + checkpoint_dir, + checkpoint_interval, + checkpoint_total_limit, + ) + if checkpointing is not None: + with self._checkpoint_context(**checkpointing): return self._train_checkpoint_epochs( self._checkpoint_progress, self._checkpoint_commit, **kwargs ) @@ -736,39 +764,7 @@ def train( if self.wandb_config is not None and not self.wandb_initialized: self._init_wandb() - - for agent_idx, agent in enumerate(self.agents): - agent.to(self.agent_devices[agent_idx]) - agent.train() - - for epoch in range(0, int(self.args.num_train_epochs)): - dl = self.get_train_dataloader() - if getattr(self, "verbose", True): - it = enumerate( - tqdm( - dl, - total=len(dl), - desc=f"Epoch {epoch+1}/{int(self.args.num_train_epochs)}", - ) - ) - else: - it = enumerate(dl) - for batch_idx, batch in it: - if int(self.args.eval_interval) > 0 and ( - batch_idx % int(self.args.eval_interval) == 0 - ): - _ = self.evaluate(num_eval_samples=int(self.args.eval_num_samples)) - - # Process single batch item (batch_size=1 enforced) - batch_item = batch[0] - self._train_step_returns(batch_item, **kwargs) - - ready_agents = [ - agent_idx - for agent_idx, buffer in enumerate(self.rollout_buffers) - if buffer - ] - self._drain_ready_buffers(ready_agents) + self._run_epochs(0, None, kwargs) if not getattr(self, "_checkpoint_context_active", False): self._checkpoint_progress = { @@ -828,7 +824,6 @@ def _on_checkpoint_train_start(self, start_epoch: int, total_epochs: int) -> Non def _train_checkpoint_epochs(self, progress, commit, **kwargs): """Run complete RL epochs, optionally inside a parent trainer's context.""" - total_epochs = int(self.args.num_train_epochs) # Preference subclasses have an outer cursor validator; this loop uses RL epochs. start_epoch = MAGRPOTrainer._checkpoint_validate_progress( self, @@ -837,10 +832,13 @@ def _train_checkpoint_epochs(self, progress, commit, **kwargs): "runtime": {"rollout_buffers": self.rollout_buffers}, }, ) - self._on_checkpoint_train_start(start_epoch, total_epochs) - + self._on_checkpoint_train_start(start_epoch, int(self.args.num_train_epochs)) self._checkpoint_init_wandb() + self._run_epochs(start_epoch, commit, kwargs) + def _run_epochs(self, start_epoch, commit, generation_kwargs): + """Train epochs ``start_epoch`` onward; ``commit`` (if given) follows each epoch.""" + total_epochs = int(self.args.num_train_epochs) for agent_idx, agent in enumerate(self.agents): agent.to(self.agent_devices[agent_idx]) agent.train() @@ -863,8 +861,9 @@ def _train_checkpoint_epochs(self, progress, commit, **kwargs): ): _ = self.evaluate(num_eval_samples=int(self.args.eval_num_samples)) + # Process single batch item (batch_size=1 enforced) batch_item = batch[0] - self._train_step_returns(batch_item, **kwargs) + self._train_step_returns(batch_item, **generation_kwargs) ready_agents = [ agent_idx @@ -872,10 +871,11 @@ def _train_checkpoint_epochs(self, progress, commit, **kwargs): if buffer ] self._drain_ready_buffers(ready_agents) - commit( - {"kind": "magrpo", "epoch": epoch + 1}, - force=epoch + 1 == total_epochs, - ) + if commit is not None: + commit( + {"kind": "magrpo", "epoch": epoch + 1}, + force=epoch + 1 == total_epochs, + ) def _train_step_returns( self, @@ -1229,7 +1229,7 @@ def _generate_completions( completion_input_ids = generation_output.sequences prompt_len = prompt_input_ids[0].shape[0] - joint_actor = getattr(self, "centralized_collaboration", None) is not None + joint_actor = self.is_centralized # Joint prompts can contain EOS even when it is also the padding token. if not joint_actor: pad_positions = (prompt_input_ids[0] == tokenizer.pad_token_id).nonzero() @@ -1430,31 +1430,7 @@ def _pack_completions_for_buffer( "reference_kls": list(completions_data.get("reference_kls") or []), } - def _reference_kl_values( - self, - agent_idx: int, - policy_model: Any, - sequences: torch.Tensor, - attention_mask: torch.Tensor, - prompt_len: int, - response_lens: Sequence[int], - ) -> List[float]: - if not reference_kl_enabled(self.args): - return [] - if not getattr(self, "reference_models", None): - raise RuntimeError( - "Reference KL is enabled but reference models are missing." - ) - reference_model = self.reference_models[agent_idx] - return _sequence_kl_values( - policy_model, - reference_model, - sequences, - attention_mask, - prompt_len, - response_lens, - sequence_kl=reference_kl_for_sequence, - ) + _reference_kl_values = reference_kl_values def _apply_reference_kl_to_returns( self, returns_tensor: torch.Tensor, completions_data: Dict[str, Any] @@ -1585,6 +1561,10 @@ def _update_from_samples(self, agent_idx: int, samples: List[NodeSample]) -> Non sample.returns, ) (loss * scale).backward() + if getattr(self.args, "skip_zero_grad_steps", False) and not _has_gradient( + self.optimizers[agent_idx] + ): + continue self.optimizers[agent_idx].step() def _compute_advantages(self, returns_tensor: torch.Tensor) -> torch.Tensor: @@ -1615,6 +1595,10 @@ def _compute_loss_with_gradients(self, agent, completions_data, returns): Returns: torch.Tensor: The computed loss with gradients attached """ + if self.is_centralized and not getattr( + self.args, "match_decentralized_loss", False + ): + return joint_policy_gradient_loss(self, agent, completions_data, returns) agent_module = unwrap_model(agent) device = next(agent_module.parameters()).device diff --git a/comlrl/utils/checkpoint.py b/comlrl/utils/checkpoint.py index e4210f6..db1f3ee 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -6,8 +6,7 @@ """ from contextlib import contextmanager -from dataclasses import fields, is_dataclass -import importlib +from dataclasses import MISSING, fields, is_dataclass import inspect import hashlib import json @@ -29,13 +28,8 @@ CHECKPOINT_VERSION = 1 STATE_FILENAME = "training_state.pt" _OWNER_FILENAME = "checkpoint_owner.json" -_DATACLASSES = { - "comlrl.trainers.reinforce.magrpo.NodeSample", - "comlrl.trainers.actor_critic.iac.RolloutSample", - "comlrl.trainers.preference.madpo.PreferencePair", - "comlrl.trainers.preference.madpo.AgentPreferenceTensors", - "comlrl.trainers.preference.iterative.PreferenceReplayShard", -} +# Trainer dataclasses allowed in checkpoints; see register_checkpoint_record. +_RECORDS = {} _MODEL_FIELDS = ("agents", "critics", "reference_models", "reward_model") _OPTIMIZER_FIELDS = ( "optimizers", @@ -172,6 +166,14 @@ def _class_name(value): return f"{cls.__module__}.{cls.__qualname__}" +def register_checkpoint_record(cls): + """Allow instances of a trainer dataclass in training checkpoints.""" + if not isinstance(cls, type) or not is_dataclass(cls): + raise TypeError("Checkpoint records must be dataclass types.") + _RECORDS[f"{cls.__module__}.{cls.__qualname__}"] = cls + return cls + + def _encode(value, *, copy_tensors=True): """Copy tensors to CPU and encode known records using safe built-in types.""" if value is None or type(value) in (bool, int, float, str, bytes): @@ -193,7 +195,7 @@ def _encode(value, *, copy_tensors=True): return {"__kind__": "text", "value": str(value)} if is_dataclass(value) and not isinstance(value, type): name = _class_name(value) - if name not in _DATACLASSES: + if name not in _RECORDS: raise TypeError(f"Unsupported checkpoint record: {name}") return { "__kind__": "record", @@ -242,10 +244,9 @@ def _decode(value): return np.frombuffer(value["data"], dtype=dtype).reshape(value["shape"]).copy() if kind == "record": name = value["class"] - if name not in _DATACLASSES: + if name not in _RECORDS: raise ValueError(f"Unsupported checkpoint record: {name}") - module, _, attr = name.rpartition(".") - cls = getattr(importlib.import_module(module), attr) + cls = _RECORDS[name] return cls(**{key: _decode(item) for key, item in value["fields"].items()}) raise ValueError("Invalid checkpoint value encoding.") @@ -293,6 +294,33 @@ def _config(trainer): } +def _config_defaults(trainer, names): + """Encoded defaults of the config fields ``names`` that a checkpoint lacks.""" + args = getattr(trainer, "args", None) + if not names or not is_dataclass(args): + return {} + defaults = {} + for field in fields(args): + if field.name not in names: + continue + if field.default is not MISSING: + default = field.default + elif field.default_factory is not MISSING: + # A factory may draw random numbers; validation must not move the RNG. + rng = capture_rng_state() + try: + default = field.default_factory() + finally: + restore_rng_state(rng) + else: + continue + try: + defaults[field.name] = _decode(_encode(default)) + except TypeError: + continue + return defaults + + def _dataset_info(trainer): dataset = getattr(trainer, "train_dataset", None) if dataset is None: @@ -750,11 +778,20 @@ def _validate_state(trainer, state, *, model_config_updates=None): for key, value in state["config"].items() if key not in _PORTABLE_CONFIG_FIELDS } - if saved_config != _config(trainer): + current_config = _config(trainer) + missing = set(current_config) - set(saved_config) + defaults = _config_defaults(trainer, missing) + # A field added after the checkpoint was written matches at its default. + for key in missing: + if key in defaults and current_config[key] == defaults[key]: + saved_config[key] = current_config[key] + if saved_config != current_config: changed = sorted( key - for key in set(saved_config) | set(_config(trainer)) - if saved_config.get(key) != _config(trainer).get(key) + for key in set(saved_config) | set(current_config) + if key not in saved_config + or key not in current_config + or saved_config[key] != current_config[key] ) raise ValueError( f"Checkpoint configuration does not match: {', '.join(changed)}" @@ -942,14 +979,14 @@ def restore_trainer_state(trainer, state, checkpoint_dir=None, restore_rng=True) assets_hook = getattr(trainer, "_checkpoint_load_assets", None) if callable(assets_hook) and "assets" in state: assets_hook(checkpoint_dir, state["assets"]) - extra_hook = getattr(trainer, "_restore_checkpoint_extra_state", None) + extra_hook = getattr(trainer, "_checkpoint_restore_extra_state", None) if callable(extra_hook): extra_hook(state["extra"], checkpoint_dir) _restore_wandb_identity(trainer, state.get("wandb")) if restore_rng: _restore_trainer_rng(trainer, state, additional_devices) trainer._checkpoint_save_blocked = False - trainer._checkpoint_untracked_preference = False + trainer._checkpoint_untracked = False def _fsync_directory(directory): @@ -1076,6 +1113,34 @@ def _preflight_checkpoint_directory(checkpoint_dir, metadata=None): class TrainingCheckpointMixin: """Explicit checkpoint APIs; ordinary training performs no checkpoint I/O.""" + def _checkpoint_arguments( + self, + resume_from_checkpoint, + checkpoint_dir, + checkpoint_interval, + checkpoint_total_limit, + ): + """Validate ``train``'s checkpoint arguments; return them when any is set.""" + if ( + isinstance(checkpoint_interval, bool) + or not isinstance(checkpoint_interval, numbers.Integral) + or checkpoint_interval < 1 + ): + raise ValueError("checkpoint_interval must be a positive integer.") + if ( + resume_from_checkpoint is None + and checkpoint_dir is None + and checkpoint_interval == 1 + and checkpoint_total_limit is None + ): + return None + return { + "resume_from_checkpoint": resume_from_checkpoint, + "checkpoint_dir": checkpoint_dir, + "checkpoint_interval": checkpoint_interval, + "checkpoint_total_limit": checkpoint_total_limit, + } + def _checkpoint_init_wandb(self): """Initialize resumed logging without advancing training RNG streams.""" if getattr(self, "wandb_config", None) is None or getattr( @@ -1113,9 +1178,9 @@ def save_checkpoint(self, output_dir): raise RuntimeError( "Training did not reach a completed checkpoint boundary; resume the last saved checkpoint." ) - if getattr(self, "_checkpoint_untracked_preference", False): + if getattr(self, "_checkpoint_untracked", False): raise RuntimeError( - "Preference checkpoint tracking was not enabled. Use train(checkpoint_dir=...) to retain resumable state." + "Checkpoint tracking was not enabled. Use train(checkpoint_dir=...) to retain resumable state." ) target = Path(output_dir).expanduser().absolute() if target.is_symlink() or ( @@ -1228,7 +1293,7 @@ def _checkpoint_context( self._checkpoint_context_active = True self._checkpoint_safe_to_save = False self._checkpoint_save_blocked = True - self._checkpoint_untracked_preference = False + self._checkpoint_untracked = False try: yield self._checkpoint_progress except BaseException: @@ -1293,9 +1358,6 @@ def _checkpoint_prune(self, latest): for _, path in older[max(0, int(limit) - 1) :]: shutil.rmtree(path) _fsync_directory(root) - if callable(getattr(self, "_checkpoint_save_assets", None)): - from comlrl.trainers.preference.checkpointing import ( - prune_preference_asset_pool, - ) - - prune_preference_asset_pool(root) + prune_assets = getattr(self, "_checkpoint_prune_assets", None) + if callable(prune_assets): + prune_assets(root) diff --git a/comlrl/utils/distributed.py b/comlrl/utils/distributed.py index dc513d0..120e8b5 100644 --- a/comlrl/utils/distributed.py +++ b/comlrl/utils/distributed.py @@ -1,5 +1,6 @@ from __future__ import annotations +import warnings from dataclasses import dataclass from typing import Any, Optional @@ -8,11 +9,20 @@ @dataclass(frozen=True) class DistributedContext: + """Deprecated: CoMLRL trains in one process; kept so existing imports still work.""" + is_main: bool device: torch.device def local_context(device: Optional[torch.device] = None) -> DistributedContext: + """Deprecated: always the main process, since CoMLRL trains in one process.""" + warnings.warn( + "local_context and DistributedContext are deprecated and will be removed in " + "a future release; CoMLRL trains in one process.", + DeprecationWarning, + stacklevel=2, + ) if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") return DistributedContext(is_main=True, device=device) diff --git a/comlrl/utils/reference_kl.py b/comlrl/utils/reference_kl.py index 15b31dc..6fb07e1 100644 --- a/comlrl/utils/reference_kl.py +++ b/comlrl/utils/reference_kl.py @@ -1,7 +1,7 @@ from __future__ import annotations import copy -from typing import Any, Callable, Mapping, List, Sequence +from typing import Any, Mapping, List, Sequence import torch import torch.nn.functional as F @@ -108,19 +108,28 @@ def response_token_logprobs( return token_log_probs[:, start_index:end_index] -def _sequence_kl_values( +def reference_kl_values( + trainer: Any, + agent_idx: int, policy_model: Any, - reference_model: Any, sequences: torch.Tensor, attention_mask: torch.Tensor, prompt_len: int, response_lens: Sequence[int], - *, - sequence_kl: Callable[..., torch.Tensor], ) -> List[float]: + """One KL estimate per sampled sequence against ``trainer``'s frozen reference. + + Trainers bind this as ``_reference_kl_values``; it returns ``[]`` while + reference KL is disabled. + """ + if not reference_kl_enabled(trainer.args): + return [] + if not getattr(trainer, "reference_models", None): + raise RuntimeError("Reference KL is enabled but reference models are missing.") + reference_model = trainer.reference_models[agent_idx] values: List[float] = [] for seq, attn, response_len in zip(sequences, attention_mask, response_lens): - kl_value = sequence_kl( + kl_value = reference_kl_for_sequence( policy_model, reference_model, seq.unsqueeze(0), diff --git a/docs/content/_index.md b/docs/content/_index.md index 6317b9a..46de149 100644 --- a/docs/content/_index.md +++ b/docs/content/_index.md @@ -46,9 +46,11 @@ During decentralized execution, agents use their own prompts and policies. Indep - **_Multi-Agent REINFORCE:_** Critic-free policy gradient methods, including [MAREINFORCE](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/mareinforce.py), [MAGRPO](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/magrpo.py), [MARLOO](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/marloo.py), and [MAReMax](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/reinforce/maremax.py). - Index-aligned joint responses with `joint_mode='aligned'`. - Cartesian-product joint responses with `joint_mode='cross'`. + - Decentralized actors or one centralized actor with `collaboration_mode='centralized'`. - **_Multi-Agent Actor-Critic:_** Actor-Critic methods, including [IAC](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/actor_critic/iac.py) and [MAAC](https://github.com/OpenMLRL/CoMLRL/blob/main/comlrl/trainers/actor_critic/maac.py). - Independent actor-critic (separate critic or value head over the LLM backbone). - Centralized critic over joint prompts with separate actors. + - Decentralized actors or one centralized actor with `collaboration_mode='centralized'`. - **_Multi-Agent PbRL:_** MADPO, MARLHF, and their iterative variants. - Direct joint preference optimization or a learned joint reward followed by RL. - Decentralized actors or one centralized actor with role-aware joint responses. diff --git a/docs/content/docs/marl/multi-agent-actor-critic.md b/docs/content/docs/marl/multi-agent-actor-critic.md index 7151732..9a33ca3 100644 --- a/docs/content/docs/marl/multi-agent-actor-critic.md +++ b/docs/content/docs/marl/multi-agent-actor-critic.md @@ -36,6 +36,7 @@ In Q mode, the implementation uses the same TD-residual actor weight, with a sam **MAACConfig** parameters: - `num_agents`: Number of agents - `num_turns`: Number of turns +- `collaboration_mode`: `"decentralized"` (default) trains one actor per agent; `"centralized"` trains one actor that writes the complete joint response for all `num_agents` roles and requires `num_turns=1`. The shared critic still reads each role's prompt and parsed output, in the same format as in decentralized training - `critic_type`: Critic target type (`v` for V(h), `q` for Q(h,a)) - `num_train_epochs`: Number of training epochs - `agent_learning_rate`: Learning rate for agents @@ -44,6 +45,7 @@ In Q mode, the implementation uses the same TD-residual actor weight, with a sam - `advantage_normalization`: Whether to normalize advantages before updates - `rollout_buffer_size`: Number of samples to collect per agent before an update - `train_batch_size`: Mini-batch size within each update +- `per_sample_backward`: Backpropagates each sample as soon as its loss is computed and steps once per mini-batch, so only one sample's activations stay in memory instead of the whole mini-batch's; the update is mathematically the same, though floating-point sums may differ slightly (default `False`) - `max_new_tokens`: Maximum tokens to generate per completion - `temperature`: Temperature for sampling - `top_p`: Top-p for nucleus sampling @@ -77,6 +79,7 @@ In Q mode, the implementation uses the same TD-residual actor weight, with a sam - `model_config`: Model loading options; `critic_value_head_hidden_dim` configures the centralized critic's value head (optional) - `wandb_config`: Weights & Biases logging configuration (optional) - `metrics_callback`: Optional callback taking rollout samples and returning a dictionary of scalar metrics +- `centralized_adapter`: Joint prompt builder and role parser (`build_prompt`, `parse_completion`) for a centralized actor; required with `collaboration_mode="centralized"`. Rewards, the critic and `metrics_callback` receive the parsed role outputs - `external_transition`: Transition function required when `num_turns > 1` {{% /hint %}} @@ -124,6 +127,7 @@ where {{< katex >}}\hat{V}_t{{< /katex >}} is the TD target, {{< katex >}}\epsil - `num_agents`: Number of agents - `num_turns`: Number of turns +- `collaboration_mode`: `"decentralized"` (default) trains one actor per agent; `"centralized"` trains one actor that writes the complete joint response for all `num_agents` roles and requires `num_turns=1`. Its critic conditions on the joint prompt - `num_train_epochs`: Number of training epochs - `agent_learning_rate`: Learning rate for agents - `critic_learning_rate`: Learning rate for separate critics; shared models use `agent_learning_rate` for both actor and value head @@ -132,6 +136,7 @@ where {{< katex >}}\hat{V}_t{{< /katex >}} is the TD target, {{< katex >}}\epsil - `advantage_normalization`: Whether to normalize advantages - `rollout_buffer_size`: Number of samples to collect before an update - `train_batch_size`: Mini-batch size for policy updates +- `per_sample_backward`: Backpropagates each sample as soon as its loss is computed and steps once per mini-batch, so only one sample's activations stay in memory instead of the whole mini-batch's; the update is mathematically the same, though floating-point sums may differ slightly (default `False`) - `max_new_tokens`: Maximum new tokens to generate - `temperature`: Temperature for sampling - `top_p`: Top-p for nucleus sampling @@ -169,6 +174,7 @@ where {{< katex >}}\hat{V}_t{{< /katex >}} is the TD target, {{< katex >}}\epsil - `model_config`: Model configuration dictionary (optional) - `wandb_config`: Configuration for Weights & Biases logging (optional) - `metrics_callback`: Optional callback taking rollout samples and returning a dictionary of scalar metrics +- `centralized_adapter`: Joint prompt builder and role parser (`build_prompt`, `parse_completion`) for a centralized actor; required with `collaboration_mode="centralized"`. Rewards and `metrics_callback` receive the parsed role outputs - `external_transition`: Optional transition function required for multi-turn training {{% /hint %}} diff --git a/docs/content/docs/marl/multi-agent-reinforce.md b/docs/content/docs/marl/multi-agent-reinforce.md index 89c1814..6afecd0 100644 --- a/docs/content/docs/marl/multi-agent-reinforce.md +++ b/docs/content/docs/marl/multi-agent-reinforce.md @@ -54,7 +54,8 @@ The current `MAGRPOTrainer` loss slice omits the final completion token, so its {{% hint info %}} **MAGRPOConfig** parameters: -- `num_agents`: Configured number of agents, which should match the trainer setup +- `num_agents`: Configured number of agents, or of task roles with a centralized actor, which should match the trainer setup +- `collaboration_mode`: `"decentralized"` (default) trains one actor per agent; `"centralized"` trains one actor that writes the complete joint response for all roles, requires `num_turns=1`, and `max_new_tokens` then covers the whole joint response - `num_turns`: Number of turns per episode - `num_train_epochs`: Number of training epochs - `agent_learning_rate`: Learning rate @@ -75,6 +76,8 @@ The current `MAGRPOTrainer` loss slice omits the final completion token, so its - `eval_batch_size`: Evaluation DataLoader batch size - `external_prompt_passthrough`: Whether to use external prompts directly in multi-turn training - `advantage_mode`: Baseline mode (`mean`, `max`, `rloo`, `raw`) +- `skip_zero_grad_steps`: Whether an optimizer step is skipped when all of its gradients are zero, for example when every advantage in the batch is zero; otherwise AdamW keeps moving the weights along its stored momentum (default `False`) +- `match_decentralized_loss`: With a centralized actor, whether the joint response is scored with the decentralized loss (the same token range, precision and non-finite fallback) instead of the joint-sequence loss (default `False`) - `parallel_training`: Agent execution mode (`none` or `mp`) - `agent_devices`: Agent device placement; required explicitly for `parallel_training="mp"` - `reference_kl_enabled`: Whether to enable a frozen-reference KL penalty (default `False`) @@ -99,6 +102,7 @@ The current `MAGRPOTrainer` loss slice omits the final completion token, so its - `external_transition`: Function providing transitions between turns; required when `num_turns > 1` (the default is `2`) - `eval_logger`: Optional task-metrics callback, called as `eval_logger(agent_completions_turns=..., batch_items=..., prompts=...)`; returns one metrics dictionary per evaluated item and is used with `eval_aggregator` - `eval_aggregator`: Optional callback, called as `eval_aggregator(metrics, num_turns=...)`; the returned `turn_k/...` metrics are logged under `eval/`. Both evaluation callbacks must be supplied for custom metrics +- `centralized_adapter`: Joint prompt builder and role parser (`build_prompt`, `parse_completion`) for a centralized actor; required with `collaboration_mode="centralized"`. Rewards and evaluation receive the parsed role outputs {{% /hint %}} {{% hint warning %}} diff --git a/docs/content/docs/marl/multi-llm-pbrl.md b/docs/content/docs/marl/multi-llm-pbrl.md index f4a214c..2e2e01e 100644 --- a/docs/content/docs/marl/multi-llm-pbrl.md +++ b/docs/content/docs/marl/multi-llm-pbrl.md @@ -14,7 +14,7 @@ aliases: - /docs/user-guide/multi-agent-preference-learning/ --- -Multi-agent PbRL trains collaborating LLM agents from preferences between their joint responses. CoMLRL builds the preference data from task items and learns from it with one of two algorithms: **MARLHF** fits a joint reward model and then trains the policies with RL, while **MADPO** optimizes the policies directly from the preferences. `MARLHFTrainer`, `MADPOTrainer`, and their iterative variants `MARLHFIterTrainer` and `MADPOIterTrainer` are available from `comlrl.trainers.preference`. +Multi-agent PbRL trains collaborating LLM agents from preferences between their joint responses. CoMLRL builds the preference data from task items and learns from it with one of two algorithms: **MARLHF** fits a joint reward model and then trains the policies with RL, while **MADPO** optimizes the policies directly from the preferences. `MARLHFTrainer`, `MADPOTrainer`, and their iterative variants `MARLHFIterTrainer` and `MADPOIterTrainer` are available from `comlrl.trainers.pbrl`. {{% hint warning %}} All four trainers currently require `num_turns=1` and aligned joint responses (`joint_mode="aligned"`). An iteration of the iterative trainers is a round of data collection and training, not a dialog turn. @@ -143,6 +143,7 @@ The `history` source saves policy checkpoints automatically. The two sides may s - `comparator_devices`: Device placement for local comparator models - `comparator_generation_mode`: `"decentralized"` (default) or `"centralized"` - `comparator_centralized_agent_index`: Role whose actor, checkpoint, or model writes a centralized comparator's joint response; defaults to 0, which a centralized actor requires +- `comparator_centralized_max_new_tokens`: Token budget of a local centralized comparator's joint response, which holds every agent's output; defaults to `max_new_tokens` - `policy_checkpoint_dir`: Optional actor-checkpoint location; setting it also saves checkpoints for the other sources {{% /hint %}} @@ -245,7 +246,7 @@ Historical modes assign each of the `preference_replay_sample_size` draws to a d `current` rejects `preference_replay_sample_size` and `preference_replay_lambda` and accepts `preference_replay_k` only as `None` or 1. The following configuration instead samples 256 pairs from all rounds, weighted toward recent ones: ```python -from comlrl.trainers.preference import MADPOIterConfig +from comlrl.trainers.pbrl import MADPOIterConfig args = MADPOIterConfig( num_agents=2, diff --git a/docs/content/docs/user-guide/training-entities.md b/docs/content/docs/user-guide/training-entities.md index 43e1a25..a123d2a 100644 --- a/docs/content/docs/user-guide/training-entities.md +++ b/docs/content/docs/user-guide/training-entities.md @@ -146,7 +146,7 @@ In PbRL, [MARLHF]({{< relref "/docs/marl/multi-llm-pbrl#marlhf" >}}) learns a jo The following example loads a _Qwen/Qwen2.5-Coder-7B_ reward model for two _Qwen/Qwen2.5-Coder-3B_ agents: ```python -from comlrl.trainers.preference import MARLHFConfig, MARLHFTrainer +from comlrl.trainers.pbrl import MARLHFConfig, MARLHFTrainer trainer = MARLHFTrainer( agent_model="Qwen/Qwen2.5-Coder-3B", diff --git a/docs/content/docs/user-guide/training-parallelization.md b/docs/content/docs/user-guide/training-parallelization.md index 0d077ad..07db460 100644 --- a/docs/content/docs/user-guide/training-parallelization.md +++ b/docs/content/docs/user-guide/training-parallelization.md @@ -11,7 +11,7 @@ Each model runs entirely on its assigned device and is not split across GPUs. De The following configuration puts each model of an iterative MARLHF run on its own GPU: two agents, a centralized critic for MAAC policy learning, the reward model, and copies of the current agents as comparators. ```python -from comlrl.trainers.preference import MARLHFIterConfig +from comlrl.trainers.pbrl import MARLHFIterConfig config = MARLHFIterConfig( num_agents=2, diff --git a/tests/checkpoint_helpers.py b/tests/checkpoint_helpers.py index 7215d33..48032f0 100644 --- a/tests/checkpoint_helpers.py +++ b/tests/checkpoint_helpers.py @@ -9,13 +9,13 @@ import numpy as np import torch -from comlrl.trainers.preference import ( +from comlrl.trainers.pbrl import ( MADPOIterTrainer, MADPOTrainer, MARLHFIterTrainer, MARLHFTrainer, ) -from comlrl.trainers.preference.madpo import AgentPreferenceTensors, PreferencePair +from comlrl.trainers.pbrl.madpo import AgentPreferenceTensors, PreferencePair from comlrl.trainers.reinforce.magrpo import MAGRPOConfig, NodeSample from comlrl.utils.checkpoint import TrainingCheckpointMixin @@ -110,7 +110,6 @@ def __init__(self, *, width=3, prompts=("first", "second"), **config): [], ] self.reward_func = StatefulReward() - self.dist_env = SimpleNamespace(is_main=True) self._checkpoint_progress = {"epoch": 1, "next_item": 0} diff --git a/tests/helpers.py b/tests/helpers.py index 3916530..ea7f218 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -12,13 +12,9 @@ import torch from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast -from comlrl.trainers.preference import ( - MADPOIterConfig, - MARLHFConfig, - MARLHFIterConfig, - TaggedCentralizedComparatorAdapter, -) -from comlrl.trainers.reinforce import CentralizedMAGRPOConfig, CentralizedMAGRPOTrainer +from comlrl.trainers.centralized import TaggedCentralizedComparatorAdapter +from comlrl.trainers.pbrl import MADPOIterConfig, MARLHFConfig, MARLHFIterConfig +from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer from comlrl.utils.reward_utils import set_reward_range GOOD = " good good " @@ -167,6 +163,8 @@ def generate(*, input_ids, num_return_sequences, **kwargs): def make_centralized_magrpo_trainer(tokenizer, **overrides): """A centralized MAGRPO trainer whose actor alternates GOOD/BAD samples.""" options = dict( + collaboration_mode="centralized", + num_turns=1, agent_devices="cpu", num_agents=2, num_train_epochs=1, @@ -199,11 +197,11 @@ def generate(*, input_ids, num_return_sequences, **kwargs): model.generate = generate dataset = Dataset.from_list([{"id": "task", "prompt": "task prompt"}]) - trainer = CentralizedMAGRPOTrainer( + trainer = MAGRPOTrainer( agents=[model], num_agents=2, tokenizer=tokenizer, - args=CentralizedMAGRPOConfig(**options), + args=MAGRPOConfig(**options), train_dataset=dataset, eval_dataset=dataset, reward_func=task_reward, diff --git a/tests/test_actor_critic.py b/tests/test_actor_critic.py index 10c2bfe..a4940c2 100644 --- a/tests/test_actor_critic.py +++ b/tests/test_actor_critic.py @@ -7,14 +7,14 @@ import torch from torch import nn import wandb +from datasets import Dataset from comlrl.trainers.actor_critic import IACTrainer, MAACTrainer, RolloutSample from comlrl.trainers.actor_critic.ac_base import ActorCriticTrainerBase from comlrl.trainers.actor_critic.iac import IACConfig from comlrl.trainers.actor_critic.maac import MAACConfig -from comlrl.utils.distributed import DistributedContext -from helpers import tiny_gpt2, zero_reward +from helpers import model_for, tiny_gpt2, zero_reward TRAINERS = [IACTrainer, MAACTrainer] @@ -47,11 +47,6 @@ def forward(self, input_ids, attention_mask, output_values): return SimpleNamespace(logits=logits, values=values) -def _cpu_context(): - return DistributedContext(is_main=True, device=torch.device("cpu")) - - -@pytest.mark.parametrize("trainer_cls", TRAINERS) @pytest.mark.parametrize( "num_agents,num_ret,rewards,expected", [ @@ -62,9 +57,9 @@ def _cpu_context(): ], ) def test_expand_rewards_preserves_generation_precedence( - trainer_cls, num_agents, num_ret, rewards, expected + num_agents, num_ret, rewards, expected ): - trainer = trainer_cls.__new__(trainer_cls) + trainer = ActorCriticTrainerBase.__new__(ActorCriticTrainerBase) trainer.args = SimpleNamespace(num_agents=num_agents) result = trainer._expand_rewards(rewards, num_ret) assert result == expected @@ -72,9 +67,8 @@ def test_expand_rewards_preserves_generation_precedence( assert result[1] == expected[1] -@pytest.mark.parametrize("trainer_cls", TRAINERS) -def test_expand_rewards_rejects_unsupported_lengths(trainer_cls): - trainer = trainer_cls.__new__(trainer_cls) +def test_expand_rewards_rejects_unsupported_lengths(): + trainer = ActorCriticTrainerBase.__new__(ActorCriticTrainerBase) trainer.args = SimpleNamespace(num_agents=2) with pytest.raises(ValueError, match="Reward function must return 1 value"): trainer._expand_rewards([1.0] * 4, 3) @@ -127,9 +121,8 @@ def _expand_rewards(self, rewards, num_ret): ] -@pytest.mark.parametrize("trainer_cls", TRAINERS) -def test_sequence_stats_exclude_prompt_and_trailing_tokens(trainer_cls): - trainer = trainer_cls.__new__(trainer_cls) +def test_sequence_stats_exclude_prompt_and_trailing_tokens(): + trainer = ActorCriticTrainerBase.__new__(ActorCriticTrainerBase) probabilities = torch.tensor( [[[0.2, 0.8], [0.3, 0.7], [0.4, 0.6], [0.9, 0.1], [0.5, 0.5]]], dtype=torch.float64, @@ -260,6 +253,7 @@ def test_ac_step_reads_mask_from_rollout_metadata(trainer_cls): assert actor.logits.grad[:, [0, 2, 4]].count_nonzero() == 0 +@pytest.mark.parametrize("per_sample", [False, True]) @pytest.mark.parametrize( "trainer_cls,separate,critic_type", [ @@ -271,10 +265,16 @@ def test_ac_step_reads_mask_from_rollout_metadata(trainer_cls): (MAACTrainer, True, "q"), ], ) -def test_ac_step_matches_analytic_gradients(trainer_cls, separate, critic_type): +def test_ac_step_matches_analytic_gradients( + trainer_cls, separate, critic_type, per_sample +): + """Per-sample backward gives the same gradients and update as the batch mean.""" trainer = trainer_cls.__new__(trainer_cls) trainer.args = SimpleNamespace( - use_separate_critic=separate, value_loss_coef=0.6, value_clip_range=0.2 + use_separate_critic=separate, + value_loss_coef=0.6, + value_clip_range=0.2, + per_sample_backward=per_sample, ) trainer.critic_type = critic_type trainer.agent_devices = trainer.critic_devices = [torch.device("cpu")] @@ -332,9 +332,31 @@ def test_ac_step_matches_analytic_gradients(trainer_cls, separate, critic_type): assert actor.calls == [(3, False), (value_length, True)] * 2 +def test_sample_loss_shares_stay_within_the_batch_mean_range(): + """fp16: 2 * 59840 overflows, but 59840 / 4 * 2 equals the weighted batch mean.""" + trainer = ActorCriticTrainerBase.__new__(ActorCriticTrainerBase) + trainer.args = SimpleNamespace(value_loss_coef=2.0) + actor_share, value_share = trainer._sample_loss_shares( + torch.tensor([1.0], dtype=torch.float16), + torch.tensor([59840.0], dtype=torch.float16), + batch_size=4, + ) + assert actor_share.item() == 0.25 + assert value_share.item() == 29920.0 + + +@pytest.mark.parametrize( + "config_cls,fifth", + [(MAACConfig, "value_loss_coef"), (IACConfig, "value_clip_range")], +) +def test_per_sample_backward_keeps_positional_config_arguments(config_cls, fifth): + config = config_cls(5e-6, 5e-6, 8, 4, 0.1) + assert getattr(config, fifth) == 0.1 + assert config.per_sample_backward is False + + def test_log_metrics_logs_each_call_directly(monkeypatch): trainer = ActorCriticTrainerBase() - trainer.dist_env = _cpu_context() trainer.wandb_initialized = True trainer.env_step = 5 trainer.args = SimpleNamespace(logging_steps=1) @@ -356,7 +378,6 @@ def __init__(self): self.args = SimpleNamespace(eval_batch_size=1, eval_num_samples=1) self.wandb_initialized = False self.env_step = 0 - self.dist_env = _cpu_context() self.verbose = False self.eval_flags = [] @@ -378,6 +399,57 @@ def _collect_rollouts(self, item): assert "eval/turn_1/reward_mean" in metrics +@pytest.mark.parametrize("trainer_cls", TRAINERS) +def test_training_updates_in_train_mode_and_samples_in_eval_mode( + trainer_cls, tokenizer +): + # Models from from_pretrained start in eval mode, where HF skips gradient + # checkpointing; the critic never samples, so only training sets its mode. + actor, critic = model_for(tokenizer).eval(), model_for(tokenizer).eval() + sampling, updates = [], {actor: [], critic: []} + + def sample(*, input_ids, num_return_sequences=1, **_kwargs): + sampling.append(actor.training) + tail = torch.tensor(tokenizer.encode("good")) + return torch.cat([input_ids[0], tail]).repeat(num_return_sequences, 1) + + def record(module, _args): + if torch.is_grad_enabled(): + updates[module].append(module.training) + + actor.generate = sample + actor.register_forward_pre_hook(record) + critic.register_forward_pre_hook(record) + options = dict( + num_agents=1, + num_turns=1, + agent_devices="cpu", + critic_devices="cpu", + num_train_epochs=1, + num_generations=2, + rollout_buffer_size=2, + max_new_tokens=4, + eval_interval=0, + ) + if trainer_cls is IACTrainer: + args = IACConfig(use_separate_critic=True, **options) + else: + args = MAACConfig(**options) + trainer = trainer_cls( + agents=[actor], + critics=[critic], + tokenizer=tokenizer, + reward_func=zero_reward, + formatters=[lambda _: "left"], + args=args, + train_dataset=Dataset.from_list([{"prompt": "left"}]), + ) + trainer.train() + assert sampling and not any(sampling) + assert updates[actor] and all(updates[actor]) + assert updates[critic] and all(updates[critic]) + + @pytest.mark.parametrize("trainer_cls", TRAINERS) def test_wandb_init_reports_actor_critic_config( trainer_cls, stub_tokenizer, monkeypatch diff --git a/tests/test_centralized_actor_critic.py b/tests/test_centralized_actor_critic.py new file mode 100644 index 0000000..0b5792c --- /dev/null +++ b/tests/test_centralized_actor_critic.py @@ -0,0 +1,190 @@ +"""Centralized IAC and MAAC train one joint actor on all roles' text.""" + +from types import SimpleNamespace + +import pytest +import torch +import wandb +from datasets import Dataset + +from comlrl.trainers.actor_critic import IACConfig, IACTrainer, MAACConfig, MAACTrainer + +from helpers import BAD, GOOD, JointAdapter, model_for, task_reward + +ITEM = {"id": "task", "prompt": "task prompt"} +OPTIONS = dict( + collaboration_mode="centralized", + num_agents=2, + num_turns=1, + agent_devices="cpu", + critic_devices="cpu", + num_train_epochs=1, + num_generations=2, + rollout_buffer_size=2, + max_new_tokens=8, + eval_interval=0, + eval_num_samples=1, + agent_learning_rate=1e-3, + critic_learning_rate=1e-3, + early_termination_threshold=None, +) + + +def _scripted_actor(tokenizer, prompts): + """A tiny actor whose samples alternate between a GOOD and a BAD joint response.""" + model = model_for(tokenizer) + + def generate(*, input_ids, num_return_sequences=1, **kwargs): + prompts.append(tokenizer.decode(input_ids[0])) + rows = [ + torch.cat( + [ + input_ids[0], + torch.tensor( + tokenizer.encode(GOOD if index % 2 == 0 else BAD), + device=input_ids.device, + ), + ] + ) + for index in range(num_return_sequences) + ] + return torch.stack(rows) + + model.generate = generate + return model + + +def _trainer(kind, tokenizer, prompts=None, metrics_callback=None, **overrides): + dataset = Dataset.from_list([ITEM]) + common = dict( + tokenizer=tokenizer, + reward_func=task_reward, + formatters=[lambda _: "left", lambda _: "right"], + train_dataset=dataset, + eval_dataset=dataset, + metrics_callback=metrics_callback, + centralized_adapter=JointAdapter(), + ) + options = {**OPTIONS, **overrides} + actor = _scripted_actor(tokenizer, [] if prompts is None else prompts) + if kind == "iac": + return IACTrainer( + agents=[actor], + args=IACConfig(use_separate_critic=False, **options), + **common, + ) + return MAACTrainer( + agents=[actor], + critics=[model_for(tokenizer)], + args=MAACConfig(**options), + **common, + ) + + +@pytest.mark.parametrize("config_cls", [IACConfig, MAACConfig]) +def test_configs_validate_the_collaboration_mode(config_cls): + assert config_cls().collaboration_mode == "decentralized" + with pytest.raises(ValueError, match="collaboration_mode"): + config_cls(collaboration_mode="joint") + with pytest.raises(ValueError, match="num_turns=1"): + config_cls(collaboration_mode="centralized", num_turns=2) + + +@pytest.mark.parametrize("kind", ["iac", "maac"]) +def test_one_joint_actor_trains_on_the_joint_prompt(kind, tokenizer): + prompts = [] + trainer = _trainer(kind, tokenizer, prompts) + assert trainer.num_agents == len(trainer.agents) == 1 + assert trainer.args.num_agents == 2 + before = [p.detach().clone() for p in trainer.agents[0].parameters()] + trainer.train() + assert prompts and set(prompts) == {"joint left right"} + assert trainer.env_step > 0 + assert any( + not torch.equal(old, new) + for old, new in zip(before, trainer.agents[0].parameters()) + ) + + +@pytest.mark.parametrize("kind", ["iac", "maac"]) +def test_metrics_callback_receives_role_outputs(kind, tokenizer): + seen = [] + + def metrics(rollouts): + seen.extend((sample.agent_idx, sample.completion) for sample in rollouts) + return {"samples": float(len(rollouts))} + + trainer = _trainer(kind, tokenizer, metrics_callback=metrics) + result = trainer.evaluate() + assert {agent for agent, _ in seen} == {0, 1} + assert {completion for _, completion in seen} == {"good"} + assert result["eval/samples"] == 2.0 + + +@pytest.mark.parametrize("kind", ["iac", "maac"]) +def test_subclasses_can_add_generate_arguments(kind, tokenizer): + seen = [] + trainer = _trainer(kind, tokenizer) + generate = trainer.agents[0].generate + + def record(**kwargs): + seen.append(kwargs.get("stop_strings")) + return generate(**kwargs) + + trainer.agents[0].generate = record + trainer._generation_kwargs = lambda agent_idx: {"stop_strings": [""]} + trainer.evaluate() + assert seen and all(stops == [""] for stops in seen) + + +@pytest.mark.parametrize("critic_type", ["v", "q"]) +def test_maac_critic_reads_each_role_like_the_decentralized_critic( + critic_type, tokenizer +): + trainer = _trainer("maac", tokenizer, critic_type=critic_type) + joint = trainer._critic_input_for(ITEM, ["joint left right"], [GOOD]) + roles = trainer._build_critic_input( + ["left", "right"], None if critic_type == "v" else ["good", "good"] + ) + assert joint == roles + assert "[Agent 0] left" in joint and "[Agent 1] right" in joint + assert ("[Agent 1 action]\ngood" in joint) == (critic_type == "q") + + +@pytest.mark.parametrize("kind", ["iac", "maac"]) +def test_centralized_mode_rejects_several_actors(kind, tokenizer): + config = ( + IACConfig(use_separate_critic=False, **OPTIONS) + if kind == "iac" + else MAACConfig(**OPTIONS) + ) + trainer_cls = IACTrainer if kind == "iac" else MAACTrainer + extra = {} if kind == "iac" else {"critics": [model_for(tokenizer)]} + with pytest.raises(ValueError, match="exactly one actor"): + trainer_cls( + agents=[model_for(tokenizer), model_for(tokenizer)], + tokenizer=tokenizer, + reward_func=task_reward, + args=config, + centralized_adapter=JointAdapter(), + **extra, + ) + + +@pytest.mark.parametrize("kind", ["iac", "maac"]) +def test_wandb_reports_the_centralized_actor(kind, tokenizer, monkeypatch): + calls = [] + + def init(**kwargs): + calls.append(kwargs) + return SimpleNamespace(id="run-id", project="project", entity="team") + + monkeypatch.setattr(wandb, "run", None) + monkeypatch.setattr(wandb, "init", init) + trainer = _trainer(kind, tokenizer) + trainer.wandb_config = {"project": "project", "entity": "team"} + trainer._init_wandb() + config = calls[0]["config"] + assert config["collaboration_mode"] == "centralized" + assert config["num_roles"] == 2 + assert config["num_actor_models"] == config["num_agents"] == 1 diff --git a/tests/test_centralized_comparator.py b/tests/test_centralized_comparator.py index 3bd9bd6..c5855bd 100644 --- a/tests/test_centralized_comparator.py +++ b/tests/test_centralized_comparator.py @@ -1,12 +1,13 @@ +from contextlib import nullcontext from types import SimpleNamespace import pytest -from comlrl.trainers.preference.centralized import ( +from comlrl.trainers.centralized import ( CentralizedComparatorParseError, TaggedCentralizedComparatorAdapter, ) -from comlrl.trainers.preference.iterative import MADPOIterConfig, MADPOIterTrainer +from comlrl.trainers.pbrl.iterative import MADPOIterConfig, MADPOIterTrainer class _StubAdapter: @@ -55,9 +56,32 @@ def test_iterative_config_allows_generic_centralized_agent_counts(): assert config.num_agents == 3 -def test_iterative_config_keeps_decentralized_as_default(): - config = MADPOIterConfig(agent_devices="cpu") - assert config.comparator_generation_mode == "decentralized" +@pytest.mark.parametrize("budget,expected", [(None, 8), (24, 24)]) +def test_centralized_comparator_answer_has_its_own_budget(budget, expected): + trainer = MADPOIterTrainer.__new__(MADPOIterTrainer) + trainer.num_agents = 2 + trainer.centralized_comparator_adapter = _StubAdapter() + trainer.args = MADPOIterConfig( + agent_devices="cpu", + max_new_tokens=8, + comparator_generation_mode="centralized", + comparator_centralized_max_new_tokens=budget, + ) + trainer._comparator_rng = lambda agent, **kwargs: nullcontext() + budgets = [] + + def generate(agent, batch_items, **kwargs): + budgets.append(kwargs["max_new_tokens"]) + return {"completions": [["joint"]]} + + trainer._generate_completions = generate + outputs = trainer._generate_centralized_policy_output_for_agent( + "agent", {"id": "item"}, prompt="centralized prompt", num_candidates=1 + ) + assert budgets == [expected] + assert outputs[1]["completions"] == [["joint-agent-1"]] + with pytest.raises(ValueError, match="comparator_centralized_max_new_tokens"): + MADPOIterConfig(agent_devices="cpu", comparator_centralized_max_new_tokens=0) def test_centralized_split_preserves_candidate_and_agent_alignment(): diff --git a/tests/test_centralized_magrpo.py b/tests/test_centralized_magrpo.py index 7fdf112..e9853de 100644 --- a/tests/test_centralized_magrpo.py +++ b/tests/test_centralized_magrpo.py @@ -3,22 +3,18 @@ import pytest import torch -from comlrl.trainers.reinforce import ( - CentralizedMAGRPOConfig, - CentralizedMAGRPOTrainer, - MAGRPOConfig, -) +from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer from helpers import BAD, GOOD, JointAdapter, make_centralized_magrpo_trainer, model_for -def test_config_rejects_wrong_mode_and_multiple_turns(): - assert CentralizedMAGRPOConfig().num_turns == 1 +def test_config_rejects_unknown_mode_and_multiple_turns(): + assert MAGRPOConfig().collaboration_mode == "decentralized" assert MAGRPOConfig().num_turns == 2 - with pytest.raises(ValueError, match="centralized collaboration"): - CentralizedMAGRPOConfig(collaboration_mode="decentralized") + with pytest.raises(ValueError, match="collaboration_mode"): + MAGRPOConfig(collaboration_mode="joint") with pytest.raises(ValueError, match="num_turns=1"): - CentralizedMAGRPOConfig(num_turns=2) + MAGRPOConfig(collaboration_mode="centralized", num_turns=2) def test_one_actor_training_uses_joint_prompt_and_task_rewards(tokenizer): @@ -31,8 +27,6 @@ def test_one_actor_training_uses_joint_prompt_and_task_rewards(tokenizer): assert trainer.env_step == 2 assert not torch.equal(before, trainer.agents[0].transformer.wte.weight) assert trainer.reward_func.reward_range == (0.0, 2.0) - assert not hasattr(trainer, "reward_model") - assert not hasattr(trainer, "comparator_policy") def test_eval_passes_separate_roles_to_eval_logger(tokenizer): @@ -47,33 +41,6 @@ def test_eval_passes_separate_roles_to_eval_logger(tokenizer): assert observed[0]["prompts"] == ["task prompt"] assert result["eval/turn_1/reward_mean"] == 2.0 assert result["eval/turn_1/domain_metric"] == 3.0 - assert not trainer._centralized_eval_items - - -def test_loss_matches_full_joint_sequence_policy_gradient(tokenizer): - trainer, _ = make_centralized_magrpo_trainer(tokenizer) - model = trainer.agents[0] - prompt = torch.tensor(tokenizer.encode("joint left right")) - sequences = [torch.tensor(tokenizer.encode(text)) for text in (GOOD, BAD)] - data = {"prompt_input_ids": prompt[None], "completion_input_ids": [sequences]} - manual = [] - for tokens, advantage in zip(sequences, (1.0, -1.0)): - logits = model(torch.cat([prompt, tokens])[None], use_cache=False).logits[ - 0, len(prompt) - 1 : -1 - ] - manual.append( - -logits.log_softmax(-1).gather(1, tokens[:, None]).sum() * advantage - ) - expected = torch.stack(manual).mean() - actual = trainer._compute_loss_with_gradients(model, data, [2.0, 0.0]) - torch.testing.assert_close(actual, expected) - expected.backward() - expected_grad = model.transformer.wte.weight.grad.clone() - model.zero_grad() - actual.backward() - torch.testing.assert_close(model.transformer.wte.weight.grad, expected_grad) - for token in ("", "", ""): - assert expected_grad[tokenizer.convert_tokens_to_ids(token)].abs().sum() > 0 def test_joint_actor_gets_one_reference_model(tokenizer): @@ -87,10 +54,12 @@ def test_joint_actor_gets_one_reference_model(tokenizer): def test_rejects_multiple_actors(tokenizer): with pytest.raises(ValueError, match="exactly one actor"): - CentralizedMAGRPOTrainer( + MAGRPOTrainer( agents=[model_for(tokenizer), model_for(tokenizer)], centralized_adapter=JointAdapter(), - args=CentralizedMAGRPOConfig(agent_devices="cpu"), + args=MAGRPOConfig( + collaboration_mode="centralized", num_turns=1, agent_devices="cpu" + ), ) @@ -107,3 +76,24 @@ def test_wandb_does_not_claim_a_comparator(tokenizer, monkeypatch): assert captured["config"]["num_roles"] == 2 assert captured["config"]["num_actor_models"] == 1 assert "comparator_generation_mode" not in captured["config"] + + +def test_matching_the_decentralized_loss_reuses_the_magrpo_objective(tokenizer): + from comlrl.trainers.reinforce import MAGRPOTrainer + + trainer, _ = make_centralized_magrpo_trainer( + tokenizer, match_decentralized_loss=True + ) + model = trainer.agents[0] + prompt = torch.tensor(tokenizer.encode("joint left right")) + sequences = [torch.tensor(tokenizer.encode(text)) for text in (GOOD, BAD)] + data = {"prompt_input_ids": prompt[None], "completion_input_ids": [sequences]} + actual = trainer._compute_loss_with_gradients(model, data, [2.0, 0.0]) + expected = MAGRPOTrainer._compute_loss_with_gradients( + trainer, model, data, [2.0, 0.0] + ) + torch.testing.assert_close(actual, expected) + trainer.args.match_decentralized_loss = False + joint = trainer._compute_loss_with_gradients(model, data, [2.0, 0.0]) + # The joint loss also scores each sequence's last token; MAGRPO's does not. + assert not torch.allclose(joint, actual) diff --git a/tests/test_centralized_preference.py b/tests/test_centralized_preference.py index 3e9a179..b851ba6 100644 --- a/tests/test_centralized_preference.py +++ b/tests/test_centralized_preference.py @@ -6,7 +6,7 @@ import pytest import torch -from comlrl.trainers.preference import ( +from comlrl.trainers.pbrl import ( MADPOConfig, MADPOIterConfig, MADPOIterTrainer, @@ -16,8 +16,8 @@ MARLHFIterTrainer, MARLHFTrainer, ) -from comlrl.trainers.preference.collaboration import CentralizedCollaboration -from comlrl.trainers.preference.marlhf import JointRewardModel +from comlrl.trainers.centralized import CentralizedCollaboration +from comlrl.trainers.pbrl.marlhf import JointRewardModel from helpers import ( BAD, @@ -68,7 +68,10 @@ def forward(self, input_ids, **kwargs): "config_cls", [MADPOConfig, MARLHFConfig, MADPOIterConfig, MARLHFIterConfig] ) def test_config_defaults_and_centralized_comparator(config_cls): - assert config_cls(agent_devices="cpu").collaboration_mode == "decentralized" + default = config_cls(agent_devices="cpu") + assert default.collaboration_mode == "decentralized" + if hasattr(default, "comparator_generation_mode"): + assert default.comparator_generation_mode == "decentralized" config = config_cls(agent_devices="cpu", collaboration_mode="centralized") assert config.num_agents == 2 if hasattr(config, "comparator_generation_mode"): @@ -89,18 +92,6 @@ def test_centralized_mode_requires_one_actor(tokenizer): ) -def test_centralized_reference_uses_one_device(tokenizer, tmp_path): - trainer, _ = make_preference_trainer( - MADPOTrainer, - MADPOConfig, - tokenizer, - tmp_path, - reference_kl_enabled=True, - reference_devices=["cpu"], - ) - assert len(trainer.reference_models) == 1 - - @pytest.mark.parametrize( "cls,config_cls", [(MADPOTrainer, MADPOConfig), (MADPOIterTrainer, MADPOIterConfig)] ) @@ -164,7 +155,10 @@ def test_replay_stores_raw_joint_text_and_current_comparator_uses_same_prompt( assert len(pairs[0].agent_tensors) == 1 torch.testing.assert_close(calls[0], calls[1]) shard = trainer._write_iteration_preference_pairs(0, pairs) - loaded = trainer._load_replay_shard(shard) + loaded = [ + trainer._preference_pair_from_record(record) + for record in trainer._load_replay_records(shard) + ] assert loaded[0].winner_completions == [GOOD] torch.testing.assert_close( loaded[0].agent_tensors[0].winner_completion_ids, @@ -201,7 +195,6 @@ def test_eval_splits_roles_and_uses_task_reward(cls, config_cls, tokenizer, tmp_ assert observed[0]["agent_completions_turns"] == [[["good"]], [["good"]]] assert result["eval/turn_1/reward_mean"] == 2.0 assert result["eval/turn_1/domain_metric"] == 1.0 - assert not trainer._centralized_eval_items @pytest.mark.parametrize("cls,config_cls", MARLHF_TRAINERS) diff --git a/tests/test_checkpoint_assets.py b/tests/test_checkpoint_assets.py index 1bbdbee..281fed9 100644 --- a/tests/test_checkpoint_assets.py +++ b/tests/test_checkpoint_assets.py @@ -12,12 +12,12 @@ import pytest import torch -from comlrl.trainers.preference import ( +from comlrl.trainers.pbrl import ( MADPOIterTrainer, MARLHFIterTrainer, - checkpointing, + checkpoint_extras, ) -from comlrl.trainers.preference.checkpointing import ( +from comlrl.trainers.pbrl.checkpoint_extras import ( _preference_resume_destinations, _resume_output_path, load_preference_assets, @@ -25,8 +25,8 @@ save_preference_assets, validate_preference_assets, ) -from comlrl.trainers.preference.iterative import PreferenceReplayShard -from comlrl.trainers.preference.output_paths import resolve_output_path +from comlrl.trainers.pbrl.iterative import PreferenceReplayShard +from comlrl.trainers.pbrl.output_paths import resolve_output_path from comlrl.utils.checkpoint import STATE_FILENAME from checkpoint_helpers import ( @@ -201,35 +201,6 @@ def test_unsafe_restore_destination_rejected_before_model_mutation(tmp_path): assert not (tmp_path / "new-run").exists() -@pytest.mark.parametrize( - "config", ["checkpoint", "wandb", "sections", "explicit-parent"] -) -def test_resume_fallback_stays_in_requested_run(tmp_path, config): - source, _, _ = source_assets(tmp_path / "original") - checkpoint = tmp_path / "source-checkpoint" - manifest = save_preference_assets(source, checkpoint) - run = tmp_path / "new-run" - target = asset_trainer(run) - target.args.preference_replay_dir = None - target.args.policy_checkpoint_dir = None - target.wandb_config = None - if config == "checkpoint": - target._checkpoint_directory = str(run) - elif config == "wandb": - target.wandb_config = {"output_dir": str(run)} - elif config == "sections": - target.wandb_config = {"config_sections": {"output": {"base_dir": str(run)}}} - else: - target.args.policy_checkpoint_dir = str(run / "policies") - load_preference_assets(target, checkpoint, manifest) - for attribute in ( - "_preference_replay_dir_path", - "_policy_checkpoint_dir_path", - "_reward_distribution_dir_path", - ): - assert Path(getattr(target, attribute)).is_relative_to(run) - - @pytest.mark.parametrize("config", ["checkpoint", "wandb", "explicit"]) def test_resumed_output_paths_match_source_tilde_conventions( tmp_path, monkeypatch, config @@ -398,14 +369,14 @@ def test_history_hash_and_copy_bytes_grow_only_with_new_assets(tmp_path, monkeyp trainer, replay, initial = source_assets(tmp_path / "run") counts = observe_hash_reads(monkeypatch) copy_bytes = 0 - copy2 = checkpointing.shutil.copy2 + copy2 = checkpoint_extras.shutil.copy2 def counted_copy(source, destination, *args, **kwargs): nonlocal copy_bytes copy_bytes += Path(source).stat().st_size return copy2(source, destination, *args, **kwargs) - monkeypatch.setattr(checkpointing.shutil, "copy2", counted_copy) + monkeypatch.setattr(checkpoint_extras.shutil, "copy2", counted_copy) unique_bytes = replay.stat().st_size + initial.stat().st_size for index in range(6): policy = initial.parents[2] / f"iteration_{index:04d}/agent_0/weights.bin" @@ -413,14 +384,14 @@ def counted_copy(source, destination, *args, **kwargs): payload = bytes([index]) * (128 * 1024) policy.write_bytes(payload) unique_bytes += len(payload) - checkpointing.save_preference_assets( + checkpoint_extras.save_preference_assets( trainer, tmp_path / "checkpoints" / f"checkpoint-{index}" ) # One source hash and one verification hash per newly copied file. assert sum(counts.values()) == 2 * unique_bytes assert copy_bytes == unique_bytes before_reads, before_copies = sum(counts.values()), copy_bytes - checkpointing.save_preference_assets( + checkpoint_extras.save_preference_assets( trainer, tmp_path / "checkpoints" / "checkpoint-unchanged" ) assert sum(counts.values()) == before_reads @@ -439,11 +410,11 @@ def test_digest_cache_invalidates_modified_or_replaced_source( cache = {} counts = observe_hash_reads(monkeypatch) assert ( - checkpointing._asset_digest(path, cache=cache) + checkpoint_extras._asset_digest(path, cache=cache) == hashlib.sha256(original).hexdigest() ) assert ( - checkpointing._asset_digest(path, cache=cache) + checkpoint_extras._asset_digest(path, cache=cache) == hashlib.sha256(original).hexdigest() ) assert sum(counts.values()) == len(original) @@ -457,38 +428,17 @@ def test_digest_cache_invalidates_modified_or_replaced_source( path.write_bytes(modified) if change == "same_size_restored_mtime": os.utime(path, ns=(previous.st_atime_ns, previous.st_mtime_ns)) - assert checkpointing._asset_stat_key(path.stat()) != checkpointing._asset_stat_key( - previous - ) + assert checkpoint_extras._asset_stat_key( + path.stat() + ) != checkpoint_extras._asset_stat_key(previous) assert ( - checkpointing._asset_digest(path, cache=cache) + checkpoint_extras._asset_digest(path, cache=cache) == hashlib.sha256(modified).hexdigest() ) assert sum(counts.values()) == len(original) + len(modified) assert len(cache) == 1 -def test_reused_blob_needs_no_byte_reads_and_changed_source_gets_new_blob( - tmp_path, monkeypatch -): - source = tmp_path / "weights.bin" - source.write_bytes(b"first") - pool, cache = tmp_path / "pool", {} - counts = observe_hash_reads(monkeypatch) - first = tmp_path / "first/weights.bin" - checkpointing._copy_checkpoint_asset(source, first, pool, digest_cache=cache) - first_reads = sum(counts.values()) - second = tmp_path / "second/weights.bin" - checkpointing._copy_checkpoint_asset(source, second, pool, digest_cache=cache) - assert sum(counts.values()) == first_reads - assert first.stat().st_ino == second.stat().st_ino - source.write_bytes(b"new weights") - third = tmp_path / "third/weights.bin" - checkpointing._copy_checkpoint_asset(source, third, pool, digest_cache=cache) - assert sum(counts.values()) == first_reads + 2 * len(b"new weights") - assert third.stat().st_ino != first.stat().st_ino - - def test_source_change_during_hash_is_rejected_without_caching(tmp_path, monkeypatch): source = tmp_path / "weights.bin" source.write_bytes(b"before") @@ -499,45 +449,32 @@ def test_source_change_during_hash_is_rejected_without_caching(tmp_path, monkeyp after_read=lambda: source.write_bytes(b"after!"), ) with pytest.raises(RuntimeError, match="changed during checkpoint saving"): - checkpointing._asset_digest(source, cache=cache) + checkpoint_extras._asset_digest(source, cache=cache) assert not cache -def test_source_change_during_initial_copy_is_rejected(tmp_path, monkeypatch): +@pytest.mark.parametrize("changed", ["source", "copy"]) +def test_initial_copy_rejects_changed_source_or_corrupted_copy( + tmp_path, monkeypatch, changed +): source = tmp_path / "weights.bin" source.write_bytes(b"original") - copy2 = checkpointing.shutil.copy2 + copy2 = checkpoint_extras.shutil.copy2 def change_after_copy(src, dst, *args, **kwargs): result = copy2(src, dst, *args, **kwargs) - if Path(src) == source: + if changed == "copy": + Path(dst).write_bytes(b"corrupt!") + elif Path(src) == source: source.write_bytes(b"modified") return result - monkeypatch.setattr(checkpointing.shutil, "copy2", change_after_copy) + monkeypatch.setattr(checkpoint_extras.shutil, "copy2", change_after_copy) destination, pool = tmp_path / "snapshot/weights.bin", tmp_path / "pool" with pytest.raises(RuntimeError, match="changed during checkpoint saving"): - checkpointing._copy_checkpoint_asset(source, destination, pool, digest_cache={}) - assert not destination.exists() - assert not list(pool.iterdir()) - - -def test_corrupted_initial_copy_is_rejected_before_pool_publication( - tmp_path, monkeypatch -): - source = tmp_path / "weights.bin" - source.write_bytes(b"original") - copy2 = checkpointing.shutil.copy2 - - def corrupt_copy(src, dst, *args, **kwargs): - result = copy2(src, dst, *args, **kwargs) - Path(dst).write_bytes(b"corrupt!") - return result - - monkeypatch.setattr(checkpointing.shutil, "copy2", corrupt_copy) - destination, pool = tmp_path / "snapshot/weights.bin", tmp_path / "pool" - with pytest.raises(RuntimeError, match="changed during checkpoint saving"): - checkpointing._copy_checkpoint_asset(source, destination, pool, digest_cache={}) + checkpoint_extras._copy_checkpoint_asset( + source, destination, pool, digest_cache={} + ) assert not destination.exists() assert not list(pool.iterdir()) @@ -546,7 +483,7 @@ def test_cached_source_change_during_link_is_rejected(tmp_path, monkeypatch): source = tmp_path / "weights.bin" source.write_bytes(b"original") cache, pool = {}, tmp_path / "pool" - checkpointing._copy_checkpoint_asset( + checkpoint_extras._copy_checkpoint_asset( source, tmp_path / "first/weights.bin", pool, digest_cache=cache ) destination = tmp_path / "second/weights.bin" @@ -560,7 +497,7 @@ def change_during_link(src, dst, *args, **kwargs): monkeypatch.setattr(os, "link", change_during_link) with pytest.raises(RuntimeError, match="changed during checkpoint saving"): - checkpointing._copy_checkpoint_asset( + checkpoint_extras._copy_checkpoint_asset( source, destination, pool, digest_cache=cache ) assert not destination.exists() @@ -571,7 +508,7 @@ def test_existing_blob_size_and_readonly_guards_remain(tmp_path, damage): source = tmp_path / "weights.bin" source.write_bytes(b"original") cache, pool = {}, tmp_path / "pool" - checkpointing._copy_checkpoint_asset( + checkpoint_extras._copy_checkpoint_asset( source, tmp_path / "first/weights.bin", pool, digest_cache=cache ) blob = pool / hashlib.sha256(b"original").hexdigest() @@ -580,7 +517,7 @@ def test_existing_blob_size_and_readonly_guards_remain(tmp_path, damage): blob.write_bytes(b"short") blob.chmod(0o444) with pytest.raises(ValueError, match="damaged blob"): - checkpointing._copy_checkpoint_asset( + checkpoint_extras._copy_checkpoint_asset( source, tmp_path / "second/weights.bin", pool, digest_cache=cache ) diff --git a/tests/test_checkpoint_files.py b/tests/test_checkpoint_files.py index 5a7f8a3..3b7f095 100644 --- a/tests/test_checkpoint_files.py +++ b/tests/test_checkpoint_files.py @@ -85,21 +85,21 @@ def test_public_save_preserves_rng_and_load_restores_state(tmp_path): assert_nested_equal(actual_random, expected_random) -@pytest.mark.parametrize("failure", ["write", "publish"]) +@pytest.mark.parametrize("failure", ["write", "sync", "publish"]) def test_failed_checkpoint_write_keeps_existing_checkpoint_loadable( tmp_path, monkeypatch, failure ): source = TinyTrainer() advance_models(source) - checkpoint = tmp_path / "checkpoint" - source.save_checkpoint(checkpoint) + destination = tmp_path / "checkpoint" + source.save_checkpoint(destination) saved_files = { - path.relative_to(checkpoint): path.read_bytes() - for path in checkpoint.rglob("*") + path.relative_to(destination): path.read_bytes() + for path in destination.rglob("*") if path.is_file() } expected = TinyTrainer() - expected.load_checkpoint(checkpoint) + expected.load_checkpoint(destination) advance_models(source) def fail_save(*args, **kwargs): @@ -110,7 +110,7 @@ def fail_save(*args, **kwargs): def fail_publish(source_path, target_path): nonlocal publish_failed - if Path(target_path) == checkpoint and not publish_failed: + if Path(target_path) == destination and not publish_failed: publish_failed = True raise OSError("simulated storage failure") return original_replace(source_path, target_path) @@ -118,18 +118,21 @@ def fail_publish(source_path, target_path): with monkeypatch.context() as patch: if failure == "write": patch.setattr(torch, "save", fail_save) + elif failure == "sync": + patch.setattr(checkpoint, "_fsync_tree", fail_save) else: patch.setattr(os, "replace", fail_publish) with pytest.raises(OSError, match="simulated storage failure"): - source.save_checkpoint(checkpoint) + source.save_checkpoint(destination) assert { - path.relative_to(checkpoint): path.read_bytes() - for path in checkpoint.rglob("*") + path.relative_to(destination): path.read_bytes() + for path in destination.rglob("*") if path.is_file() } == saved_files + assert not list(tmp_path.glob(".checkpoint-*")) restored = TinyTrainer() - restored.load_checkpoint(checkpoint) + restored.load_checkpoint(destination) assert_models_equal(restored, expected) @@ -167,23 +170,6 @@ def replace(source, destination): assert events[-1] == ("sync", tmp_path.stat().st_ino) -def test_asset_sync_failure_does_not_replace_previous_checkpoint(tmp_path, monkeypatch): - trainer = TinyTrainer() - destination = tmp_path / "checkpoint" - trainer.save_checkpoint(destination) - previous = (destination / checkpoint.STATE_FILENAME).read_bytes() - monkeypatch.setattr( - checkpoint, - "_fsync_tree", - lambda root: (_ for _ in ()).throw(OSError("sync failed")), - ) - trainer.env_step += 1 - with pytest.raises(OSError, match="sync failed"): - trainer.save_checkpoint(destination) - assert (destination / checkpoint.STATE_FILENAME).read_bytes() == previous - assert not list(tmp_path.glob(".checkpoint-*")) - - def test_public_load_rejects_unallowlisted_pickle_without_mutation(tmp_path): target = TinyTrainer() before = capture_trainer_state(target) @@ -196,7 +182,7 @@ def test_public_load_rejects_unallowlisted_pickle_without_mutation(tmp_path): assert_nested_equal(capture_trainer_state(target), before) -def test_public_load_rejects_unallowlisted_record_without_mutation(tmp_path): +def test_public_load_rejects_unregistered_record_without_mutation(tmp_path): target = TinyTrainer() before = capture_trainer_state(target) path = tmp_path / "unknown-record.pt" @@ -260,19 +246,6 @@ def test_invalid_retention_limit_rejected_before_training(tmp_path, limit): assert_nested_equal(capture_trainer_state(trainer), before) -def test_automatic_save_does_not_overwrite_foreign_run(tmp_path): - first, second = TinyTrainer(), TinyTrainer() - with first._checkpoint_context(checkpoint_dir=tmp_path): - first._checkpoint_commit({"epoch": 1}) - saved = (tmp_path / "checkpoint-1" / checkpoint.STATE_FILENAME).read_bytes() - with pytest.raises(FileExistsError, match="another run"): - with second._checkpoint_context( - checkpoint_dir=tmp_path, checkpoint_total_limit=1 - ): - second._checkpoint_commit({"epoch": 1}) - assert (tmp_path / "checkpoint-1" / checkpoint.STATE_FILENAME).read_bytes() == saved - - def test_automatic_save_expands_user_directory(tmp_path, monkeypatch): original = Path.expanduser monkeypatch.setattr( @@ -301,10 +274,19 @@ def test_new_run_rejects_any_existing_snapshot_before_mutation(tmp_path, kind): @pytest.mark.parametrize( - "kind", ["foreign", "legacy", "bad-sequence", "file", "symlink"] + "kind, match", + [ + ("foreign", "another run"), + ("legacy", "Legacy or unowned.*ownership metadata"), + ("bad-sequence", "inconsistent ownership metadata"), + ("file", "not a checkpoint directory"), + ("symlink", "not a checkpoint directory"), + ], ) @pytest.mark.parametrize("source_is_file", [False, True]) -def test_resume_rejects_future_conflicts_before_loading(tmp_path, kind, source_is_file): +def test_resume_rejects_future_conflicts_before_loading( + tmp_path, kind, match, source_is_file +): source_root = tmp_path / "source" source = _save_run(source_root) output = tmp_path / "output" @@ -314,7 +296,7 @@ def test_resume_rejects_future_conflicts_before_loading(tmp_path, kind, source_i if source_is_file: path = path / checkpoint.STATE_FILENAME trainer, attributes, state = _untouched_trainer() - with pytest.raises(FileExistsError): + with pytest.raises(FileExistsError, match=match): with trainer._checkpoint_context( resume_from_checkpoint=path, checkpoint_dir=output ): @@ -322,20 +304,6 @@ def test_resume_rejects_future_conflicts_before_loading(tmp_path, kind, source_i _assert_untouched(trainer, attributes, state) -def test_legacy_destination_reports_missing_ownership_before_restore(tmp_path): - source_root = tmp_path / "source" - _save_run(source_root) - output = tmp_path / "output" - _snapshot(output, 2, kind="legacy") - trainer, attributes, state = _untouched_trainer() - with pytest.raises(FileExistsError, match="Legacy or unowned.*ownership metadata"): - with trainer._checkpoint_context( - resume_from_checkpoint=source_root / "checkpoint-1", checkpoint_dir=output - ): - pytest.fail("legacy destination was not rejected") - _assert_untouched(trainer, attributes, state) - - def test_same_run_future_snapshots_can_be_replaced_on_resume(tmp_path): source = _save_run(tmp_path, count=3) trainer = TinyTrainer() diff --git a/tests/test_checkpoint_metadata.py b/tests/test_checkpoint_metadata.py index 3b45312..092f6d1 100644 --- a/tests/test_checkpoint_metadata.py +++ b/tests/test_checkpoint_metadata.py @@ -100,11 +100,9 @@ def load(*args, **kwargs): monkeypatch.setattr( wandb, "init", lambda **kwargs: pytest.fail("metadata initialized W&B") ) - # No training records, callbacks or model classes need decoding/importing. + # Metadata needs only primitives, so no training state is decoded. monkeypatch.setattr( - checkpoint.importlib, - "import_module", - lambda *a: pytest.fail("decoded training record"), + checkpoint, "_decode", lambda *a: pytest.fail("decoded training state") ) metadata = read_checkpoint_metadata(state_path if file_path else path) assert metadata == { diff --git a/tests/test_checkpoint_models.py b/tests/test_checkpoint_models.py index dccc29b..3bb3713 100644 --- a/tests/test_checkpoint_models.py +++ b/tests/test_checkpoint_models.py @@ -9,8 +9,6 @@ from comlrl.trainers.actor_critic import IACConfig, IACTrainer, MAACConfig, MAACTrainer from comlrl.trainers.reinforce import ( - CentralizedMAGRPOConfig, - CentralizedMAGRPOTrainer, MAGRPOConfig, MAGRPOTrainer, MAREINFORCEConfig, @@ -134,10 +132,12 @@ def make_trainer(variant, tokenizer, seed, device): formatters=format_prompt, ) if variant == "centralized-magrpo": - trainer = CentralizedMAGRPOTrainer( + trainer = MAGRPOTrainer( agents=[model_for(tokenizer)], num_agents=2, - args=CentralizedMAGRPOConfig(**common), + args=MAGRPOConfig( + **{**common, "collaboration_mode": "centralized", "num_turns": 1} + ), centralized_adapter=SplitRoles(), **kwargs, ) diff --git a/tests/test_checkpoint_nested.py b/tests/test_checkpoint_nested.py index 44e42d4..37a0528 100644 --- a/tests/test_checkpoint_nested.py +++ b/tests/test_checkpoint_nested.py @@ -7,7 +7,7 @@ import pytest import torch -from comlrl.trainers.preference import marlhf +from comlrl.trainers.pbrl import marlhf from comlrl.utils.checkpoint import ( STATE_FILENAME, _decode, diff --git a/tests/test_checkpoint_preference.py b/tests/test_checkpoint_preference.py index 24e7400..c153cd7 100644 --- a/tests/test_checkpoint_preference.py +++ b/tests/test_checkpoint_preference.py @@ -7,14 +7,14 @@ import torch import wandb -from comlrl.trainers.preference import ( +from comlrl.trainers.pbrl import ( MADPOIterTrainer, MADPOTrainer, MARLHFIterTrainer, MARLHFTrainer, marlhf, ) -from comlrl.trainers.preference.checkpointing import train_iterations_checkpointed +from comlrl.trainers.pbrl.checkpoint_extras import train_iterations_checkpointed from checkpoint_helpers import ( Interrupted, @@ -116,22 +116,6 @@ def test_madpo_commits_generated_pairs_before_first_update_and_large_interval(tm assert_nested_equal(rng_state(), expected_rng) -def test_madpo_first_update_failure_leaves_generated_pairs_resumable(tmp_path): - trainer = toy_trainer() - - def fail(batch): - raise Interrupted("first optimizer update failed") - - trainer._update_from_preference_batch = fail - with pytest.raises(Interrupted): - trainer.train(checkpoint_dir=tmp_path) - resumed = toy_trainer() - progress = resumed.load_checkpoint(tmp_path / "checkpoint-1") - assert progress["epoch"] == progress["updates_seen"] == 0 - assert len(progress["preference_pairs"]) == 6 - assert not resumed.optimizers[0].state - - def test_madpo_empty_generation_has_one_complete_boundary(tmp_path): trainer = toy_trainer() trainer._build_preference_dataset = lambda **kwargs: [] @@ -377,7 +361,7 @@ def test_plain_preference_training_does_not_retain_checkpoint_data(tmp_path, kin trainer = toy_trainer(MARLHFIterTrainer) trainer.args.num_iterations = 0 trainer.train() - assert trainer._checkpoint_untracked_preference + assert trainer._checkpoint_untracked assert not getattr(trainer, "_checkpoint_progress", {}).get("preference_pairs") with pytest.raises(RuntimeError, match="checkpoint_dir"): trainer.save_checkpoint(tmp_path / "manual") @@ -386,7 +370,7 @@ def test_plain_preference_training_does_not_retain_checkpoint_data(tmp_path, kin def test_extra_state_rejects_arbitrary_attributes(): trainer = toy_trainer() with pytest.raises(ValueError, match="extra-state"): - trainer._restore_checkpoint_extra_state({"agents": []}, ".") + trainer._checkpoint_restore_extra_state({"agents": []}, ".") assert len(trainer.agents) == 1 @@ -404,14 +388,13 @@ def test_missing_replay_asset_rejected_before_model_mutation(tmp_path): assert torch.equal(before, resumed.agents[0].weight) -@pytest.mark.parametrize("value", [True, "1", 0.5, -1]) +@pytest.mark.parametrize("value", [True, "1", 0.5, -1, None]) @pytest.mark.parametrize("kind", ["madpo", "marlhf", "madpo_iter"]) -def test_preference_checkpoint_cursors_reject_noninteger_or_negative_values( - tmp_path, kind, value -): +def test_preference_checkpoint_cursors_reject_invalid_progress(tmp_path, kind, value): + """Cursors must be nonnegative integers; ``None`` leaves only the progress kind.""" if kind == "madpo": trainer = toy_trainer() - trainer._checkpoint_progress = { + progress = { "kind": kind, "epoch": value, "updates_seen": 0, @@ -421,7 +404,7 @@ def test_preference_checkpoint_cursors_reject_noninteger_or_negative_values( run = trainer._train_madpo_checkpointed elif kind == "marlhf": trainer = marlhf_trainer() - trainer._checkpoint_progress = { + progress = { "kind": kind, "phase": "online_rl", "rl": {"kind": "magrpo", "epoch": value}, @@ -429,7 +412,7 @@ def test_preference_checkpoint_cursors_reject_noninteger_or_negative_values( run = trainer._train_marlhf_checkpointed else: trainer = iterative_trainer(tmp_path) - trainer._checkpoint_progress = { + progress = { "kind": kind, "iteration": value, "updates_seen": 0, @@ -437,22 +420,7 @@ def test_preference_checkpoint_cursors_reject_noninteger_or_negative_values( "complete": False, } run = lambda: train_iterations_checkpointed(trainer, reward_model=False) - with pytest.raises(ValueError): - run() - - -@pytest.mark.parametrize("kind", ["madpo", "marlhf", "madpo_iter"]) -def test_preference_checkpoint_cursors_require_complete_shape(tmp_path, kind): - if kind == "madpo": - trainer = toy_trainer() - run = trainer._train_madpo_checkpointed - elif kind == "marlhf": - trainer = marlhf_trainer() - run = trainer._train_marlhf_checkpointed - else: - trainer = iterative_trainer(tmp_path) - run = lambda: train_iterations_checkpointed(trainer, reward_model=False) - trainer._checkpoint_progress = {"kind": kind} + trainer._checkpoint_progress = {"kind": kind} if value is None else progress with pytest.raises(ValueError): run() diff --git a/tests/test_checkpoint_preference_specials.py b/tests/test_checkpoint_preference_specials.py index 963561d..bb31a08 100644 --- a/tests/test_checkpoint_preference_specials.py +++ b/tests/test_checkpoint_preference_specials.py @@ -5,7 +5,7 @@ import pytest import torch -from comlrl.trainers.preference import ( +from comlrl.trainers.pbrl import ( MADPOConfig, MADPOIterConfig, MADPOIterTrainer, diff --git a/tests/test_checkpoint_rng.py b/tests/test_checkpoint_rng.py index c7663b2..df94c6f 100644 --- a/tests/test_checkpoint_rng.py +++ b/tests/test_checkpoint_rng.py @@ -11,7 +11,7 @@ import pytest import torch -from comlrl.trainers.preference.iterative import MADPOIterTrainer +from comlrl.trainers.pbrl.iterative import MADPOIterTrainer import comlrl.utils.checkpoint as checkpoint from comlrl.utils.checkpoint import capture_trainer_state, restore_trainer_state @@ -244,19 +244,6 @@ def test_nested_rng_routes_restore_latest_parent_rng_values(monkeypatch): assert [value.item() for value in restored[0]["cuda"]] == [11, 22, 22] -def test_nested_rng_routes_cannot_conflict_with_parent_streams(monkeypatch): - trainer = SimpleNamespace(_checkpoint_rng_devices_override={"agents.0": "cuda:0"}) - state = { - "rng": {"cuda": [torch.tensor([11]), torch.tensor([22])]}, - "rng_devices": {"agents.0": "cuda:0"}, - } - nested = [{"saved": {"critics.0": "cuda:1"}, "current": {"critics.0": "cuda:0"}}] - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) - with pytest.raises(ValueError, match="merge or split"): - checkpoint._cuda_rng_routes(trainer, state, nested) - - @pytest.mark.parametrize( "seed, target, comparator", [(7, 7, -7), (0, 0, 1), (2**64 - 5, -5, 5)] ) diff --git a/tests/test_checkpoint_state.py b/tests/test_checkpoint_state.py index 73d22c3..a0bd8b9 100644 --- a/tests/test_checkpoint_state.py +++ b/tests/test_checkpoint_state.py @@ -1,17 +1,16 @@ """In-memory trainer state capture/restore; failures never partially mutate.""" import copy +from dataclasses import dataclass, field import numpy as np import pytest from tokenizers.models import WordLevel import torch -from comlrl.trainers.reinforce.magrpo import NodeSample +from comlrl.trainers.reinforce.magrpo import MAGRPOConfig, NodeSample import comlrl.utils.checkpoint as checkpoint from comlrl.utils.checkpoint import ( - _decode, - _encode, capture_trainer_state, restore_trainer_state, ) @@ -155,20 +154,6 @@ def test_restored_training_does_not_mutate_reusable_in_memory_snapshot(tmp_path) assert_nested_equal(state, before) -def test_state_is_compatible_with_weights_only_loading(tmp_path): - trainer = TinyTrainer() - advance_models(trainer) - path = tmp_path / "state.pt" - state = capture_trainer_state(trainer) - torch.save(_encode(state), path) - - loaded = _decode(torch.load(path, weights_only=True, map_location="cpu")) - - restored = TinyTrainer() - restore_trainer_state(restored, loaded, tmp_path) - assert_models_equal(restored, trainer) - - def test_stateful_callback_restores_its_state(tmp_path): source = TinyTrainer() source.reward_func.calls = 19 @@ -457,6 +442,51 @@ def test_same_shape_model_metadata_mismatch_precedes_any_mutation( assert_nested_equal(capture_trainer_state(target), before) +@pytest.mark.parametrize("value,resumes", [(None, True), (4, False)]) +def test_config_field_newer_than_the_checkpoint_resumes_at_its_default( + tmp_path, value, resumes +): + state = capture_trainer_state(TinyTrainer()) + # A checkpoint written before this field existed has no entry for it. + del state["config"]["train_batch_size"] + target = TinyTrainer() + target.args.train_batch_size = value + if resumes: + restore_trainer_state(target, state, tmp_path) + else: + with pytest.raises(ValueError, match="does not match: train_batch_size"): + restore_trainer_state(target, state, tmp_path) + + +def test_config_defaults_are_built_only_for_fields_the_checkpoint_lacks(tmp_path): + calls = [] + + def factory(): + calls.append(torch.rand(1)) # draws from the global RNG + return [] + + @dataclass + class FactoryConfig(MAGRPOConfig): + extra_items: list = field(default_factory=factory) + + def trainer(): + target = TinyTrainer() + target.args = FactoryConfig(num_agents=2, num_turns=1, agent_devices="cpu") + return target + + state = capture_trainer_state(trainer()) + target = trainer() # building the config itself runs the factory + calls.clear() + restore_trainer_state(target, state, tmp_path, restore_rng=False) + assert calls == [] + del state["config"]["extra_items"] + target = trainer() + calls.clear() + before = torch.get_rng_state() + restore_trainer_state(target, state, tmp_path, restore_rng=False) + assert len(calls) == 1 and torch.equal(torch.get_rng_state(), before) + + @pytest.mark.parametrize("changed", ["vocab", "chat_template", "eos"]) def test_tokenizer_metadata_mismatch_precedes_any_mutation( tmp_path, tokenizer, changed diff --git a/tests/test_checkpoint_tokenizer_specials.py b/tests/test_checkpoint_tokenizer_specials.py index 90a27fb..28ea090 100644 --- a/tests/test_checkpoint_tokenizer_specials.py +++ b/tests/test_checkpoint_tokenizer_specials.py @@ -8,8 +8,6 @@ from transformers import GPT2Config, GPT2LMHeadModel from comlrl.trainers.reinforce import ( - CentralizedMAGRPOConfig, - CentralizedMAGRPOTrainer, MAGRPOConfig, MAGRPOTrainer, ) @@ -88,8 +86,10 @@ def _make_trainer(variant="magrpo", *, eos_mismatch=False, seed=37): formatters=lambda item: item["prompt"], ) if variant == "centralized": - trainer = CentralizedMAGRPOTrainer( - args=CentralizedMAGRPOConfig(**common), + trainer = MAGRPOTrainer( + args=MAGRPOConfig( + **{**common, "collaboration_mode": "centralized", "num_turns": 1} + ), centralized_adapter=_SplitRoles(), **kwargs, ) diff --git a/tests/test_config.py b/tests/test_config.py index e26d125..9f270b4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4,8 +4,8 @@ from comlrl.trainers.actor_critic.iac import IACConfig from comlrl.trainers.actor_critic.maac import MAACConfig -from comlrl.trainers.preference.madpo import MADPOConfig -from comlrl.trainers.preference.marlhf import MARLHFConfig +from comlrl.trainers.pbrl.madpo import MADPOConfig +from comlrl.trainers.pbrl.marlhf import MARLHFConfig from comlrl.trainers.reinforce.magrpo import MAGRPOConfig ACTOR_CRITIC = (IACConfig, MAACConfig) @@ -40,11 +40,18 @@ def _invalid(config_cls, value, *fields, match=None, **options): "rollout_buffer_size", "train_batch_size", "num_agents", - "num_generations", "num_turns", "eval_batch_size", "logging_steps", ), + # Single-turn, so the lower bound is checked rather than the multi-turn rule. + *_invalid( + config_cls, + 0, + "num_generations", + match="num_generations must be >= 1", + num_turns=1, + ), *_invalid(config_cls, -1, "eval_interval", "eval_num_samples"), *_invalid(config_cls, "x", "critic_type"), *_invalid(config_cls, "invalid", "parallel_training"), @@ -109,7 +116,6 @@ def _invalid(config_cls, value, *fields, match=None, **options): for config_cls in ACTOR_CRITIC for options in [ {}, - {"parallel_training": "none"}, {"parallel_training": "null"}, { "parallel_training": "mp", @@ -117,7 +123,6 @@ def _invalid(config_cls, value, *fields, match=None, **options): "critic_devices": "cpu", }, {"agent_devices": "cpu", "critic_devices": "cpu"}, - {"num_turns": 2, "num_generations": 1}, {"critic_type": "q"}, ] ], @@ -125,7 +130,6 @@ def _invalid(config_cls, value, *fields, match=None, **options): (MAGRPOConfig, options) for options in [ {}, - {"parallel_training": "none"}, {"parallel_training": "null"}, {"parallel_training": "mp", "agent_devices": "cpu"}, {"agent_devices": "cpu"}, diff --git a/tests/test_joint_policy.py b/tests/test_joint_policy.py index 8399e12..2d321b4 100644 --- a/tests/test_joint_policy.py +++ b/tests/test_joint_policy.py @@ -6,10 +6,12 @@ import pytest import torch -from comlrl.trainers.preference import MARLHFConfig, MARLHFTrainer -from comlrl.trainers.preference.centralized import TaggedCentralizedComparatorAdapter -from comlrl.trainers.preference.collaboration import joint_sequence_log_prob -from comlrl.trainers.reinforce import CentralizedMAGRPOConfig, CentralizedMAGRPOTrainer +from comlrl.trainers.pbrl import MARLHFConfig, MARLHFTrainer +from comlrl.trainers.centralized import ( + TaggedCentralizedComparatorAdapter, + joint_sequence_log_prob, +) +from comlrl.trainers.reinforce import MAGRPOConfig, MAGRPOTrainer from helpers import ( BAD, @@ -205,11 +207,13 @@ def reward(*roles, batch_items, **kwargs): observed = [] dataset = Dataset.from_list([{"prompt": "task"}]) - trainer = CentralizedMAGRPOTrainer( + trainer = MAGRPOTrainer( agents=[model], num_agents=num_roles, tokenizer=tokenizer, - args=CentralizedMAGRPOConfig( + args=MAGRPOConfig( + collaboration_mode="centralized", + num_turns=1, num_agents=num_roles, agent_devices="cpu", num_train_epochs=1, diff --git a/tests/test_magrpo.py b/tests/test_magrpo.py index a1ae3e4..fc550ac 100644 --- a/tests/test_magrpo.py +++ b/tests/test_magrpo.py @@ -224,3 +224,26 @@ def test_compute_rewards_wraps_completions_in_place(): ) == [2.0] assert completions == [["0:0:0"], ["0:1:1"]] assert completions[1] is second + + +@pytest.mark.parametrize("skip", [False, True]) +def test_zero_gradient_updates_can_skip_the_optimizer_step(tokenizer, skip): + from comlrl.trainers.reinforce.magrpo import NodeSample + from helpers import make_centralized_magrpo_trainer + + trainer, _ = make_centralized_magrpo_trainer(tokenizer, skip_zero_grad_steps=skip) + weights = trainer.agents[0].transformer.wte.weight + optimizer = trainer.optimizers[0] + steps = [] + original_step = optimizer.step + optimizer.step = lambda *args, **kwargs: steps.append(1) or original_step() + sample = NodeSample(0, 0, {}, [1.0, 1.0], 1.0, 1.0, 2) + # Equal returns leave every advantage, and so every gradient, at zero. + trainer._compute_loss_with_gradients = ( + lambda agent, data, returns: weights.sum() * 0 + ) + trainer._update_from_samples(0, [sample]) + assert steps == ([] if skip else [1]) + trainer._compute_loss_with_gradients = lambda agent, data, returns: weights.sum() + trainer._update_from_samples(0, [sample]) + assert steps == ([1] if skip else [1, 1]) diff --git a/tests/test_preference.py b/tests/test_preference.py index 8240efc..c013b74 100644 --- a/tests/test_preference.py +++ b/tests/test_preference.py @@ -9,7 +9,7 @@ import pytest import torch -from comlrl.trainers.preference import ( +from comlrl.trainers.pbrl import ( MADPOConfig, MADPOIterTrainer, MADPOTrainer, diff --git a/tests/test_reference_kl.py b/tests/test_reference_kl.py index be6607e..70c4605 100644 --- a/tests/test_reference_kl.py +++ b/tests/test_reference_kl.py @@ -1,17 +1,17 @@ """Reference-model KL penalties for MAGRPO, IAC and MAAC.""" -import importlib from types import SimpleNamespace import pytest import torch from transformers import GPT2LMHeadModel -from comlrl.trainers.actor_critic import IACTrainer +from comlrl.trainers.actor_critic import ActorCriticTrainerBase, IACTrainer from comlrl.trainers.actor_critic.iac import IACConfig from comlrl.trainers.actor_critic.maac import MAACConfig, MAACTrainer from comlrl.trainers.reinforce import MAGRPOTrainer from comlrl.trainers.reinforce.magrpo import MAGRPOConfig +from comlrl.utils import reference_kl from comlrl.utils.reference_kl import ( reference_kl_for_sequence, resolve_reference_devices, @@ -23,24 +23,6 @@ TRAINERS = [IACTrainer, MAACTrainer, MAGRPOTrainer] -def _actor_critic_trainer(trainer_cls, tokenizer, **options): - options.update(num_agents=1, num_turns=1, agent_devices="cpu", critic_devices="cpu") - if trainer_cls is IACTrainer: - return IACTrainer( - agents=[tiny_gpt2()], - tokenizer=tokenizer, - reward_func=zero_reward, - args=IACConfig(use_separate_critic=False, **options), - ) - return MAACTrainer( - agents=[tiny_gpt2()], - critics=[tiny_gpt2()], - tokenizer=tokenizer, - reward_func=zero_reward, - args=MAACConfig(**options), - ) - - @pytest.mark.parametrize("config_cls", CONFIGS) def test_reference_kl_config_defaults_off(config_cls): config = config_cls() @@ -136,29 +118,18 @@ def test_reference_kl_for_identical_model_is_zero(): assert torch.allclose(kl, torch.zeros_like(kl), atol=1e-6) -@pytest.mark.parametrize("trainer_cls", [IACTrainer, MAACTrainer]) -def test_actor_critic_kl_shaping_is_noop_when_disabled(trainer_cls, stub_tokenizer): - trainer = _actor_critic_trainer(trainer_cls, stub_tokenizer) +@pytest.mark.parametrize("enabled", [False, True]) +def test_actor_critic_kl_shaping(enabled): + # IAC and MAAC both use the base-class method. + trainer = ActorCriticTrainerBase.__new__(ActorCriticTrainerBase) + trainer.args = SimpleNamespace(reference_kl_enabled=enabled, reference_kl_coef=0.25) shaped_reward, metadata = trainer._kl_shaped_reward( 1.0, {"reference_kls": [0.4]}, 0 ) - assert shaped_reward == pytest.approx(1.0) - assert metadata == {} - - -@pytest.mark.parametrize("trainer_cls", [IACTrainer, MAACTrainer]) -def test_actor_critic_kl_shaping_penalizes_reward(trainer_cls, stub_tokenizer): - trainer = _actor_critic_trainer( - trainer_cls, - stub_tokenizer, - reference_kl_enabled=True, - reference_kl_coef=0.25, - ) - - shaped_reward, metadata = trainer._kl_shaped_reward( - 1.0, {"reference_kls": [0.4]}, 0 - ) - + if not enabled: + assert shaped_reward == pytest.approx(1.0) + assert metadata == {} + return assert shaped_reward == pytest.approx(0.9) assert metadata["reference_kl"] == pytest.approx(0.4) assert metadata["reference_kl_penalty"] == pytest.approx(0.1) @@ -187,24 +158,26 @@ def test_magrpo_applies_reference_kl_to_returns(stub_tokenizer): assert torch.allclose(adjusted, torch.tensor([0.9, 1.8])) -@pytest.mark.parametrize("trainer_cls", TRAINERS) -def test_reference_kl_values_disabled_or_missing_models(trainer_cls): - trainer = trainer_cls.__new__(trainer_cls) - trainer.args = SimpleNamespace(reference_kl_enabled=False) - assert trainer._reference_kl_values(0, None, None, None, 2, []) == [] +def test_trainers_share_one_reference_kl_helper(): + assert all( + cls._reference_kl_values is reference_kl.reference_kl_values for cls in TRAINERS + ) + + +def test_reference_kl_values_disabled_or_missing_models(): + trainer = SimpleNamespace(args=SimpleNamespace(reference_kl_enabled=False)) + assert reference_kl.reference_kl_values(trainer, 0, None, None, None, 2, []) == [] trainer.args.reference_kl_enabled = True with pytest.raises(RuntimeError, match="reference models are missing"): - trainer._reference_kl_values(0, None, None, None, 2, []) + reference_kl.reference_kl_values(trainer, 0, None, None, None, 2, []) -@pytest.mark.parametrize("trainer_cls", TRAINERS) @pytest.mark.parametrize("response_lens", [[], [2], [2, 1, 3]]) -def test_reference_kl_values_keep_order_and_module_hook( - trainer_cls, response_lens, monkeypatch -): - trainer = trainer_cls.__new__(trainer_cls) - trainer.args = SimpleNamespace(reference_kl_enabled=True) - trainer.reference_models = [object(), object()] +def test_reference_kl_values_keep_order(response_lens, monkeypatch): + trainer = SimpleNamespace( + args=SimpleNamespace(reference_kl_enabled=True), + reference_models=[object(), object()], + ) policy = object() ids = torch.tensor([[2, 3, 4, 5], [6, 7, 8, 9]]) mask = torch.tensor([[1, 1, 1, 1], [1, 1, 1, 0]]) @@ -221,28 +194,28 @@ def sequence_kl( calls.append((prompt_len, response_len)) return torch.tensor([idx + 0.25, 100.0]) - module = importlib.import_module(trainer_cls.__module__) - monkeypatch.setattr(module, "reference_kl_for_sequence", sequence_kl) - values = trainer._reference_kl_values(1, policy, ids, mask, 2, response_lens) + monkeypatch.setattr(reference_kl, "reference_kl_for_sequence", sequence_kl) + values = reference_kl.reference_kl_values( + trainer, 1, policy, ids, mask, 2, response_lens + ) count = min(len(ids), len(response_lens)) assert values == [idx + 0.25 for idx in range(count)] assert calls == [(2, length) for length in response_lens[:count]] -@pytest.mark.parametrize("trainer_cls", TRAINERS) -def test_reference_kl_values_restore_training_modes_and_rng(trainer_cls): - trainer = trainer_cls.__new__(trainer_cls) - trainer.args = SimpleNamespace(reference_kl_enabled=True) +def test_reference_kl_values_restore_training_modes_and_rng(): policy, reference = tiny_gpt2(), tiny_gpt2() policy.train() reference.eval() - trainer.reference_models = [reference] + trainer = SimpleNamespace( + args=SimpleNamespace(reference_kl_enabled=True), reference_models=[reference] + ) ids = torch.tensor([[2, 3, 4, 5], [6, 7, 8, 9]]) mask = torch.ones_like(ids) rng_before = torch.get_rng_state().clone() - first = trainer._reference_kl_values(0, policy, ids, mask, 2, [2, 1]) - second = trainer._reference_kl_values(0, policy, ids, mask, 2, [2, 1]) + first = reference_kl.reference_kl_values(trainer, 0, policy, ids, mask, 2, [2, 1]) + second = reference_kl.reference_kl_values(trainer, 0, policy, ids, mask, 2, [2, 1]) assert first == second assert len(first) == 2 and all(value >= 0.0 for value in first) diff --git a/tests/test_trainer_init.py b/tests/test_trainer_init.py index 2003ee1..e5cdb5e 100644 --- a/tests/test_trainer_init.py +++ b/tests/test_trainer_init.py @@ -7,9 +7,9 @@ from comlrl.trainers.actor_critic import IACTrainer, MAACTrainer from comlrl.trainers.actor_critic.iac import IACConfig from comlrl.trainers.actor_critic.maac import MAACConfig -from comlrl.trainers.preference import MADPOTrainer, MARLHFTrainer -from comlrl.trainers.preference.madpo import MADPOConfig -from comlrl.trainers.preference.marlhf import MARLHFConfig +from comlrl.trainers.pbrl import MADPOTrainer, MARLHFTrainer +from comlrl.trainers.pbrl.madpo import MADPOConfig +from comlrl.trainers.pbrl.marlhf import MARLHFConfig from comlrl.trainers.reinforce import MAGRPOTrainer from comlrl.trainers.reinforce.magrpo import MAGRPOConfig @@ -289,17 +289,6 @@ def test_iac_multiturn_accepts_external_transition(stub_tokenizer, tiny_model_a) assert len(trainer.agents) == 1 -def test_iac_multiturn_rejects_multiple_generations(stub_tokenizer, tiny_model_a): - with pytest.raises(ValueError, match="num_generations"): - IACTrainer( - agents=[tiny_model_a], - tokenizer=stub_tokenizer, - reward_func=zero_reward, - external_transition=_external_transition, - args=_iac_cfg(num_agents=1, num_turns=2, num_generations=2), - ) - - def test_maac_loads_agent_and_critic_by_name(model_names): trainer = MAACTrainer( agent_model=model_names[0], @@ -359,18 +348,6 @@ def test_maac_multiturn_accepts_external_transition(stub_tokenizer, tiny_model_a assert len(trainer.critics) == 1 -def test_maac_multiturn_rejects_multiple_generations(stub_tokenizer, tiny_model_a): - with pytest.raises(ValueError, match="num_generations"): - MAACTrainer( - agents=[tiny_model_a], - critics=[tiny_model_a], - tokenizer=stub_tokenizer, - reward_func=zero_reward, - external_transition=_external_transition, - args=_maac_cfg(num_agents=1, num_turns=2, num_generations=2), - ) - - def test_magrpo_loads_agent_model_by_name(model_names): trainer = MAGRPOTrainer( agent_model=model_names[0], diff --git a/tests/test_utils.py b/tests/test_utils.py index 7be57b2..3bf1053 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,6 +6,7 @@ import torch from comlrl.utils import ( + DistributedContext, RewardProcessors, build_formatters, call_reward_function, @@ -158,7 +159,7 @@ def test_wandb_init_kwargs_keeps_every_config_section(): } -def test_local_context_defaults_to_single_process(): - ctx = local_context(torch.device("cpu")) - assert ctx.is_main is True - assert ctx.device == torch.device("cpu") +def test_local_context_is_a_deprecated_single_process_context(): + with pytest.warns(DeprecationWarning, match="deprecated"): + ctx = local_context(torch.device("cpu")) + assert ctx == DistributedContext(is_main=True, device=torch.device("cpu"))