feature: support eager as attn

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,

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

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

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
 
@@ -1210,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,
             }
 
@@ -1229,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
 
@@ -1282,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,
@@ -1304,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(
@@ -1693,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)