From 6f50766a5375bc4bc0d387f10335fde2303d10f4 Mon Sep 17 00:00:00 2001 From: N!no Date: Sun, 18 Oct 2026 22:12:42 -0400 Subject: [PATCH 01/21] feat: add MAGRPO switches for zero-gradient steps and a shared loss skip_zero_grad_steps (MAGRPOConfig, default False) skips optimizer.step() when every gradient is zero, for example when all advantages in a batch are zero; AdamW would otherwise keep moving the weights along its stored momentum. match_decentralized_loss (CentralizedMAGRPOConfig, default False) scores the joint response with MAGRPOTrainer's decentralized loss (same token range, precision and non-finite fallback), so centralized and decentralized runs share one objective implementation. The defaults keep the current behavior. Co-Authored-By: Claude Opus 5.5 --- .../trainers/reinforce/centralized_magrpo.py | 8 +++++++ comlrl/trainers/reinforce/magrpo.py | 17 ++++++++++++++ .../docs/marl/multi-agent-reinforce.md | 1 + tests/test_centralized_magrpo.py | 21 +++++++++++++++++ tests/test_magrpo.py | 23 +++++++++++++++++++ 5 files changed, 70 insertions(+) diff --git a/comlrl/trainers/reinforce/centralized_magrpo.py b/comlrl/trainers/reinforce/centralized_magrpo.py index 69458b9..92faca5 100644 --- a/comlrl/trainers/reinforce/centralized_magrpo.py +++ b/comlrl/trainers/reinforce/centralized_magrpo.py @@ -9,6 +9,10 @@ class CentralizedMAGRPOConfig(MAGRPOConfig): num_turns: int = 1 collaboration_mode: str = "centralized" + # Score the joint response with MAGRPOTrainer's decentralized loss (the same + # token range, precision and non-finite fallback), so both modes share one + # objective implementation. The default keeps the joint-sequence loss. + match_decentralized_loss: bool = False def __post_init__(self): super().__post_init__() @@ -92,6 +96,10 @@ def _log_eval_metrics(self, all_agent_completions_turns, *args, **kwargs): return super()._log_eval_metrics(role_outputs, *args, **kwargs) def _compute_loss_with_gradients(self, agent, completions_data, returns): + if getattr(self.args, "match_decentralized_loss", False): + return MAGRPOTrainer._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..b1bfc7b 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -76,6 +76,9 @@ 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 reference_kl_enabled: bool = False reference_kl_coef: float = 0.1 reference_devices: Optional[Union[str, Sequence[str]]] = None @@ -128,6 +131,16 @@ class NodeSample: node_env_step: int +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(TrainingCheckpointMixin): """ Multi-Agent Group Relative Policy Optimization Trainer (MAGRPO). @@ -1585,6 +1598,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: diff --git a/docs/content/docs/marl/multi-agent-reinforce.md b/docs/content/docs/marl/multi-agent-reinforce.md index 89c1814..890bedc 100644 --- a/docs/content/docs/marl/multi-agent-reinforce.md +++ b/docs/content/docs/marl/multi-agent-reinforce.md @@ -75,6 +75,7 @@ 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`) - `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`) diff --git a/tests/test_centralized_magrpo.py b/tests/test_centralized_magrpo.py index 7fdf112..ba3a845 100644 --- a/tests/test_centralized_magrpo.py +++ b/tests/test_centralized_magrpo.py @@ -107,3 +107,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_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]) From 70c4bc25016f534a55fbe59bcb64652e2f7f39b8 Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 19:42:33 -0400 Subject: [PATCH 02/21] refactor: rename preference to pbrl and share centralized collaboration - Move the centralized adapters and CentralizedCollaboration out of the preference package into comlrl/trainers/centralized, and add CollaborationModeMixin, which every trainer family can use. - Replace CentralizedMAGRPOTrainer/Config with MAGRPOConfig.collaboration_mode ("decentralized" | "centralized"), plus the centralized_adapter trainer argument. MAREINFORCE, MAReMax and MARLOO inherit the mode. - MADPO and MARLHF reuse the MAGRPO collaboration setup instead of repeating it. Evaluation splits role outputs from the batch items that MAGRPO already collects, so _centralized_eval_items is gone. - Rename preference/ to pbrl/ and pbrl/checkpointing.py to pbrl/checkpoint_extras.py; export every trainer from comlrl.trainers. Decentralized and centralized MAGRPO trainings on tiny CPU models give bit-identical losses, rewards, evaluation metrics and parameters before and after this change. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- comlrl/trainers/__init__.py | 31 +++++- comlrl/trainers/centralized/__init__.py | 29 +++++ .../adapters.py} | 0 .../collaboration.py | 2 +- comlrl/trainers/centralized/mode.py | 82 ++++++++++++++ .../trainers/{preference => pbrl}/__init__.py | 8 -- .../checkpoint_extras.py} | 0 .../{preference => pbrl}/iterative.py | 12 +- comlrl/trainers/{preference => pbrl}/madpo.py | 66 +++-------- .../trainers/{preference => pbrl}/marlhf.py | 4 +- .../{preference => pbrl}/output_paths.py | 0 comlrl/trainers/reinforce/__init__.py | 3 - .../trainers/reinforce/centralized_magrpo.py | 105 ------------------ comlrl/trainers/reinforce/magrpo.py | 49 ++++++-- comlrl/utils/checkpoint.py | 8 +- .../docs/marl/multi-agent-reinforce.md | 5 +- docs/content/docs/marl/multi-llm-pbrl.md | 2 +- .../docs/user-guide/training-entities.md | 2 +- .../user-guide/training-parallelization.md | 2 +- tests/checkpoint_helpers.py | 4 +- tests/helpers.py | 16 ++- tests/test_centralized_comparator.py | 4 +- tests/test_centralized_magrpo.py | 23 ++-- tests/test_centralized_preference.py | 7 +- tests/test_checkpoint_assets.py | 62 ++++++----- tests/test_checkpoint_models.py | 8 +- tests/test_checkpoint_nested.py | 2 +- tests/test_checkpoint_preference.py | 4 +- tests/test_checkpoint_preference_specials.py | 2 +- tests/test_checkpoint_rng.py | 2 +- tests/test_checkpoint_tokenizer_specials.py | 8 +- tests/test_config.py | 4 +- tests/test_joint_policy.py | 16 ++- tests/test_preference.py | 2 +- tests/test_trainer_init.py | 6 +- 36 files changed, 299 insertions(+), 283 deletions(-) create mode 100644 comlrl/trainers/centralized/__init__.py rename comlrl/trainers/{preference/centralized.py => centralized/adapters.py} (100%) rename comlrl/trainers/{preference => centralized}/collaboration.py (99%) create mode 100644 comlrl/trainers/centralized/mode.py rename comlrl/trainers/{preference => pbrl}/__init__.py (63%) rename comlrl/trainers/{preference/checkpointing.py => pbrl/checkpoint_extras.py} (100%) rename comlrl/trainers/{preference => pbrl}/iterative.py (99%) rename comlrl/trainers/{preference => pbrl}/madpo.py (90%) rename comlrl/trainers/{preference => pbrl}/marlhf.py (99%) rename comlrl/trainers/{preference => pbrl}/output_paths.py (100%) delete mode 100644 comlrl/trainers/reinforce/centralized_magrpo.py diff --git a/README.md b/README.md index c74e4d5..205fd7b 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ conda install -c conda-forge comlrl - **_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). + - **_Multi-Agent Preference Learning:_** [MADPO](comlrl/trainers/pbrl/madpo.py), [MARLHF](comlrl/trainers/pbrl/marlhf.py), and their [iterative variants](comlrl/trainers/preference/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/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..4f1bc91 --- /dev/null +++ b/comlrl/trainers/centralized/mode.py @@ -0,0 +1,82 @@ +"""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_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 100% rename from comlrl/trainers/preference/checkpointing.py rename to comlrl/trainers/pbrl/checkpoint_extras.py diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/pbrl/iterative.py similarity index 99% rename from comlrl/trainers/preference/iterative.py rename to comlrl/trainers/pbrl/iterative.py index 09f78b1..a0ee5cf 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -33,7 +33,7 @@ MADPOTrainer, PreferencePair, ) -from .centralized import CentralizedComparatorAdapter +from ..centralized import CentralizedComparatorAdapter from .marlhf import MARLHFConfig, MARLHFTrainer from .output_paths import resolve_output_path @@ -348,7 +348,7 @@ class MADPOIterTrainer(MADPOTrainer): algorithm_name = "MADPOIter" def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress + from .checkpoint_extras import validate_preference_progress validate_preference_progress(self, state, "madpo_iter") @@ -623,7 +623,7 @@ def train( or checkpoint_interval != 1 or checkpoint_total_limit is not None ): - from .checkpointing import train_iterations_checkpointed + from .checkpoint_extras import train_iterations_checkpointed with self._checkpoint_context( resume_from_checkpoint=resume_from_checkpoint, @@ -2016,7 +2016,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. @@ -2860,7 +2860,7 @@ class MARLHFIterTrainer(MADPOIterTrainer, MARLHFTrainer): algorithm_name = "MARLHFIter" def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress + from .checkpoint_extras import validate_preference_progress validate_preference_progress(self, state, "marlhf_iter") @@ -2906,7 +2906,7 @@ def train( or checkpoint_interval != 1 or checkpoint_total_limit is not None ): - from .checkpointing import train_iterations_checkpointed + from .checkpoint_extras import train_iterations_checkpointed with self._checkpoint_context( resume_from_checkpoint=resume_from_checkpoint, diff --git a/comlrl/trainers/preference/madpo.py b/comlrl/trainers/pbrl/madpo.py similarity index 90% rename from comlrl/trainers/preference/madpo.py rename to comlrl/trainers/pbrl/madpo.py index cee83a2..e636589 100644 --- a/comlrl/trainers/preference/madpo.py +++ b/comlrl/trainers/pbrl/madpo.py @@ -15,11 +15,12 @@ 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 .collaboration import CentralizedCollaboration, joint_sequence_log_prob @dataclass @@ -87,8 +88,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 +191,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 +208,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, @@ -383,37 +343,37 @@ def commit_epoch(next_epoch, updates_seen, *, force=False): commit_epoch(start_epoch, updates_seen) def _checkpoint_extra_state(self): - from .checkpointing import preference_extra_state + from .checkpoint_extras 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 + from .checkpoint_extras 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 + from .checkpoint_extras 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 + from .checkpoint_extras 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 + from .checkpoint_extras import validate_preference_assets validate_preference_assets(self, checkpoint_dir, assets) def _checkpoint_validate_extra_state(self, state, checkpoint_dir): - from .checkpointing import validate_preference_extra_state + from .checkpoint_extras import validate_preference_extra_state validate_preference_extra_state(state) def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress + from .checkpoint_extras import validate_preference_progress validate_preference_progress(self, state, "madpo") diff --git a/comlrl/trainers/preference/marlhf.py b/comlrl/trainers/pbrl/marlhf.py similarity index 99% rename from comlrl/trainers/preference/marlhf.py rename to comlrl/trainers/pbrl/marlhf.py index abebc32..cbadd1d 100644 --- a/comlrl/trainers/preference/marlhf.py +++ b/comlrl/trainers/pbrl/marlhf.py @@ -15,7 +15,7 @@ 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 @@ -241,7 +241,7 @@ def train( self._train_online_rl(**kwargs) def _checkpoint_validate_progress(self, state): - from .checkpointing import validate_preference_progress + from .checkpoint_extras import validate_preference_progress validate_preference_progress(self, state, "marlhf") 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 92faca5..0000000 --- a/comlrl/trainers/reinforce/centralized_magrpo.py +++ /dev/null @@ -1,105 +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" - # Score the joint response with MAGRPOTrainer's decentralized loss (the same - # token range, precision and non-finite fallback), so both modes share one - # objective implementation. The default keeps the joint-sequence loss. - match_decentralized_loss: bool = False - - 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): - if getattr(self.args, "match_decentralized_loss", False): - return MAGRPOTrainer._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 b1bfc7b..c0d6ba2 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -42,6 +42,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: @@ -79,6 +85,13 @@ class MAGRPOConfig: # 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 @@ -102,6 +115,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: @@ -112,11 +128,7 @@ 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) @@ -141,7 +153,7 @@ def _has_gradient(optimizer) -> bool: ) -class MAGRPOTrainer(TrainingCheckpointMixin): +class MAGRPOTrainer(CollaborationModeMixin, TrainingCheckpointMixin): """ Multi-Agent Group Relative Policy Optimization Trainer (MAGRPO). Supports both single-turn and multi-turn training with external transitions. @@ -167,6 +179,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 @@ -192,8 +208,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") ) @@ -458,7 +484,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) @@ -676,6 +702,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 @@ -1242,7 +1271,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() @@ -1632,6 +1661,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..aa9808e 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -32,9 +32,9 @@ _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", + "comlrl.trainers.pbrl.madpo.PreferencePair", + "comlrl.trainers.pbrl.madpo.AgentPreferenceTensors", + "comlrl.trainers.pbrl.iterative.PreferenceReplayShard", } _MODEL_FIELDS = ("agents", "critics", "reference_models", "reward_model") _OPTIMIZER_FIELDS = ( @@ -1294,7 +1294,7 @@ def _checkpoint_prune(self, latest): shutil.rmtree(path) _fsync_directory(root) if callable(getattr(self, "_checkpoint_save_assets", None)): - from comlrl.trainers.preference.checkpointing import ( + from comlrl.trainers.pbrl.checkpoint_extras import ( prune_preference_asset_pool, ) diff --git a/docs/content/docs/marl/multi-agent-reinforce.md b/docs/content/docs/marl/multi-agent-reinforce.md index 890bedc..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 @@ -76,6 +77,7 @@ The current `MAGRPOTrainer` loss slice omits the final completion token, so its - `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`) @@ -100,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..50f4d0b 100644 --- a/docs/content/docs/marl/multi-llm-pbrl.md +++ b/docs/content/docs/marl/multi-llm-pbrl.md @@ -245,7 +245,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..58dff49 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 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_centralized_comparator.py b/tests/test_centralized_comparator.py index 3bd9bd6..cfd398f 100644 --- a/tests/test_centralized_comparator.py +++ b/tests/test_centralized_comparator.py @@ -2,11 +2,11 @@ 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: diff --git a/tests/test_centralized_magrpo.py b/tests/test_centralized_magrpo.py index ba3a845..e297d92 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): @@ -47,7 +43,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): @@ -87,10 +82,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" + ), ) diff --git a/tests/test_centralized_preference.py b/tests/test_centralized_preference.py index 3e9a179..0a5c827 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, @@ -201,7 +201,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..4b57ec0 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 ( @@ -398,14 +398,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 +413,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 +439,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,11 +457,11 @@ 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) @@ -476,15 +476,15 @@ def test_reused_blob_needs_no_byte_reads_and_changed_source_gets_new_blob( 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) + checkpoint_extras._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) + checkpoint_extras._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) + checkpoint_extras._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 @@ -499,14 +499,14 @@ 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): 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) @@ -514,10 +514,12 @@ def change_after_copy(src, dst, *args, **kwargs): 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={}) + checkpoint_extras._copy_checkpoint_asset( + source, destination, pool, digest_cache={} + ) assert not destination.exists() assert not list(pool.iterdir()) @@ -527,17 +529,19 @@ def test_corrupted_initial_copy_is_rejected_before_pool_publication( ): source = tmp_path / "weights.bin" source.write_bytes(b"original") - copy2 = checkpointing.shutil.copy2 + copy2 = checkpoint_extras.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) + monkeypatch.setattr(checkpoint_extras.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 +550,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 +564,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 +575,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 +584,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_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..f7db919 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, 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..235fbc5 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 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..90e9218 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) 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_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_trainer_init.py b/tests/test_trainer_init.py index 2003ee1..aed7839 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 From a24771fc2c8951ae279a9067e49e040c84cebbc6 Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 19:56:33 -0400 Subject: [PATCH 03/21] feat: add the centralized collaboration mode to IAC and MAAC IACConfig and MAACConfig gain collaboration_mode, and both trainers accept centralized_adapter. With collaboration_mode="centralized", one joint actor answers for all num_agents roles. Rewards and metrics_callback see the parsed role outputs. MAAC's shared critic reads each role's prompt and output in the same format as in decentralized training, and IAC's critic conditions on the joint prompt. ActorCriticTrainerBase now uses CollaborationModeMixin. Its num_agents property counts the actor models (args.num_agents, or one joint actor), and the actor-critic code uses it in place of args.num_agents. Decentralized IAC and MAAC trainings on tiny CPU models stay bit-identical. The docs and README list the centralized mode for every trainer family. Co-Authored-By: Claude Opus 5.5 --- README.md | 4 +- comlrl/trainers/actor_critic/ac_base.py | 26 ++- comlrl/trainers/actor_critic/iac.py | 83 ++++++--- comlrl/trainers/actor_critic/maac.py | 97 +++++++--- comlrl/trainers/centralized/mode.py | 6 + docs/content/_index.md | 2 + .../docs/marl/multi-agent-actor-critic.md | 4 + docs/content/docs/marl/multi-llm-pbrl.md | 2 +- tests/test_centralized_actor_critic.py | 174 ++++++++++++++++++ 9 files changed, 331 insertions(+), 67 deletions(-) create mode 100644 tests/test_centralized_actor_critic.py diff --git a/README.md b/README.md index 205fd7b..6215e48 100644 --- a/README.md +++ b/README.md @@ -55,10 +55,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/pbrl/madpo.py), [MARLHF](comlrl/trainers/pbrl/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/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index f0193eb..69ef0cd 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -13,10 +13,24 @@ from comlrl.utils.checkpoint import TrainingCheckpointMixin 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: @@ -105,7 +119,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 +135,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 @@ -425,9 +439,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] = [] diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index 552ea1d..c079800 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -27,6 +27,7 @@ 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 @@ -84,6 +87,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,7 +119,9 @@ 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 + ) @dataclass @@ -162,8 +170,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 +194,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 +208,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 +216,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,8 +226,8 @@ 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) @@ -221,9 +239,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 +252,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 +298,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 +340,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,7 +386,7 @@ 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 @@ -391,7 +409,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 +424,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)), @@ -559,7 +582,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 +601,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 +668,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 +681,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 +703,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 +723,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 +782,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 +794,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()) @@ -1043,7 +1066,7 @@ 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 +1108,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..38b0bbf 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -27,6 +27,7 @@ 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 +54,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 @@ -81,6 +84,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 +111,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 +144,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 +180,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 +188,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,7 +199,7 @@ 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) @@ -192,7 +210,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 +219,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 +303,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 +329,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,7 +348,7 @@ 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 @@ -346,7 +364,7 @@ def _init_wandb(self) -> 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 +376,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 +406,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, @@ -505,7 +544,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 +572,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 +591,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 +678,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 +691,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 +699,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 +715,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 +746,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 +825,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 +837,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()) diff --git a/comlrl/trainers/centralized/mode.py b/comlrl/trainers/centralized/mode.py index 4f1bc91..ec447d3 100644 --- a/comlrl/trainers/centralized/mode.py +++ b/comlrl/trainers/centralized/mode.py @@ -73,6 +73,12 @@ def _setup_collaboration( 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: 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..280b6e0 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 @@ -77,6 +78,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 +126,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 @@ -169,6 +172,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-llm-pbrl.md b/docs/content/docs/marl/multi-llm-pbrl.md index 50f4d0b..fab6a16 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. diff --git a/tests/test_centralized_actor_critic.py b/tests/test_centralized_actor_critic.py new file mode 100644 index 0000000..5608e12 --- /dev/null +++ b/tests/test_centralized_actor_critic.py @@ -0,0 +1,174 @@ +"""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("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 From 3b0102d0e2b0fdf96d229665f7877b1e4bb3731b Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:09:17 -0400 Subject: [PATCH 04/21] refactor: make checkpointing independent of pbrl - Trainers register their checkpoint dataclasses with register_checkpoint_record instead of a hard-coded allowlist of pbrl module paths; decoding no longer imports modules by name. - The untracked-state flag and its error are generic (_checkpoint_untracked). - pbrl prunes its asset pool through a _checkpoint_prune_assets hook. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/iac.py | 2 ++ comlrl/trainers/pbrl/iterative.py | 6 +++-- comlrl/trainers/pbrl/madpo.py | 10 ++++++- comlrl/trainers/pbrl/marlhf.py | 2 +- comlrl/trainers/reinforce/magrpo.py | 3 ++- comlrl/utils/checkpoint.py | 42 ++++++++++++++--------------- tests/test_checkpoint_files.py | 2 +- tests/test_checkpoint_metadata.py | 6 ++--- tests/test_checkpoint_preference.py | 2 +- 9 files changed, 42 insertions(+), 33 deletions(-) diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index c079800..982203b 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -11,6 +11,7 @@ from comlrl.models.actor_critic import CausalLMWithValueHead from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training +from comlrl.utils.checkpoint import register_checkpoint_record from comlrl.utils.distributed import local_context, unwrap_model from comlrl.utils.formatters import build_formatters from comlrl.utils.model_loading import resolve_model_kwargs, resolve_model_sources @@ -124,6 +125,7 @@ def __post_init__(self) -> None: ) +@register_checkpoint_record @dataclass class RolloutSample: agent_idx: int diff --git a/comlrl/trainers/pbrl/iterative.py b/comlrl/trainers/pbrl/iterative.py index a0ee5cf..b3e150a 100644 --- a/comlrl/trainers/pbrl/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -18,6 +18,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, @@ -328,6 +329,7 @@ def __post_init__(self) -> None: ) +@register_checkpoint_record @dataclass class PreferenceReplayShard: iteration: int @@ -633,7 +635,7 @@ def train( ): return train_iterations_checkpointed(self, reward_model=False, **kwargs) - self._checkpoint_untracked_preference = True + self._checkpoint_untracked = True if int(self.args.num_turns) != 1: raise ValueError("MADPOIter currently supports num_turns=1 only.") @@ -2916,7 +2918,7 @@ def train( ): return train_iterations_checkpointed(self, reward_model=True, **kwargs) - self._checkpoint_untracked_preference = True + self._checkpoint_untracked = True if int(self.args.num_turns) != 1: raise ValueError("MARLHFIter currently supports num_turns=1 only.") diff --git a/comlrl/trainers/pbrl/madpo.py b/comlrl/trainers/pbrl/madpo.py index e636589..4a5174f 100644 --- a/comlrl/trainers/pbrl/madpo.py +++ b/comlrl/trainers/pbrl/madpo.py @@ -10,6 +10,7 @@ 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 @@ -23,6 +24,7 @@ ) +@register_checkpoint_record @dataclass class AgentPreferenceTensors: prompt_input_ids: torch.Tensor @@ -30,6 +32,7 @@ class AgentPreferenceTensors: loser_completion_ids: torch.Tensor +@register_checkpoint_record @dataclass class PreferencePair: prompts: List[str] @@ -239,7 +242,7 @@ def train( ): 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.") @@ -362,6 +365,11 @@ def _checkpoint_load_assets(self, checkpoint_dir, assets): load_preference_assets(self, checkpoint_dir, assets) + def _checkpoint_prune_assets(self, checkpoint_root): + from .checkpoint_extras import prune_preference_asset_pool + + prune_preference_asset_pool(checkpoint_root) + def _checkpoint_validate_assets(self, checkpoint_dir, assets): from .checkpoint_extras import validate_preference_assets diff --git a/comlrl/trainers/pbrl/marlhf.py b/comlrl/trainers/pbrl/marlhf.py index cbadd1d..4b44415 100644 --- a/comlrl/trainers/pbrl/marlhf.py +++ b/comlrl/trainers/pbrl/marlhf.py @@ -225,7 +225,7 @@ def train( ): 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.") diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index c0d6ba2..d43a1c7 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -17,7 +17,7 @@ from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training -from comlrl.utils.checkpoint import TrainingCheckpointMixin +from comlrl.utils.checkpoint import TrainingCheckpointMixin, register_checkpoint_record from comlrl.utils.distributed import local_context, unwrap_model from comlrl.utils.formatters import build_formatters from comlrl.utils.model_loading import ( @@ -132,6 +132,7 @@ def __post_init__(self) -> None: validate_reference_kl_config(self, actor_count) +@register_checkpoint_record @dataclass class NodeSample: agent_idx: int diff --git a/comlrl/utils/checkpoint.py b/comlrl/utils/checkpoint.py index aa9808e..708e22d 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -7,7 +7,6 @@ from contextlib import contextmanager from dataclasses import fields, is_dataclass -import importlib 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.pbrl.madpo.PreferencePair", - "comlrl.trainers.pbrl.madpo.AgentPreferenceTensors", - "comlrl.trainers.pbrl.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.") @@ -949,7 +950,7 @@ def restore_trainer_state(trainer, state, checkpoint_dir=None, restore_rng=True) 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): @@ -1113,9 +1114,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 +1229,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 +1294,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.pbrl.checkpoint_extras 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/tests/test_checkpoint_files.py b/tests/test_checkpoint_files.py index 5a7f8a3..fd4a328 100644 --- a/tests/test_checkpoint_files.py +++ b/tests/test_checkpoint_files.py @@ -196,7 +196,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" 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_preference.py b/tests/test_checkpoint_preference.py index f7db919..c73ed59 100644 --- a/tests/test_checkpoint_preference.py +++ b/tests/test_checkpoint_preference.py @@ -377,7 +377,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") From f2662157eaa4c4e1b250839515407acd87db1184 Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:33:09 -0400 Subject: [PATCH 05/21] refactor: bind pbrl checkpoint hooks to their helpers MADPOTrainer's checkpoint hooks were seven methods that each imported one checkpoint_extras function at call time; nothing needs the import to be lazy, so the hooks are now the functions themselves. The four _checkpoint_validate_progress overrides become one method that reads a per-class _checkpoint_progress_kind. Function-level imports in marlhf, iterative and checkpoint_extras move to module level, except the two that break the madpo/iterative import cycle and MARLHF's lookup of the built-in actor-critic classes. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/pbrl/checkpoint_extras.py | 10 ++--- comlrl/trainers/pbrl/iterative.py | 17 ++------ comlrl/trainers/pbrl/madpo.py | 53 ++++++++--------------- comlrl/trainers/pbrl/marlhf.py | 40 ++++++++--------- 4 files changed, 44 insertions(+), 76 deletions(-) diff --git a/comlrl/trainers/pbrl/checkpoint_extras.py b/comlrl/trainers/pbrl/checkpoint_extras.py index dc6e6bc..2eb450a 100644 --- a/comlrl/trainers/pbrl/checkpoint_extras.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"])) ) diff --git a/comlrl/trainers/pbrl/iterative.py b/comlrl/trainers/pbrl/iterative.py index b3e150a..9713569 100644 --- a/comlrl/trainers/pbrl/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -35,6 +35,7 @@ PreferencePair, ) from ..centralized import CentralizedComparatorAdapter +from .checkpoint_extras import train_iterations_checkpointed from .marlhf import MARLHFConfig, MARLHFTrainer from .output_paths import resolve_output_path @@ -348,11 +349,7 @@ class MADPOIterTrainer(MADPOTrainer): default_config_cls = MADPOIterConfig algorithm_name = "MADPOIter" - - def _checkpoint_validate_progress(self, state): - from .checkpoint_extras import validate_preference_progress - - validate_preference_progress(self, state, "madpo_iter") + _checkpoint_progress_kind = "madpo_iter" def __init__( self, @@ -625,8 +622,6 @@ def train( or checkpoint_interval != 1 or checkpoint_total_limit is not None ): - from .checkpoint_extras import train_iterations_checkpointed - with self._checkpoint_context( resume_from_checkpoint=resume_from_checkpoint, checkpoint_dir=checkpoint_dir, @@ -2860,11 +2855,7 @@ class MARLHFIterTrainer(MADPOIterTrainer, MARLHFTrainer): default_config_cls = MARLHFIterConfig algorithm_name = "MARLHFIter" - - def _checkpoint_validate_progress(self, state): - from .checkpoint_extras import validate_preference_progress - - validate_preference_progress(self, state, "marlhf_iter") + _checkpoint_progress_kind = "marlhf_iter" def _preference_scoring_raw_and_processed_rewards( self, @@ -2908,8 +2899,6 @@ def train( or checkpoint_interval != 1 or checkpoint_total_limit is not None ): - from .checkpoint_extras import train_iterations_checkpointed - with self._checkpoint_context( resume_from_checkpoint=resume_from_checkpoint, checkpoint_dir=checkpoint_dir, diff --git a/comlrl/trainers/pbrl/madpo.py b/comlrl/trainers/pbrl/madpo.py index 4a5174f..e459e8a 100644 --- a/comlrl/trainers/pbrl/madpo.py +++ b/comlrl/trainers/pbrl/madpo.py @@ -22,6 +22,16 @@ 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, +) @register_checkpoint_record @@ -345,45 +355,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 .checkpoint_extras import preference_extra_state - - return preference_extra_state(self) - - def _restore_checkpoint_extra_state(self, state, checkpoint_dir): - from .checkpoint_extras import restore_preference_extra_state - - restore_preference_extra_state(self, state, checkpoint_dir) - - def _checkpoint_save_assets(self, checkpoint_dir): - from .checkpoint_extras import save_preference_assets - - return save_preference_assets(self, checkpoint_dir) - - def _checkpoint_load_assets(self, checkpoint_dir, assets): - from .checkpoint_extras import load_preference_assets - - load_preference_assets(self, checkpoint_dir, assets) - - def _checkpoint_prune_assets(self, checkpoint_root): - from .checkpoint_extras import prune_preference_asset_pool - - prune_preference_asset_pool(checkpoint_root) - - def _checkpoint_validate_assets(self, checkpoint_dir, assets): - from .checkpoint_extras 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 + _restore_checkpoint_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 .checkpoint_extras import validate_preference_extra_state - validate_preference_extra_state(state) def _checkpoint_validate_progress(self, state): - from .checkpoint_extras 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/pbrl/marlhf.py b/comlrl/trainers/pbrl/marlhf.py index 4b44415..49dd60e 100644 --- a/comlrl/trainers/pbrl/marlhf.py +++ b/comlrl/trainers/pbrl/marlhf.py @@ -11,8 +11,23 @@ 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 ..centralized import joint_policy_gradient_loss @@ -166,6 +181,7 @@ class MARLHFTrainer(MADPOTrainer): default_config_cls = MARLHFConfig algorithm_name = "MARLHF" + _checkpoint_progress_kind = "marlhf" def __init__( self, @@ -240,11 +256,6 @@ def train( self._reward_model_active = True self._train_online_rl(**kwargs) - def _checkpoint_validate_progress(self, state): - from .checkpoint_extras 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 +264,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 +376,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 +413,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 +595,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 +602,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, From ddf4d112f1003f85af3c04835f3e31ef9accf625 Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:34:27 -0400 Subject: [PATCH 06/21] refactor: check train() checkpoint arguments in one place MAGRPO, the actor-critic base, MADPO, MARLHF, MADPOIter and MARLHFIter each validated checkpoint_interval and decided whether checkpointing was requested with the same thirteen lines. TrainingCheckpointMixin now does this in _checkpoint_arguments, which returns the arguments for _checkpoint_context or None. Error messages and checks are unchanged. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 27 ++++--------- comlrl/trainers/pbrl/iterative.py | 53 ++++++++----------------- comlrl/trainers/pbrl/madpo.py | 27 ++++--------- comlrl/trainers/pbrl/marlhf.py | 27 ++++--------- comlrl/trainers/reinforce/magrpo.py | 27 ++++--------- comlrl/utils/checkpoint.py | 28 +++++++++++++ 6 files changed, 76 insertions(+), 113 deletions(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index 69ef0cd..e7d0cfd 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 @@ -536,24 +535,14 @@ 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 ) diff --git a/comlrl/trainers/pbrl/iterative.py b/comlrl/trainers/pbrl/iterative.py index 9713569..f67d419 100644 --- a/comlrl/trainers/pbrl/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 @@ -610,24 +609,14 @@ 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 train_iterations_checkpointed(self, reward_model=False, **kwargs) self._checkpoint_untracked = True @@ -2887,24 +2876,14 @@ 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 train_iterations_checkpointed(self, reward_model=True, **kwargs) self._checkpoint_untracked = True diff --git a/comlrl/trainers/pbrl/madpo.py b/comlrl/trainers/pbrl/madpo.py index e459e8a..cb524cc 100644 --- a/comlrl/trainers/pbrl/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 @@ -232,24 +231,14 @@ 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 = True diff --git a/comlrl/trainers/pbrl/marlhf.py b/comlrl/trainers/pbrl/marlhf.py index 49dd60e..de60020 100644 --- a/comlrl/trainers/pbrl/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 @@ -221,24 +220,14 @@ 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 = True diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index d43a1c7..913c294 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 @@ -752,24 +751,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 ) diff --git a/comlrl/utils/checkpoint.py b/comlrl/utils/checkpoint.py index 708e22d..ab809b9 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -1077,6 +1077,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( From 7305437d6b7e55631e16f170c32676e32cf6be37 Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:39:18 -0400 Subject: [PATCH 07/21] refactor: run default and checkpointed training through one loop MAGRPOTrainer and ActorCriticTrainerBase each had two copies of their epoch loop, one for plain train() and one for checkpointed training; both now call _run_epochs, which commits after an epoch only when given a callback. MADPOIter and MARLHFIter had three copies of the iteration loop (two plain train() methods and the checkpointed one in checkpoint_extras); MADPOIterTrainer._train_iterations now runs all of them, and MARLHFIter inherits train(). Call order, logs and console output are unchanged. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 37 ++--- comlrl/trainers/pbrl/checkpoint_extras.py | 76 ++-------- comlrl/trainers/pbrl/iterative.py | 160 ++++++++-------------- comlrl/trainers/reinforce/magrpo.py | 54 ++------ 4 files changed, 95 insertions(+), 232 deletions(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index e7d0cfd..731b2d9 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -549,29 +549,12 @@ def train( 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 @@ -612,7 +595,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, { @@ -620,9 +602,13 @@ 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.""" + total_epochs = self.args.num_train_epochs for epoch in range(start_epoch, total_epochs): epoch_metrics = defaultdict(list) dataloader = self.get_train_dataloader() @@ -638,7 +624,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/pbrl/checkpoint_extras.py b/comlrl/trainers/pbrl/checkpoint_extras.py index 2eb450a..b863d0e 100644 --- a/comlrl/trainers/pbrl/checkpoint_extras.py +++ b/comlrl/trainers/pbrl/checkpoint_extras.py @@ -641,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( { @@ -662,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/pbrl/iterative.py b/comlrl/trainers/pbrl/iterative.py index f67d419..42ecd4e 100644 --- a/comlrl/trainers/pbrl/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -615,34 +615,63 @@ def train( 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=False, **kwargs) + return train_iterations_checkpointed( + self, reward_model=reward_model, **kwargs + ) 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, @@ -659,25 +688,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, @@ -2867,81 +2905,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, - ): - 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 train_iterations_checkpointed(self, reward_model=True, **kwargs) - - self._checkpoint_untracked = 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/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index 913c294..ed8212e 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -768,39 +768,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 = { @@ -860,7 +828,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, @@ -869,10 +836,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() @@ -895,8 +865,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 @@ -904,10 +875,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, From 9fa15626d34b6f76b0d7a539492ea5c54d4b202f Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:40:27 -0400 Subject: [PATCH 08/21] refactor: drop the single-process DistributedContext local_context() always returned is_main=True, so every dist_env.is_main check in MAGRPO, IAC, MAAC and the actor-critic base was a no-op. DistributedContext, local_context and the trainers' dist_env attribute are removed; comlrl.utils.distributed keeps unwrap_model. The test of local_context goes with it (CoMLRL tests 735 -> 734). Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 11 ++--------- comlrl/trainers/actor_critic/iac.py | 9 ++------- comlrl/trainers/actor_critic/maac.py | 9 ++------- comlrl/trainers/reinforce/magrpo.py | 7 ++----- comlrl/utils/__init__.py | 8 +------- comlrl/utils/distributed.py | 19 +------------------ tests/checkpoint_helpers.py | 1 - tests/test_actor_critic.py | 7 ------- tests/test_utils.py | 15 ++------------- 9 files changed, 12 insertions(+), 74 deletions(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index 731b2d9..86b3219 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -284,9 +284,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, @@ -312,12 +310,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( diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index 982203b..fe795ca 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -12,7 +12,7 @@ from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training from comlrl.utils.checkpoint import register_checkpoint_record -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 ( @@ -231,7 +231,6 @@ def __init__( 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] = [] @@ -392,7 +391,7 @@ def __init__( 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) @@ -402,8 +401,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: @@ -1065,8 +1062,6 @@ 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.num_agents == 1: actor = unwrap_model(self.agents[0]) diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index 38b0bbf..547e0ac 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -11,7 +11,7 @@ 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 ( @@ -202,7 +202,6 @@ def __init__( 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) @@ -353,13 +352,11 @@ def __init__( 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] = { @@ -968,8 +965,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/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index ed8212e..4f9b5ac 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -17,7 +17,7 @@ from comlrl.schedulers import DeviceScheduler from comlrl.schedulers.device_scheduler import normalize_parallel_training from comlrl.utils.checkpoint import TrainingCheckpointMixin, register_checkpoint_record -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 ( infer_model_name, @@ -284,7 +284,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 @@ -376,7 +375,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) @@ -462,8 +461,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 = {} diff --git a/comlrl/utils/__init__.py b/comlrl/utils/__init__.py index 5e6d721..04f5b08 100644 --- a/comlrl/utils/__init__.py +++ b/comlrl/utils/__init__.py @@ -16,11 +16,7 @@ resolve_reward_range, set_reward_range, ) -from .distributed import ( - DistributedContext, - local_context, - unwrap_model, -) +from .distributed import unwrap_model from .tokenizer_utils import ( apply_tokenizer_specials, ensure_pad_token, @@ -50,7 +46,5 @@ "normalize_reward_lengths", "resolve_reward_range", "set_reward_range", - "DistributedContext", - "local_context", "unwrap_model", ] diff --git a/comlrl/utils/distributed.py b/comlrl/utils/distributed.py index dc513d0..695d007 100644 --- a/comlrl/utils/distributed.py +++ b/comlrl/utils/distributed.py @@ -1,21 +1,4 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Optional - -import torch - - -@dataclass(frozen=True) -class DistributedContext: - is_main: bool - device: torch.device - - -def local_context(device: Optional[torch.device] = None) -> DistributedContext: - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - return DistributedContext(is_main=True, device=device) +from typing import Any def unwrap_model(model: Any) -> Any: diff --git a/tests/checkpoint_helpers.py b/tests/checkpoint_helpers.py index 58dff49..48032f0 100644 --- a/tests/checkpoint_helpers.py +++ b/tests/checkpoint_helpers.py @@ -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/test_actor_critic.py b/tests/test_actor_critic.py index 10c2bfe..07d8a5e 100644 --- a/tests/test_actor_critic.py +++ b/tests/test_actor_critic.py @@ -12,7 +12,6 @@ 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 @@ -47,10 +46,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", @@ -334,7 +329,6 @@ def test_ac_step_matches_analytic_gradients(trainer_cls, separate, critic_type): 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 +350,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 = [] diff --git a/tests/test_utils.py b/tests/test_utils.py index 7be57b2..0872a56 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,16 +1,11 @@ -"""Reward, formatter, W&B and distributed helpers in comlrl.utils.""" +"""Reward, formatter and W&B helpers in comlrl.utils.""" import inspect import pytest import torch -from comlrl.utils import ( - RewardProcessors, - build_formatters, - call_reward_function, - local_context, -) +from comlrl.utils import RewardProcessors, build_formatters, call_reward_function from comlrl.utils.wandb_utils import wandb_init_kwargs @@ -156,9 +151,3 @@ def test_wandb_init_kwargs_keeps_every_config_section(): "dataset_name": "data", "dataset_type": "task", } - - -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") From be4152d4c71a6b619762f00474a9cbebe647f736 Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:42:58 -0400 Subject: [PATCH 09/21] refactor: share one reference-KL helper across trainers MAGRPO, IAC and MAAC each defined the same _reference_kl_values method, which forwarded to _sequence_kl_values with the strategy fixed to reference_kl_for_sequence. comlrl.utils.reference_kl.reference_kl_values now does both steps and the three trainers bind it. The KL test patches the helper's own module instead of each trainer module. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/iac.py | 29 ++-------------------------- comlrl/trainers/actor_critic/maac.py | 29 ++-------------------------- comlrl/trainers/reinforce/magrpo.py | 29 ++-------------------------- comlrl/utils/reference_kl.py | 21 ++++++++++++++------ tests/test_reference_kl.py | 9 +++------ 5 files changed, 24 insertions(+), 93 deletions(-) diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index fe795ca..284dcfd 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -16,12 +16,11 @@ 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, ) @@ -543,31 +542,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))) diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index 547e0ac..79d426c 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -15,12 +15,11 @@ 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, ) @@ -503,31 +502,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))) diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index 4f9b5ac..ee361d9 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -25,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, ) @@ -1431,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] 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/tests/test_reference_kl.py b/tests/test_reference_kl.py index be6607e..672cb5d 100644 --- a/tests/test_reference_kl.py +++ b/tests/test_reference_kl.py @@ -1,6 +1,5 @@ """Reference-model KL penalties for MAGRPO, IAC and MAAC.""" -import importlib from types import SimpleNamespace import pytest @@ -12,6 +11,7 @@ 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, @@ -199,9 +199,7 @@ def test_reference_kl_values_disabled_or_missing_models(trainer_cls): @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 -): +def test_reference_kl_values_keep_order(trainer_cls, response_lens, monkeypatch): trainer = trainer_cls.__new__(trainer_cls) trainer.args = SimpleNamespace(reference_kl_enabled=True) trainer.reference_models = [object(), object()] @@ -221,8 +219,7 @@ 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) + monkeypatch.setattr(reference_kl, "reference_kl_for_sequence", sequence_kl) values = trainer._reference_kl_values(1, policy, ids, mask, 2, response_lens) count = min(len(ids), len(response_lens)) assert values == [idx + 0.25 for idx in range(count)] From a27dc35fb11a59112d3f10c7468ae8f4a881940d Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:44:47 -0400 Subject: [PATCH 10/21] chore: remove a test-only replay helper and a typo alias MADPOIterTrainer._load_replay_shard was called only by one test; replay sampling loads records and rebuilds pairs directly, and the test now does the same. MARLHF also accepted rl_algorithm="mareforce" as a misspelling of "mareinforce"; no config, test or document uses it, so it is now rejected like any other unknown name. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/pbrl/iterative.py | 11 ----------- comlrl/trainers/pbrl/marlhf.py | 6 +----- tests/test_centralized_preference.py | 5 ++++- 3 files changed, 5 insertions(+), 17 deletions(-) diff --git a/comlrl/trainers/pbrl/iterative.py b/comlrl/trainers/pbrl/iterative.py index 42ecd4e..4227289 100644 --- a/comlrl/trainers/pbrl/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -1357,17 +1357,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: diff --git a/comlrl/trainers/pbrl/marlhf.py b/comlrl/trainers/pbrl/marlhf.py index de60020..cb7fb29 100644 --- a/comlrl/trainers/pbrl/marlhf.py +++ b/comlrl/trainers/pbrl/marlhf.py @@ -35,9 +35,6 @@ from ..reinforce.magrpo import MAGRPOTrainer -_RL_ALGORITHM_ALIASES = { - "mareforce": "mareinforce", -} _MAGRPO_ADVANTAGE_MODES = { "magrpo": "mean", "mareinforce": "raw", @@ -48,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: diff --git a/tests/test_centralized_preference.py b/tests/test_centralized_preference.py index 0a5c827..aae08c9 100644 --- a/tests/test_centralized_preference.py +++ b/tests/test_centralized_preference.py @@ -164,7 +164,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, From 1b28c16594dc9e2622e24d0e4043bd37f7d1097b Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 20:51:08 -0400 Subject: [PATCH 11/21] refactor: name the extra-state restore hook like the other checkpoint hooks Every optional trainer hook that comlrl.utils.checkpoint looks up is named _checkpoint_* except _restore_checkpoint_extra_state, which becomes _checkpoint_restore_extra_state. Only MADPOTrainer and one test define or call it; no LLM_Collab repository refers to it, and checkpoint files store no hook names. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/pbrl/madpo.py | 2 +- comlrl/utils/checkpoint.py | 2 +- tests/test_checkpoint_preference.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/comlrl/trainers/pbrl/madpo.py b/comlrl/trainers/pbrl/madpo.py index cb524cc..054ac7c 100644 --- a/comlrl/trainers/pbrl/madpo.py +++ b/comlrl/trainers/pbrl/madpo.py @@ -347,7 +347,7 @@ def commit_epoch(next_epoch, updates_seen, *, force=False): # Checkpoint hooks looked up by comlrl.utils.checkpoint (see checkpoint_extras). _checkpoint_progress_kind = "madpo" _checkpoint_extra_state = preference_extra_state - _restore_checkpoint_extra_state = restore_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) diff --git a/comlrl/utils/checkpoint.py b/comlrl/utils/checkpoint.py index ab809b9..cac25b2 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -943,7 +943,7 @@ 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")) diff --git a/tests/test_checkpoint_preference.py b/tests/test_checkpoint_preference.py index c73ed59..780163c 100644 --- a/tests/test_checkpoint_preference.py +++ b/tests/test_checkpoint_preference.py @@ -386,7 +386,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 From cdf3977ca18d8587b3ce4e7a5d3d492647d9e7ea Mon Sep 17 00:00:00 2001 From: N!no Date: Mon, 19 Oct 2026 21:09:57 -0400 Subject: [PATCH 12/21] test: drop redundant tests and fix one config case Five read-only reviews, per-test line coverage and mutation runs found tests that repeat a stronger test of the same behavior: - deleted: the MADPO first-update failure, nested RNG route conflict, weights-only load, centralized joint-loss and reference-device, default comparator mode (folded into the config defaults test), two multi-turn IAC/MAAC trainer tests that only reach config validation, a foreign-run save that only reaches the new-run preflight, a reused-blob test, and a resume-fallback test whose cases other tests already pin; - merged by parametrization: two preference cursor tests, resume conflict kinds (now with each error message), checkpoint write/sync/publish failures (staging cleanup now checked for all three) and the two initial asset copy failures; - base-class methods (_expand_rewards, _compute_sequence_stats, _kl_shaped_reward) and the shared reference_kl_values helper are tested once instead of once per trainer, plus a check that all trainers bind the helper; valid-config cases equal to the defaults are dropped. The IAC/MAAC num_generations=0 case ran with num_turns=2 and so hit the multi-turn rule; it is now single-turn and reaches the lower-bound check. Covered comlrl lines: 5875 before, 5876 after (none lost). Tests: 734 -> 697. Co-Authored-By: Claude Opus 5.5 --- tests/test_actor_critic.py | 15 ++--- tests/test_centralized_comparator.py | 5 -- tests/test_centralized_magrpo.py | 28 --------- tests/test_centralized_preference.py | 17 ++--- tests/test_checkpoint_assets.py | 81 +++--------------------- tests/test_checkpoint_files.py | 84 ++++++++----------------- tests/test_checkpoint_preference.py | 46 +++----------- tests/test_checkpoint_rng.py | 13 ---- tests/test_checkpoint_state.py | 16 ----- tests/test_config.py | 12 ++-- tests/test_reference_kl.py | 92 ++++++++++------------------ tests/test_trainer_init.py | 23 ------- 12 files changed, 92 insertions(+), 340 deletions(-) diff --git a/tests/test_actor_critic.py b/tests/test_actor_critic.py index 07d8a5e..8d83bfc 100644 --- a/tests/test_actor_critic.py +++ b/tests/test_actor_critic.py @@ -46,7 +46,6 @@ def forward(self, input_ids, attention_mask, output_values): return SimpleNamespace(logits=logits, values=values) -@pytest.mark.parametrize("trainer_cls", TRAINERS) @pytest.mark.parametrize( "num_agents,num_ret,rewards,expected", [ @@ -57,9 +56,9 @@ def forward(self, input_ids, attention_mask, output_values): ], ) 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 @@ -67,9 +66,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) @@ -122,9 +120,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, diff --git a/tests/test_centralized_comparator.py b/tests/test_centralized_comparator.py index cfd398f..8b6c204 100644 --- a/tests/test_centralized_comparator.py +++ b/tests/test_centralized_comparator.py @@ -55,11 +55,6 @@ 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" - - def test_centralized_split_preserves_candidate_and_agent_alignment(): trainer = MADPOIterTrainer.__new__(MADPOIterTrainer) trainer.num_agents = 2 diff --git a/tests/test_centralized_magrpo.py b/tests/test_centralized_magrpo.py index e297d92..e9853de 100644 --- a/tests/test_centralized_magrpo.py +++ b/tests/test_centralized_magrpo.py @@ -27,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): @@ -45,32 +43,6 @@ def test_eval_passes_separate_roles_to_eval_logger(tokenizer): assert result["eval/turn_1/domain_metric"] == 3.0 -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): trainer, _ = make_centralized_magrpo_trainer( tokenizer, reference_kl_enabled=True, reference_devices=["cpu"] diff --git a/tests/test_centralized_preference.py b/tests/test_centralized_preference.py index aae08c9..b851ba6 100644 --- a/tests/test_centralized_preference.py +++ b/tests/test_centralized_preference.py @@ -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)] ) diff --git a/tests/test_checkpoint_assets.py b/tests/test_checkpoint_assets.py index 4b57ec0..281fed9 100644 --- a/tests/test_checkpoint_assets.py +++ b/tests/test_checkpoint_assets.py @@ -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 @@ -468,27 +439,6 @@ def test_digest_cache_invalidates_modified_or_replaced_source( 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" - checkpoint_extras._copy_checkpoint_asset(source, first, pool, digest_cache=cache) - first_reads = sum(counts.values()) - second = tmp_path / "second/weights.bin" - checkpoint_extras._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" - checkpoint_extras._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") @@ -503,14 +453,19 @@ def test_source_change_during_hash_is_rejected_without_caching(tmp_path, monkeyp 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 = 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 @@ -524,28 +479,6 @@ def change_after_copy(src, dst, *args, **kwargs): 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 = checkpoint_extras.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(checkpoint_extras.shutil, "copy2", corrupt_copy) - destination, pool = tmp_path / "snapshot/weights.bin", tmp_path / "pool" - with pytest.raises(RuntimeError, match="changed during checkpoint saving"): - checkpoint_extras._copy_checkpoint_asset( - source, destination, pool, digest_cache={} - ) - assert not destination.exists() - assert not list(pool.iterdir()) - - def test_cached_source_change_during_link_is_rejected(tmp_path, monkeypatch): source = tmp_path / "weights.bin" source.write_bytes(b"original") diff --git a/tests/test_checkpoint_files.py b/tests/test_checkpoint_files.py index fd4a328..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) @@ -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_preference.py b/tests/test_checkpoint_preference.py index 780163c..c153cd7 100644 --- a/tests/test_checkpoint_preference.py +++ b/tests/test_checkpoint_preference.py @@ -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: [] @@ -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_rng.py b/tests/test_checkpoint_rng.py index 235fbc5..df94c6f 100644 --- a/tests/test_checkpoint_rng.py +++ b/tests/test_checkpoint_rng.py @@ -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..204868d 100644 --- a/tests/test_checkpoint_state.py +++ b/tests/test_checkpoint_state.py @@ -10,8 +10,6 @@ from comlrl.trainers.reinforce.magrpo import NodeSample import comlrl.utils.checkpoint as checkpoint from comlrl.utils.checkpoint import ( - _decode, - _encode, capture_trainer_state, restore_trainer_state, ) @@ -155,20 +153,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 diff --git a/tests/test_config.py b/tests/test_config.py index 90e9218..9f270b4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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_reference_kl.py b/tests/test_reference_kl.py index 672cb5d..70c4605 100644 --- a/tests/test_reference_kl.py +++ b/tests/test_reference_kl.py @@ -6,7 +6,7 @@ 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 @@ -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,22 +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(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]]) @@ -220,26 +195,27 @@ def sequence_kl( return torch.tensor([idx + 0.25, 100.0]) monkeypatch.setattr(reference_kl, "reference_kl_for_sequence", sequence_kl) - values = trainer._reference_kl_values(1, policy, ids, mask, 2, response_lens) + 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 aed7839..e5cdb5e 100644 --- a/tests/test_trainer_init.py +++ b/tests/test_trainer_init.py @@ -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], From 3eb2449e4c7baa1949737f1ebd3e8139d21d43b7 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 01:32:11 -0400 Subject: [PATCH 13/21] fix: resume checkpoints written before a config field existed Resume compared the saved and current configs as whole dicts, so adding any config field made every older checkpoint unresumable (with an empty list of changed fields). A field missing from the checkpoint now matches while it keeps its dataclass default, and the error names fields present on only one side. Co-Authored-By: Claude Opus 5.5 --- comlrl/utils/checkpoint.py | 36 ++++++++++++++++++++++++++++++---- tests/test_checkpoint_state.py | 16 +++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/comlrl/utils/checkpoint.py b/comlrl/utils/checkpoint.py index cac25b2..a2577e3 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -6,7 +6,7 @@ """ from contextlib import contextmanager -from dataclasses import fields, is_dataclass +from dataclasses import MISSING, fields, is_dataclass import inspect import hashlib import json @@ -294,6 +294,26 @@ def _config(trainer): } +def _config_defaults(trainer): + """Encoded config defaults, for fields added after a checkpoint was written.""" + args = getattr(trainer, "args", None) + if not is_dataclass(args): + return {} + defaults = {} + for field in fields(args): + if field.default is not MISSING: + default = field.default + elif field.default_factory is not MISSING: + default = field.default_factory() + 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: @@ -751,11 +771,19 @@ 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) + defaults = _config_defaults(trainer) + # A field added after the checkpoint was written matches at its default. + for key in set(current_config) - set(saved_config): + 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)}" diff --git a/tests/test_checkpoint_state.py b/tests/test_checkpoint_state.py index 204868d..2a4cdd4 100644 --- a/tests/test_checkpoint_state.py +++ b/tests/test_checkpoint_state.py @@ -441,6 +441,22 @@ 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) + + @pytest.mark.parametrize("changed", ["vocab", "chat_template", "eos"]) def test_tokenizer_metadata_mismatch_precedes_any_mutation( tmp_path, tokenizer, changed From f94e2d726694cec88b389c123966265914dd0b93 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 01:32:12 -0400 Subject: [PATCH 14/21] feat: give the centralized comparator its own answer budget comparator_centralized_max_new_tokens sets the token budget of a local centralized comparator's joint response, which holds every agent's output. It defaults to max_new_tokens, so existing runs are unchanged. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/pbrl/iterative.py | 15 +++++++++++- docs/content/docs/marl/multi-llm-pbrl.md | 1 + tests/test_centralized_comparator.py | 29 ++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/comlrl/trainers/pbrl/iterative.py b/comlrl/trainers/pbrl/iterative.py index 4227289..861d40b 100644 --- a/comlrl/trainers/pbrl/iterative.py +++ b/comlrl/trainers/pbrl/iterative.py @@ -121,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.") @@ -275,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 @@ -1913,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, ) diff --git a/docs/content/docs/marl/multi-llm-pbrl.md b/docs/content/docs/marl/multi-llm-pbrl.md index fab6a16..2e2e01e 100644 --- a/docs/content/docs/marl/multi-llm-pbrl.md +++ b/docs/content/docs/marl/multi-llm-pbrl.md @@ -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 %}} diff --git a/tests/test_centralized_comparator.py b/tests/test_centralized_comparator.py index 8b6c204..c5855bd 100644 --- a/tests/test_centralized_comparator.py +++ b/tests/test_centralized_comparator.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -55,6 +56,34 @@ def test_iterative_config_allows_generic_centralized_agent_counts(): assert config.num_agents == 3 +@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(): trainer = MADPOIterTrainer.__new__(MADPOIterTrainer) trainer.num_agents = 2 From c0b23fbf540a44c6712476e3c644d740dd2388b1 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 01:36:58 -0400 Subject: [PATCH 15/21] feat: let actor-critic subclasses add generate arguments IAC and MAAC build a closed set of generate keyword arguments. ActorCriticTrainerBase._generation_kwargs(agent_idx) now adds to them; it returns nothing by default, so generation is unchanged. A domain can use it for stop strings, as MAGRPO subclasses already can through their generation kwargs. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 4 ++++ comlrl/trainers/actor_critic/iac.py | 1 + comlrl/trainers/actor_critic/maac.py | 1 + tests/test_centralized_actor_critic.py | 16 ++++++++++++++++ 4 files changed, 22 insertions(+) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index 86b3219..ac6e710 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -191,6 +191,10 @@ 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 _get_tokenizer(self, agent_idx: Optional[int] = None): tokenizers = getattr(self, "tokenizers", None) if isinstance(tokenizers, list) and tokenizers: diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index 284dcfd..e694868 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -463,6 +463,7 @@ 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) diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index 79d426c..1b0b137 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -462,6 +462,7 @@ 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) diff --git a/tests/test_centralized_actor_critic.py b/tests/test_centralized_actor_critic.py index 5608e12..0b5792c 100644 --- a/tests/test_centralized_actor_critic.py +++ b/tests/test_centralized_actor_critic.py @@ -121,6 +121,22 @@ def metrics(rollouts): 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 From e90bca0e0ef3f865e385880b3c3cad06f2d3547a Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 09:39:01 -0400 Subject: [PATCH 16/21] fix: build config defaults only for fields an old checkpoint lacks Resume validation built every config field's default, even when the checkpoint held them all, so a default_factory that draws random numbers could move the RNG. Defaults are now built only for the missing fields, with the RNG state kept around any factory call. Found in Codex's review. Co-Authored-By: Claude Opus 5.5 --- comlrl/utils/checkpoint.py | 20 ++++++++++++++------ tests/test_checkpoint_state.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/comlrl/utils/checkpoint.py b/comlrl/utils/checkpoint.py index a2577e3..db1f3ee 100644 --- a/comlrl/utils/checkpoint.py +++ b/comlrl/utils/checkpoint.py @@ -294,17 +294,24 @@ def _config(trainer): } -def _config_defaults(trainer): - """Encoded config defaults, for fields added after a checkpoint was written.""" +def _config_defaults(trainer, names): + """Encoded defaults of the config fields ``names`` that a checkpoint lacks.""" args = getattr(trainer, "args", None) - if not is_dataclass(args): + 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: - default = field.default_factory() + # 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: @@ -772,9 +779,10 @@ def _validate_state(trainer, state, *, model_config_updates=None): if key not in _PORTABLE_CONFIG_FIELDS } current_config = _config(trainer) - defaults = _config_defaults(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 set(current_config) - set(saved_config): + 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: diff --git a/tests/test_checkpoint_state.py b/tests/test_checkpoint_state.py index 2a4cdd4..a0bd8b9 100644 --- a/tests/test_checkpoint_state.py +++ b/tests/test_checkpoint_state.py @@ -1,13 +1,14 @@ """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 ( capture_trainer_state, @@ -457,6 +458,35 @@ def test_config_field_newer_than_the_checkpoint_resumes_at_its_default( 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 From c5b5c8c3ef5e027af13faa45598ddf160b88edb7 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 09:46:34 -0400 Subject: [PATCH 17/21] fix: keep DistributedContext and local_context importable, deprecated v1.4.3 exports both from comlrl.utils, so removing them in 9fa1562 broke `from comlrl.utils import local_context` for code written against that release. They are restored as they were, and local_context() now emits a DeprecationWarning. The trainers still do not use them; that part of the cleanup stays. CoMLRL tests 704 -> 705. Co-Authored-By: Claude Opus 5.5 --- comlrl/utils/__init__.py | 8 +++++++- comlrl/utils/distributed.py | 29 ++++++++++++++++++++++++++++- tests/test_utils.py | 16 ++++++++++++++-- 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/comlrl/utils/__init__.py b/comlrl/utils/__init__.py index 04f5b08..5e6d721 100644 --- a/comlrl/utils/__init__.py +++ b/comlrl/utils/__init__.py @@ -16,7 +16,11 @@ resolve_reward_range, set_reward_range, ) -from .distributed import unwrap_model +from .distributed import ( + DistributedContext, + local_context, + unwrap_model, +) from .tokenizer_utils import ( apply_tokenizer_specials, ensure_pad_token, @@ -46,5 +50,7 @@ "normalize_reward_lengths", "resolve_reward_range", "set_reward_range", + "DistributedContext", + "local_context", "unwrap_model", ] diff --git a/comlrl/utils/distributed.py b/comlrl/utils/distributed.py index 695d007..120e8b5 100644 --- a/comlrl/utils/distributed.py +++ b/comlrl/utils/distributed.py @@ -1,4 +1,31 @@ -from typing import Any +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import Any, Optional + +import torch + + +@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) def unwrap_model(model: Any) -> Any: diff --git a/tests/test_utils.py b/tests/test_utils.py index 0872a56..3bf1053 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,11 +1,17 @@ -"""Reward, formatter and W&B helpers in comlrl.utils.""" +"""Reward, formatter, W&B and distributed helpers in comlrl.utils.""" import inspect import pytest import torch -from comlrl.utils import RewardProcessors, build_formatters, call_reward_function +from comlrl.utils import ( + DistributedContext, + RewardProcessors, + build_formatters, + call_reward_function, + local_context, +) from comlrl.utils.wandb_utils import wandb_init_kwargs @@ -151,3 +157,9 @@ def test_wandb_init_kwargs_keeps_every_config_section(): "dataset_name": "data", "dataset_type": "task", } + + +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")) From 8db315b34bf603b6a491f6a7553d24b71b010083 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 11:18:16 -0400 Subject: [PATCH 18/21] docs: point the README to the LeetCode and TL;DR examples examples/leetcode-func-print.py (MAGRPO) and examples/tldr-len-ratio.py (IAC) were not referenced anywhere. Both still run on the current API: an offline smoke with a tiny GPT-2 trains each for one epoch. Co-Authored-By: Claude Opus 5.5 --- README.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/README.md b/README.md index 6215e48..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 From e927f8159c1151f0c68265e5ac193bdbbf02895b Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 14:48:41 -0400 Subject: [PATCH 19/21] feat: optional per-sample backward for actor-critic updates MAAC and IAC hold every sample's graph of a mini-batch until one backward pass. With long prompts that alone overflows a GPU: a Travel MAAC smoke ran out of the critic's 178 GB on four samples. per_sample_backward (default False) backpropagates each sample's share of the batch-mean losses as soon as it is computed and steps once per mini-batch, so only one sample's activations stay in memory. The gradients are the same sums. The default path is unchanged: the analytic gradient test now runs in both modes, and the CPU equivalence scenarios match the 3b0102d baselines (checkpoint scenarios differ only in the new config key and the checkpoint hash). CoMLRL tests 705 -> 711. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 12 ++++++++++ comlrl/trainers/actor_critic/iac.py | 24 +++++++++++++++++++ comlrl/trainers/actor_critic/maac.py | 17 +++++++++++++ .../docs/marl/multi-agent-actor-critic.md | 2 ++ tests/test_actor_critic.py | 11 +++++++-- 5 files changed, 64 insertions(+), 2 deletions(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index ac6e710..76dcdcb 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -58,6 +58,18 @@ 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 equal those of the batch means. + """ + actor_share = policy_loss.mean() / batch_size + value_share = self.args.value_loss_coef * value_error.mean() / batch_size + 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, diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index e694868..c0b8e3b 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -45,6 +45,8 @@ class IACConfig: critic_learning_rate: Optional[float] = 5e-6 rollout_buffer_size: int = 8 train_batch_size: Optional[int] = None + # Same update with one sample's activations in memory instead of a mini-batch's. + per_sample_backward: bool = False value_clip_range: Optional[float] = 0.2 value_loss_coef: float = 0.6 advantage_normalization: bool = True @@ -876,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) @@ -958,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): diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index 1b0b137..1ff436a 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -43,6 +43,8 @@ class MAACConfig: critic_learning_rate: float = 5e-6 rollout_buffer_size: int = 8 train_batch_size: Optional[int] = None + # Same update with one sample's activations in memory instead of a mini-batch's. + per_sample_backward: bool = False value_loss_coef: float = 0.6 advantage_normalization: bool = True max_new_tokens: int = 256 @@ -861,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) @@ -908,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 diff --git a/docs/content/docs/marl/multi-agent-actor-critic.md b/docs/content/docs/marl/multi-agent-actor-critic.md index 280b6e0..47f9626 100644 --- a/docs/content/docs/marl/multi-agent-actor-critic.md +++ b/docs/content/docs/marl/multi-agent-actor-critic.md @@ -45,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, which gives the same update with one sample's activations in memory instead of the whole mini-batch's (default `False`) - `max_new_tokens`: Maximum tokens to generate per completion - `temperature`: Temperature for sampling - `top_p`: Top-p for nucleus sampling @@ -135,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, which gives the same update with one sample's activations in memory instead of the whole mini-batch's (default `False`) - `max_new_tokens`: Maximum new tokens to generate - `temperature`: Temperature for sampling - `top_p`: Top-p for nucleus sampling diff --git a/tests/test_actor_critic.py b/tests/test_actor_critic.py index 8d83bfc..3f6c64b 100644 --- a/tests/test_actor_critic.py +++ b/tests/test_actor_critic.py @@ -252,6 +252,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", [ @@ -263,10 +264,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")] From 123437014823c26b9cd242a20ea20bc87d1f9b44 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 20 Oct 2026 15:01:44 -0400 Subject: [PATCH 20/21] fix: per-sample backward review fixes from Codex - Divide each sample's value loss by the batch size before weighting it, so every share stays within the range of the batch total. Weighting first could overflow fp16 where the batch mean does not (2 * 59840). - Put per_sample_backward last in MAACConfig and IACConfig, so existing positional arguments keep their meaning. - Say "mathematically the same update" in the docs and comments, since floating-point sums may differ slightly. New tests cover the fp16 range and positional configs. CoMLRL tests 711 -> 714; the CPU equivalence scenarios still match the baselines (the checkpoint hash changes with the field order). Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 7 ++++-- comlrl/trainers/actor_critic/iac.py | 5 ++-- comlrl/trainers/actor_critic/maac.py | 5 ++-- .../docs/marl/multi-agent-actor-critic.md | 4 ++-- tests/test_actor_critic.py | 23 +++++++++++++++++++ 5 files changed, 36 insertions(+), 8 deletions(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index 76dcdcb..4a71e92 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -62,10 +62,13 @@ 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 equal those of the batch means. + 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 = self.args.value_loss_coef * value_error.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 diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index c0b8e3b..0901c86 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -45,8 +45,6 @@ class IACConfig: critic_learning_rate: Optional[float] = 5e-6 rollout_buffer_size: int = 8 train_batch_size: Optional[int] = None - # Same update with one sample's activations in memory instead of a mini-batch's. - per_sample_backward: bool = False value_clip_range: Optional[float] = 0.2 value_loss_coef: float = 0.6 advantage_normalization: bool = True @@ -77,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: diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index 1ff436a..d88a81e 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -43,8 +43,6 @@ class MAACConfig: critic_learning_rate: float = 5e-6 rollout_buffer_size: int = 8 train_batch_size: Optional[int] = None - # Same update with one sample's activations in memory instead of a mini-batch's. - per_sample_backward: bool = False value_loss_coef: float = 0.6 advantage_normalization: bool = True max_new_tokens: int = 256 @@ -71,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: diff --git a/docs/content/docs/marl/multi-agent-actor-critic.md b/docs/content/docs/marl/multi-agent-actor-critic.md index 47f9626..9a33ca3 100644 --- a/docs/content/docs/marl/multi-agent-actor-critic.md +++ b/docs/content/docs/marl/multi-agent-actor-critic.md @@ -45,7 +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, which gives the same update with one sample's activations in memory instead of the whole mini-batch's (default `False`) +- `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 @@ -136,7 +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, which gives the same update with one sample's activations in memory instead of the whole mini-batch's (default `False`) +- `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 diff --git a/tests/test_actor_critic.py b/tests/test_actor_critic.py index 3f6c64b..4636141 100644 --- a/tests/test_actor_critic.py +++ b/tests/test_actor_critic.py @@ -331,6 +331,29 @@ def test_ac_step_matches_analytic_gradients( 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.wandb_initialized = True From 45cb2c2f08cfc899535d4b1004ce9b9f2136fd53 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 27 Oct 2026 18:03:19 -0400 Subject: [PATCH 21/21] fix: train actor-critic models in training mode IAC and MAAC never switched their models to training mode, and from_pretrained returns eval mode, in which HF skips the gradient checkpointing a launcher enables. Travel enables it for every actor and critic. Its generation helper restores the wrapper's training flag, so the actors trained in training mode, but the critic never samples and stayed in eval mode without checkpointing: 116 GiB at peak in the Travel MAAC runs. A one-epoch test with the critic in training mode fit both actors and the critic in 103.0 GiB on one B200. _run_epochs now puts every actor and critic in training mode, as MAGRPO's does, and sampling switches the actor to eval mode and back (_generate_in_eval_mode), so dropout stays off and HF keeps its KV cache there. This also covers MARLHF's actor-critic online phase. Models without dropout train exactly as before: on CPU the parameters after two epochs are bit-identical for IAC and MAAC, with and without checkpointing, which now runs for both actor and critic. The new test fails on the old code. CoMLRL tests 714 -> 716; the four LLM_Collab suites pass unchanged. Co-Authored-By: Claude Opus 5.5 --- comlrl/trainers/actor_critic/ac_base.py | 20 +++++++++ comlrl/trainers/actor_critic/iac.py | 3 +- comlrl/trainers/actor_critic/maac.py | 3 +- tests/test_actor_critic.py | 54 ++++++++++++++++++++++++- 4 files changed, 75 insertions(+), 5 deletions(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index 4a71e92..497cc10 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -10,6 +10,7 @@ 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 @@ -210,6 +211,22 @@ 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: @@ -620,6 +637,9 @@ def _train_checkpoint_epochs(self, progress, commit) -> None: 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) diff --git a/comlrl/trainers/actor_critic/iac.py b/comlrl/trainers/actor_critic/iac.py index 0901c86..e6f24b6 100644 --- a/comlrl/trainers/actor_critic/iac.py +++ b/comlrl/trainers/actor_critic/iac.py @@ -468,8 +468,7 @@ def _generate_rollout( 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.") diff --git a/comlrl/trainers/actor_critic/maac.py b/comlrl/trainers/actor_critic/maac.py index d88a81e..5fbfd46 100644 --- a/comlrl/trainers/actor_critic/maac.py +++ b/comlrl/trainers/actor_critic/maac.py @@ -467,8 +467,7 @@ def _generate(self, agent_model, prompt: str, agent_idx: int) -> Dict[str, Any]: 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.") diff --git a/tests/test_actor_critic.py b/tests/test_actor_critic.py index 4636141..a4940c2 100644 --- a/tests/test_actor_critic.py +++ b/tests/test_actor_critic.py @@ -7,13 +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 helpers import tiny_gpt2, zero_reward +from helpers import model_for, tiny_gpt2, zero_reward TRAINERS = [IACTrainer, MAACTrainer] @@ -398,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