feature: support eager as attn

README.md

@@ -34,8 +34,9 @@ Key fields:
 
 `native_parallel/` contains the native single-turn BFCL parallel task:
 
-- categories: `parallel`, `parallel_multiple`, `live_parallel`,
-  `live_parallel_multiple`
+- default categories: `parallel`, `parallel_multiple`
+- live categories remain supported via `dataset.categories` overrides:
+  `live_parallel`, `live_parallel_multiple`
 - reward: flat aggregate joint reward
 - dataset: `OpenMLRL/BFCL-V4-Parallel-Native`
 - config: `native_parallel/configs/native_parallel_magrpo_config.yaml`

multiturn_flat/baseline/single_agent/config.yaml

@@ -4,9 +4,9 @@ agent_model:
   temperature: 0.7
   top_p: 0.8
   top_k: null
-  max_length: 4096
   special_tokens: {}
-  torch_dtype: null
+  torch_dtype: bfloat16
+  attn_implementation: eager
 
 dataset:
   name: OpenMLRL/BFCL-V4-Parallel-Multi-Turn

multiturn_flat/baseline/single_agent/eval_single_agent.py

@@ -262,6 +262,8 @@ def main() -> None:
     dtype = _torch_dtype(model_config.torch_dtype)
     if dtype is not None:
         model_kwargs["torch_dtype"] = dtype
+    if model_config.attn_implementation is not None:
+        model_kwargs["attn_implementation"] = model_config.attn_implementation
     model = AutoModelForCausalLM.from_pretrained(model_config.name, **model_kwargs).to(device)
     model.eval()
 

multiturn_flat/config.py

@@ -17,9 +17,9 @@ class ModelConfig:
     temperature: Optional[float] = None
     top_p: Optional[float] = None
     top_k: Optional[int] = None
-    max_length: int = 4096
     special_tokens: Dict[str, str] = field(default_factory=dict)
     torch_dtype: Optional[str] = None
+    attn_implementation: Optional[str] = "sdpa"
 
     @classmethod
     def from_dict(
@@ -69,9 +69,9 @@ def _as_optional_int(value: Any) -> Optional[int]:
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            max_length=config_dict.get("max_length", 4096),
             special_tokens=config_dict.get("special_tokens", {}),
             torch_dtype=(config_dict.get("torch_dtype") or config_dict.get("dtype")),
+            attn_implementation=config_dict.get("attn_implementation", "sdpa"),
         )
 
 

multiturn_flat/configs/multiturn_flat_magrpo_config.yaml

@@ -4,9 +4,9 @@ agent_model:
   temperature: 0.7
   top_p: 0.8
   top_k: null
-  max_length: 4096
   special_tokens: {}
-  torch_dtype: null
+  torch_dtype: bfloat16
+  attn_implementation: eager
 
 agents: null
 

multiturn_flat/train/train_magrpo.py

@@ -239,6 +239,7 @@ def main() -> None:
         tokenizer=tokenizers if agent_names else tokenizers[0],
         model_config={
             "torch_dtype": model_config.torch_dtype,
+            "attn_implementation": model_config.attn_implementation,
             "special_tokens": model_config.special_tokens,
         },
         train_dataset=train_dataset,

native_parallel/baseline/single_agent/config.yaml

@@ -4,9 +4,9 @@ agent_model:
   temperature: 0.7
   top_p: 0.8
   top_k: null
-  max_length: 4096
   special_tokens: {}
-  torch_dtype: null
+  torch_dtype: bfloat16
+  attn_implementation: eager
 
 dataset:
   name: OpenMLRL/BFCL-V4-Parallel-Native
@@ -16,8 +16,6 @@ dataset:
   categories:
   - parallel
   - parallel_multiple
-  - live_parallel
-  - live_parallel_multiple
   task_types: null
   max_eval_samples: null
 
@@ -39,6 +37,7 @@ bfcl_reward:
   overlap_penalty: 0.20
   lazy_agent_penalty: 0.15
   extra_call_penalty: 0.05
+  self_duplicate_penalty: 0.10
   min_reward: -0.4
   max_reward: 1.2
 

native_parallel/baseline/single_agent/eval_single_agent.py

@@ -156,6 +156,9 @@ def _summarize(records: List[Dict[str, Any]]) -> Dict[str, float]:
         "gold_call_count",
         "balance_score",
         "overlap_rate",
+        "self_duplicate_count",
+        "self_duplicate_rate",
+        "self_duplicate_penalty",
         "lazy_rate",
         "extra_call_rate",
     ]
@@ -242,6 +245,8 @@ def main() -> None:
     dtype = _torch_dtype(model_config.torch_dtype)
     if dtype is not None:
         model_kwargs["torch_dtype"] = dtype
+    if model_config.attn_implementation is not None:
+        model_kwargs["attn_implementation"] = model_config.attn_implementation
     model = AutoModelForCausalLM.from_pretrained(model_config.name, **model_kwargs).to(device)
     model.eval()
 

native_parallel/config.py

@@ -17,9 +17,9 @@ class ModelConfig:
     temperature: Optional[float] = None
     top_p: Optional[float] = None
     top_k: Optional[int] = None
-    max_length: int = 4096
     special_tokens: Dict[str, str] = field(default_factory=dict)
     torch_dtype: Optional[str] = None
+    attn_implementation: Optional[str] = "sdpa"
 
     @classmethod
     def from_dict(
@@ -69,9 +69,9 @@ def _as_optional_int(value: Any) -> Optional[int]:
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            max_length=config_dict.get("max_length", 4096),
             special_tokens=config_dict.get("special_tokens", {}),
             torch_dtype=(config_dict.get("torch_dtype") or config_dict.get("dtype")),
+            attn_implementation=config_dict.get("attn_implementation", "sdpa"),
         )
 
 

native_parallel/configs/native_parallel_magrpo_config.yaml

@@ -4,9 +4,9 @@ agent_model:
   temperature: 0.7
   top_p: 0.8
   top_k: null
-  max_length: 4096
   special_tokens: {}
-  torch_dtype: null
+  torch_dtype: bfloat16
+  attn_implementation: eager
 
 agents: null
 
@@ -19,8 +19,6 @@ dataset:
   categories:
   - parallel
   - parallel_multiple
-  - live_parallel
-  - live_parallel_multiple
   task_types: null
   max_train_samples: null
   max_eval_samples: null
@@ -45,6 +43,7 @@ bfcl_reward:
   overlap_penalty: 0.20
   lazy_agent_penalty: 0.15
   extra_call_penalty: 0.05
+  self_duplicate_penalty: 0.10
   min_reward: -0.4
   max_reward: 1.2
 

native_parallel/data.py

@@ -13,6 +13,10 @@
 DEFAULT_NATIVE_CATEGORIES = (
     "parallel",
     "parallel_multiple",
+)
+
+ALLOWED_NATIVE_CATEGORIES = (
+    *DEFAULT_NATIVE_CATEGORIES,
     "live_parallel",
     "live_parallel_multiple",
 )
@@ -71,11 +75,11 @@ def _filter_rows(
     max_samples: Optional[int] = None,
 ) -> List[Dict[str, Any]]:
     category_set = set(categories)
-    invalid = sorted(category_set - set(DEFAULT_NATIVE_CATEGORIES))
+    invalid = sorted(category_set - set(ALLOWED_NATIVE_CATEGORIES))
     if invalid:
         raise ValueError(
             "Native parallel categories must be one of "
-            f"{list(DEFAULT_NATIVE_CATEGORIES)}; got {invalid}."
+            f"{list(ALLOWED_NATIVE_CATEGORIES)}; got {invalid}."
         )
     task_type_set = set(task_types) if task_types else None
 

native_parallel/logger.py

@@ -87,6 +87,9 @@ def aggregate_native_parallel_metrics(
         "balance_score",
         "overlap_count",
         "overlap_rate",
+        "self_duplicate_count",
+        "self_duplicate_rate",
+        "self_duplicate_penalty",
         "lazy_agents",
         "lazy_rate",
         "extra_call_rate",

native_parallel/rewards/native_parallel_reward.py

@@ -115,6 +115,17 @@ def _cross_agent_overlap(agent_calls: Sequence[Sequence[ToolCall]]) -> int:
     return sum(max(0, count - 1) for count in key_counts.values())
 
 
+def _same_agent_duplicate_count(
+    raw_agent_calls: Sequence[Sequence[ToolCall]],
+    deduped_agent_calls: Sequence[Sequence[ToolCall]],
+) -> Tuple[int, List[int]]:
+    per_agent = [
+        max(0, len(raw_calls) - len(deduped_calls))
+        for raw_calls, deduped_calls in zip(raw_agent_calls, deduped_agent_calls)
+    ]
+    return sum(per_agent), per_agent
+
+
 def _clamp(value: float, lower: float, upper: float) -> float:
     return max(lower, min(upper, value))
 
@@ -130,6 +141,7 @@ class NativeParallelRewardConfig:
     overlap_penalty: float = 0.20
     lazy_agent_penalty: float = 0.15
     extra_call_penalty: float = 0.05
+    self_duplicate_penalty: float = 0.10
     min_reward: float = -0.4
     max_reward: float = 1.2
 
@@ -200,6 +212,11 @@ def score_native_parallel_response(
     balance_score = _balance_score(agent_counts, gold_count)
     overlap_count = _cross_agent_overlap(per_agent_calls)
     overlap_rate = min(1.0, overlap_count / max(1, gold_count))
+    self_duplicate_count, per_agent_duplicate_counts = _same_agent_duplicate_count(
+        raw_per_agent_calls, per_agent_calls
+    )
+    self_duplicate_rate = min(1.0, self_duplicate_count / max(1, gold_count))
+    self_duplicate_penalty = cfg.self_duplicate_penalty * self_duplicate_rate
     lazy_agents = sum(1 for count in agent_counts if count == 0)
     lazy_rate = lazy_agents / max(1, len(agent_counts))
     extra_calls = max(0, pred_count - gold_count)
@@ -217,6 +234,7 @@ def score_native_parallel_response(
         - cfg.overlap_penalty * overlap_rate
         - cfg.lazy_agent_penalty * lazy_rate
         - cfg.extra_call_penalty * extra_rate
+        - self_duplicate_penalty
     )
     reward = _clamp(reward, cfg.min_reward, cfg.max_reward)
 
@@ -235,10 +253,14 @@ def score_native_parallel_response(
         "balance_score": float(balance_score),
         "overlap_count": float(overlap_count),
         "overlap_rate": float(overlap_rate),
+        "self_duplicate_count": float(self_duplicate_count),
+        "self_duplicate_rate": float(self_duplicate_rate),
+        "self_duplicate_penalty": float(self_duplicate_penalty),
         "lazy_agents": float(lazy_agents),
         "lazy_rate": float(lazy_rate),
         "extra_call_rate": float(extra_rate),
         "agent_call_counts": agent_counts,
+        "per_agent_duplicate_counts": per_agent_duplicate_counts,
         "combined_calls": [
             {"name": call.name, "arguments": call.arguments} for call in combined_calls
         ],

native_parallel/train/train_magrpo.py

@@ -239,6 +239,7 @@ def main() -> None:
         tokenizer=tokenizers if agent_names else tokenizers[0],
         model_config={
             "torch_dtype": model_config.torch_dtype,
+            "attn_implementation": model_config.attn_implementation,
             "special_tokens": model_config.special_tokens,
         },
         train_dataset=train_dataset,