Repository navigation
feature: enable and test different forms of centralized MAS - #73
Open
LovelyBuggies wants to merge 21 commits into
Open
LovelyBuggies wants to merge 21 commits into
LovelyBuggies wants to merge 21 commits into
Conversation
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
… 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Member
Author
|
Given the different coding styles across developers, our massive codebase, and our shortage of hands, I've had to start using Claude to keep our code style and feature set consistent.This PR is primarily to compare the performance of two different centralized multi-agent systems under various algorithms. Also, I've polished the style of the entire codebase using Claude Code. |
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.