From 252d36499cbffdb42cb11ba52f9c89650dbc3679 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 7 Jul 2026 21:27:44 -0400 Subject: [PATCH 1/3] fix: add eager attn --- comlrl/trainers/actor_critic/ac_base.py | 22 +++++++++++++++++++++- comlrl/trainers/preference/iterative.py | 12 ++++++++++++ comlrl/trainers/preference/marlhf.py | 16 ++++++++++++++++ comlrl/trainers/reinforce/magrpo.py | 12 ++++++++++++ 4 files changed, 61 insertions(+), 1 deletion(-) diff --git a/comlrl/trainers/actor_critic/ac_base.py b/comlrl/trainers/actor_critic/ac_base.py index 6eebe7d..d05c39a 100644 --- a/comlrl/trainers/actor_critic/ac_base.py +++ b/comlrl/trainers/actor_critic/ac_base.py @@ -58,13 +58,33 @@ def _run_agent_tasks( def _filter_model_kwargs(self, cfg: Optional[Dict[str, Any]]) -> Dict[str, Any]: torch_dtype = None + attn_implementation = "sdpa" + cfg_has_attn = isinstance(cfg, dict) and "attn_implementation" in cfg if isinstance(cfg, dict): torch_dtype = cfg.get("torch_dtype") or cfg.get("dtype") + if cfg_has_attn: + attn_implementation = cfg.get("attn_implementation") if torch_dtype is None: model_cfg = getattr(self, "model_config", None) if isinstance(model_cfg, dict): torch_dtype = model_cfg.get("torch_dtype") or model_cfg.get("dtype") - return {"torch_dtype": torch_dtype} if torch_dtype is not None else {} + if not cfg_has_attn: + model_cfg = getattr(self, "model_config", None) + if isinstance(model_cfg, dict): + nested_model_kwargs = model_cfg.get("model_kwargs") + if "attn_implementation" in model_cfg: + attn_implementation = model_cfg.get("attn_implementation") + elif ( + isinstance(nested_model_kwargs, dict) + and "attn_implementation" in nested_model_kwargs + ): + attn_implementation = nested_model_kwargs.get("attn_implementation") + model_kwargs = {} + if torch_dtype is not None: + model_kwargs["torch_dtype"] = torch_dtype + if attn_implementation is not None: + model_kwargs["attn_implementation"] = attn_implementation + return model_kwargs def _format_prompt( self, diff --git a/comlrl/trainers/preference/iterative.py b/comlrl/trainers/preference/iterative.py index 11c8048..4165841 100644 --- a/comlrl/trainers/preference/iterative.py +++ b/comlrl/trainers/preference/iterative.py @@ -2672,6 +2672,18 @@ def _comparator_model_kwargs(self) -> Dict[str, Any]: ) if torch_dtype is not None: model_kwargs["torch_dtype"] = torch_dtype + attn_implementation = "sdpa" + if isinstance(self.model_config, dict): + nested_model_kwargs = self.model_config.get("model_kwargs") + if "attn_implementation" in self.model_config: + attn_implementation = self.model_config.get("attn_implementation") + elif ( + isinstance(nested_model_kwargs, dict) + and "attn_implementation" in nested_model_kwargs + ): + attn_implementation = nested_model_kwargs.get("attn_implementation") + if attn_implementation is not None: + model_kwargs["attn_implementation"] = attn_implementation return model_kwargs @staticmethod diff --git a/comlrl/trainers/preference/marlhf.py b/comlrl/trainers/preference/marlhf.py index 702d729..1469f8c 100644 --- a/comlrl/trainers/preference/marlhf.py +++ b/comlrl/trainers/preference/marlhf.py @@ -405,6 +405,22 @@ def _model_kwargs_from_config( torch_dtype = nested.get("torch_dtype") or nested.get("dtype") if torch_dtype is not None: model_kwargs["torch_dtype"] = torch_dtype + attn_implementation = "sdpa" + nested = config.get("model_kwargs") + if "attn_implementation" in config: + attn_implementation = config.get("attn_implementation") + elif isinstance(nested, dict) and "attn_implementation" in nested: + attn_implementation = nested.get("attn_implementation") + elif source_config is not None and isinstance(self.model_config, dict): + nested = self.model_config.get("model_kwargs") + if "attn_implementation" in self.model_config: + attn_implementation = self.model_config.get("attn_implementation") + elif isinstance(nested, dict) and "attn_implementation" in nested: + attn_implementation = nested.get("attn_implementation") + if attn_implementation is not None: + model_kwargs["attn_implementation"] = attn_implementation + else: + model_kwargs["attn_implementation"] = "sdpa" return model_kwargs def _resolve_critic_model_source(self): diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index 7a8d432..785ca73 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -249,6 +249,18 @@ def __init__( ) if torch_dtype is not None: model_kwargs["torch_dtype"] = torch_dtype + attn_implementation = "sdpa" + if isinstance(self.model_config, dict): + nested_model_kwargs = self.model_config.get("model_kwargs") + if "attn_implementation" in self.model_config: + attn_implementation = self.model_config.get("attn_implementation") + elif ( + isinstance(nested_model_kwargs, dict) + and "attn_implementation" in nested_model_kwargs + ): + attn_implementation = nested_model_kwargs.get("attn_implementation") + if attn_implementation is not None: + model_kwargs["attn_implementation"] = attn_implementation if actor_sources and all(isinstance(src, str) for src in actor_sources): from transformers import AutoModelForCausalLM From be2b590bfd3fee4f699e07bfac18bc56352db1e9 Mon Sep 17 00:00:00 2001 From: N!no Date: Tue, 7 Jul 2026 21:37:32 -0400 Subject: [PATCH 2/3] Update magrpo.py --- comlrl/trainers/reinforce/magrpo.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index 785ca73..ce5d775 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -1222,7 +1222,6 @@ def _generate_completions( "input_ids": prompt_input_ids, "attention_mask": prompt_attention_mask, "max_new_tokens": max_new_tokens, - "output_scores": True, "return_dict_in_generate": True, } @@ -1241,14 +1240,15 @@ def _generate_completions( kwargs.pop("do_sample", None) generation_kwargs.update(kwargs) - generation_output = agent_module.generate(**generation_kwargs) + with torch.inference_mode(): + generation_output = agent_module.generate(**generation_kwargs) except Exception as e: raise ValueError(f"Generation failed: {str(e)}") - - agent.train(training_mode) - for name, param in agent.named_parameters(): - if name in original_requires_grad: - param.requires_grad = original_requires_grad[name] + finally: + agent.train(training_mode) + for name, param in agent.named_parameters(): + if name in original_requires_grad: + param.requires_grad = original_requires_grad[name] completion_input_ids = generation_output.sequences @@ -1294,9 +1294,6 @@ def _generate_completions( batch_masks.append(mask) completion_attention_masks.append(batch_masks) - logits = ( - generation_output.scores if hasattr(generation_output, "scores") else [] - ) reference_kls = self._reference_kl_values( agent_idx, agent_module, @@ -1316,7 +1313,6 @@ def _generate_completions( "completion_attention_mask": completion_attention_masks, "response_lens": batch_response_lens, "reference_kls": reference_kls, - "logits": logits, } def _generate_completions_with_external_prompts( From 6281345b0803317e8c587aad1ebe4357c43fae50 Mon Sep 17 00:00:00 2001 From: N!no Date: Wed, 8 Jul 2026 10:49:00 -0400 Subject: [PATCH 3/3] use_cache=False --- comlrl/trainers/reinforce/magrpo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/comlrl/trainers/reinforce/magrpo.py b/comlrl/trainers/reinforce/magrpo.py index ce5d775..f171971 100644 --- a/comlrl/trainers/reinforce/magrpo.py +++ b/comlrl/trainers/reinforce/magrpo.py @@ -1701,6 +1701,7 @@ def _compute_loss_with_gradients(self, agent, completions_data, returns): outputs = agent( input_ids=input_ids.unsqueeze(0), # Add batch dimension attention_mask=attention_mask.unsqueeze(0), # Add batch dimension + use_cache=False, ) # Get logits for the completion part (excluding prompt)