diff --git a/README.md b/README.md index 9abb710..714f224 100644 --- a/README.md +++ b/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` diff --git a/multiturn_flat/baseline/single_agent/config.yaml b/multiturn_flat/baseline/single_agent/config.yaml index d335da2..c31ce50 100644 --- a/multiturn_flat/baseline/single_agent/config.yaml +++ b/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 diff --git a/multiturn_flat/baseline/single_agent/eval_single_agent.py b/multiturn_flat/baseline/single_agent/eval_single_agent.py index 780fc77..84e9efd 100644 --- a/multiturn_flat/baseline/single_agent/eval_single_agent.py +++ b/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() diff --git a/multiturn_flat/config.py b/multiturn_flat/config.py index 78da109..1e3bac7 100644 --- a/multiturn_flat/config.py +++ b/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"), ) diff --git a/multiturn_flat/configs/multiturn_flat_magrpo_config.yaml b/multiturn_flat/configs/multiturn_flat_magrpo_config.yaml index 31a981d..026eaaa 100644 --- a/multiturn_flat/configs/multiturn_flat_magrpo_config.yaml +++ b/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 diff --git a/multiturn_flat/train/train_magrpo.py b/multiturn_flat/train/train_magrpo.py index af4e367..10b345c 100644 --- a/multiturn_flat/train/train_magrpo.py +++ b/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, diff --git a/native_parallel/baseline/single_agent/config.yaml b/native_parallel/baseline/single_agent/config.yaml index 6a78a33..8d0fe02 100644 --- a/native_parallel/baseline/single_agent/config.yaml +++ b/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 diff --git a/native_parallel/baseline/single_agent/eval_single_agent.py b/native_parallel/baseline/single_agent/eval_single_agent.py index 710dfe6..dc45f4c 100644 --- a/native_parallel/baseline/single_agent/eval_single_agent.py +++ b/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() diff --git a/native_parallel/config.py b/native_parallel/config.py index 7526031..38b8cbc 100644 --- a/native_parallel/config.py +++ b/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"), ) diff --git a/native_parallel/configs/native_parallel_magrpo_config.yaml b/native_parallel/configs/native_parallel_magrpo_config.yaml index f19123e..543223f 100644 --- a/native_parallel/configs/native_parallel_magrpo_config.yaml +++ b/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 diff --git a/native_parallel/data.py b/native_parallel/data.py index c42dfdf..1e4d99e 100644 --- a/native_parallel/data.py +++ b/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 diff --git a/native_parallel/logger.py b/native_parallel/logger.py index 115f625..9cf3e1f 100644 --- a/native_parallel/logger.py +++ b/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", diff --git a/native_parallel/rewards/native_parallel_reward.py b/native_parallel/rewards/native_parallel_reward.py index ccd2c39..3755c9b 100644 --- a/native_parallel/rewards/native_parallel_reward.py +++ b/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 ], diff --git a/native_parallel/train/train_magrpo.py b/native_parallel/train/train_magrpo.py index 59d10fe..117e61a 100644 --- a/native_parallel/train/train_magrpo.py +++ b/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,