3653 lines
168 KiB
Python
3653 lines
168 KiB
Python
import os
|
||
import sys
|
||
|
||
__package__ = "trainer"
|
||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||
|
||
import argparse
|
||
import ast
|
||
import gc
|
||
import json
|
||
import math
|
||
import queue
|
||
import random
|
||
import re
|
||
import signal
|
||
import threading
|
||
import time
|
||
import warnings
|
||
from collections import deque
|
||
from contextlib import nullcontext
|
||
from dataclasses import dataclass
|
||
from datetime import timedelta
|
||
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
||
|
||
import torch
|
||
import torch.distributed as dist
|
||
import torch.nn.functional as F
|
||
from torch import optim
|
||
from torch.nn.parallel import DistributedDataParallel
|
||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||
from torch.utils.data import Dataset
|
||
|
||
from model.model_minimind import MiniMindConfig
|
||
from trainer.rollout_engine import create_rollout_engine, RolloutResult
|
||
from trainer.trainer_utils import (
|
||
LMForRewardModel,
|
||
Logger,
|
||
init_model,
|
||
is_main_process,
|
||
lm_checkpoint,
|
||
setup_seed,
|
||
)
|
||
|
||
warnings.filterwarnings("ignore")
|
||
|
||
|
||
MATH_AGENT_SYSTEM_INSTRUCTION = (
|
||
"你是一个简洁、严格的工具调用助手。遇到算术表达式时，必须先调用 calculate_math 工具；"
|
||
"不要用其它工具解算术题，不要自己心算。工具调用格式必须严格为："
|
||
'<tool_call>{"name":"calculate_math","arguments":{"expression":"1+1"}}</tool_call>'
|
||
"。每个独立算式调用一次工具，按用户问题顺序调用。拿到全部工具结果后，最终只输出 "
|
||
"<answer>结果1, 结果2</answer>；多个结果按问题顺序用英文逗号分隔，不要输出额外解释。"
|
||
)
|
||
|
||
CALCULATE_MATH_TOOL_SPEC = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "calculate_math",
|
||
"description": "Evaluate an arithmetic expression. 用于计算加减乘除、幂、sqrt/log/sin/cos 等简单数学表达式。",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"expression": {"type": "string", "description": "Arithmetic expression, e.g. 123+456 or 2**10."}
|
||
},
|
||
"required": ["expression"],
|
||
},
|
||
},
|
||
}
|
||
|
||
UNIT_CONVERTER_TOOL_SPEC = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "unit_converter",
|
||
"description": "Convert common units. 仅在用户明确要求单位换算时使用。",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"value": {"type": "number"},
|
||
"from_unit": {"type": "string"},
|
||
"to_unit": {"type": "string"},
|
||
},
|
||
"required": ["value", "from_unit", "to_unit"],
|
||
},
|
||
},
|
||
}
|
||
|
||
TIME_TOOL_SPEC = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_current_time",
|
||
"description": "Return a fixed or live current time for a timezone. 仅在用户明确询问时间时使用。",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {"timezone": {"type": "string"}},
|
||
"required": [],
|
||
},
|
||
},
|
||
}
|
||
|
||
CALCULATOR_ALIAS_TOOL_SPEC = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "calculator",
|
||
"description": "Backward-compatible alias for calculate_math. Prefer calculate_math.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {"expression": {"type": "string"}},
|
||
"required": ["expression"],
|
||
},
|
||
},
|
||
}
|
||
|
||
|
||
# -------------------------- generic helpers --------------------------
|
||
|
||
def str2bool(v):
|
||
if isinstance(v, bool):
|
||
return v
|
||
return str(v).lower() in ("true", "1", "yes", "y", "on")
|
||
|
||
|
||
def unwrap_model(model):
|
||
raw = model.module if isinstance(model, DistributedDataParallel) else model
|
||
return getattr(raw, "_orig_mod", raw)
|
||
|
||
|
||
def finite_float(x: Any, default: float = 0.0) -> float:
|
||
try:
|
||
y = float(x)
|
||
return y if math.isfinite(y) else default
|
||
except Exception:
|
||
return default
|
||
|
||
|
||
def safe_mean(xs: Iterable[float], default: float = 0.0) -> float:
|
||
vals = [finite_float(x, default=float("nan")) for x in xs]
|
||
vals = [x for x in vals if math.isfinite(x)]
|
||
return sum(vals) / max(len(vals), 1) if vals else default
|
||
|
||
|
||
def percentile(xs: Iterable[float], q: float, default: float = 0.0) -> float:
|
||
vals = sorted(finite_float(x, default=float("nan")) for x in xs)
|
||
vals = [x for x in vals if math.isfinite(x)]
|
||
if not vals:
|
||
return default
|
||
if len(vals) == 1:
|
||
return float(vals[0])
|
||
q = min(max(float(q), 0.0), 1.0)
|
||
pos = q * (len(vals) - 1)
|
||
lo = int(math.floor(pos))
|
||
hi = int(math.ceil(pos))
|
||
if lo == hi:
|
||
return float(vals[lo])
|
||
frac = pos - lo
|
||
return float(vals[lo] * (1.0 - frac) + vals[hi] * frac)
|
||
|
||
|
||
def ddp_world_size() -> int:
|
||
return dist.get_world_size() if dist.is_initialized() else 1
|
||
|
||
|
||
def trainer_rank() -> int:
|
||
return 0
|
||
|
||
|
||
def is_trainer_rollout_mode(args) -> bool:
|
||
return bool(getattr(args, "multi_gpu_mode", "trainer_rollout") == "trainer_rollout")
|
||
|
||
|
||
def is_trainer_process(args) -> bool:
|
||
rank = dist.get_rank() if dist.is_initialized() else 0
|
||
return rank == trainer_rank() if is_trainer_rollout_mode(args) else True
|
||
|
||
|
||
def ddp_all_true(flag: bool, device: str) -> bool:
|
||
if not dist.is_initialized():
|
||
return bool(flag)
|
||
x = torch.tensor([1 if flag else 0], device=device, dtype=torch.int64)
|
||
dist.all_reduce(x, op=dist.ReduceOp.MIN)
|
||
return bool(x.item() > 0)
|
||
|
||
|
||
def ddp_mean_scalar(value: float, device: str) -> float:
|
||
if not dist.is_initialized():
|
||
return float(value)
|
||
x = torch.tensor([finite_float(value)], device=device, dtype=torch.float32)
|
||
dist.all_reduce(x, op=dist.ReduceOp.SUM)
|
||
x /= max(ddp_world_size(), 1)
|
||
return float(x.item())
|
||
|
||
|
||
def ddp_sum_scalar(value: float, device: str) -> float:
|
||
if not dist.is_initialized():
|
||
return float(value)
|
||
x = torch.tensor([finite_float(value)], device=device, dtype=torch.float32)
|
||
dist.all_reduce(x, op=dist.ReduceOp.SUM)
|
||
return float(x.item())
|
||
|
||
|
||
def ddp_average_metrics(metrics: Dict[str, Any], device: str) -> Dict[str, Any]:
|
||
if not dist.is_initialized():
|
||
return metrics
|
||
keys = [k for k, v in metrics.items() if isinstance(v, (int, float)) and math.isfinite(float(v))]
|
||
if not keys:
|
||
return metrics
|
||
values = torch.tensor([float(metrics[k]) for k in keys], device=device, dtype=torch.float32)
|
||
dist.all_reduce(values, op=dist.ReduceOp.SUM)
|
||
values /= max(ddp_world_size(), 1)
|
||
out = dict(metrics)
|
||
for key, value in zip(keys, values.tolist()):
|
||
out[key] = float(value)
|
||
return out
|
||
|
||
|
||
def assert_rollout_alignment(response_ids, response_mask, response_old_logps, response_versions):
|
||
n = len(response_ids)
|
||
if not (n == len(response_mask) == len(response_old_logps) == len(response_versions)):
|
||
raise AssertionError(
|
||
f"rollout alignment mismatch: ids={len(response_ids)}, mask={len(response_mask)}, "
|
||
f"logps={len(response_old_logps)}, versions={len(response_versions)}"
|
||
)
|
||
|
||
|
||
def init_distributed_mode_compat(backend: str = "auto"):
|
||
if int(os.environ.get("RANK", -1)) == -1:
|
||
return 0
|
||
|
||
local_rank = int(os.environ["LOCAL_RANK"])
|
||
if torch.cuda.is_available():
|
||
torch.cuda.set_device(local_rank)
|
||
|
||
def _init_pg(pg_backend: str):
|
||
dist.init_process_group(backend=pg_backend, timeout=timedelta(minutes=30))
|
||
if pg_backend == "nccl":
|
||
probe = torch.ones(1, device=f"cuda:{local_rank}")
|
||
dist.all_reduce(probe)
|
||
torch.cuda.synchronize(local_rank)
|
||
|
||
backend = str(backend or "auto").lower()
|
||
if backend == "auto":
|
||
preferred = "nccl" if torch.cuda.is_available() else "gloo"
|
||
candidates = [preferred] if preferred == "gloo" else ["nccl", "gloo"]
|
||
else:
|
||
candidates = [backend]
|
||
|
||
for idx, candidate in enumerate(candidates):
|
||
try:
|
||
_init_pg(candidate)
|
||
if candidate != "nccl" and int(os.environ.get("RANK", "0")) == 0:
|
||
Logger(f"[WARN] distributed backend fallback to {candidate}")
|
||
return local_rank
|
||
except Exception as exc:
|
||
if dist.is_initialized():
|
||
try:
|
||
dist.destroy_process_group()
|
||
except Exception:
|
||
pass
|
||
if idx + 1 >= len(candidates):
|
||
raise
|
||
if int(os.environ.get("RANK", "0")) == 0:
|
||
Logger(f"[WARN] distributed backend {candidate} unavailable: {repr(exc)}; fallback to {candidates[idx + 1]}")
|
||
return local_rank
|
||
|
||
|
||
def ensure_tokenizer_ids(tokenizer, lm_config=None):
|
||
if tokenizer.eos_token_id is None:
|
||
tokenizer.eos_token_id = getattr(lm_config, "eos_token_id", 2)
|
||
if tokenizer.pad_token_id is None:
|
||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||
return tokenizer.pad_token_id, tokenizer.eos_token_id
|
||
|
||
|
||
@torch.no_grad()
|
||
def model_has_nonfinite(model: torch.nn.Module, max_checks: int = 32) -> bool:
|
||
raw = unwrap_model(model)
|
||
checked = 0
|
||
for p in raw.parameters():
|
||
if p is None:
|
||
continue
|
||
if not torch.isfinite(p.detach()).all():
|
||
return True
|
||
checked += 1
|
||
if checked >= max_checks:
|
||
break
|
||
return False
|
||
|
||
|
||
@torch.no_grad()
|
||
def copy_policy_weights(dst_model: torch.nn.Module, src_model: torch.nn.Module) -> bool:
|
||
if dst_model is None:
|
||
return False
|
||
if model_has_nonfinite(src_model):
|
||
Logger("[WARN] train model has non-finite weights; skip rollout policy sync")
|
||
return False
|
||
src = unwrap_model(src_model)
|
||
dst = unwrap_model(dst_model)
|
||
dst.load_state_dict(src.state_dict(), strict=False)
|
||
dst.to(next(src.parameters()).device)
|
||
dst.eval().requires_grad_(False)
|
||
return True
|
||
|
||
|
||
def save_checkpoint(args, lm_config, model, optimizer, scheduler, epoch: int, step: int, opt_step: int, version: int, wandb=None):
|
||
if not is_main_process():
|
||
return
|
||
if model_has_nonfinite(model):
|
||
Logger("[WARN] checkpoint skipped because model has non-finite weights")
|
||
return
|
||
os.makedirs(args.save_dir, exist_ok=True)
|
||
moe_suffix = "_moe" if lm_config.use_moe else ""
|
||
ckp = f"{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth"
|
||
raw_model = unwrap_model(model)
|
||
state_dict = raw_model.state_dict()
|
||
tmp = ckp + ".tmp"
|
||
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, tmp)
|
||
os.replace(tmp, ckp)
|
||
lm_checkpoint(
|
||
lm_config,
|
||
weight=args.save_weight,
|
||
model=model,
|
||
optimizer=optimizer,
|
||
epoch=epoch,
|
||
step=step,
|
||
wandb=wandb,
|
||
save_dir="../checkpoints",
|
||
scheduler=scheduler,
|
||
opt_step=opt_step,
|
||
policy_version=version,
|
||
reward_baseline=float(getattr(args, "reward_baseline", 0.0)),
|
||
)
|
||
Logger(f"checkpoint saved: {ckp}")
|
||
del state_dict
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
|
||
def apply_chat_template_compat(tokenizer, messages, tools=None, add_generation_prompt=True, open_thinking=False) -> str:
|
||
if hasattr(tokenizer, "apply_chat_template"):
|
||
call_orders = [
|
||
dict(tokenize=False, add_generation_prompt=add_generation_prompt, tools=tools, open_thinking=open_thinking),
|
||
dict(tokenize=False, add_generation_prompt=add_generation_prompt, tools=tools),
|
||
dict(tokenize=False, add_generation_prompt=add_generation_prompt),
|
||
]
|
||
for kwargs in call_orders:
|
||
try:
|
||
return tokenizer.apply_chat_template(messages, **kwargs)
|
||
except TypeError:
|
||
continue
|
||
except Exception:
|
||
continue
|
||
parts = []
|
||
if tools:
|
||
parts.append("<|tools|>" + json.dumps(tools, ensure_ascii=False, separators=(",", ":")) + "<|/tools|>")
|
||
for msg in messages:
|
||
parts.append(f"<|im_start|>{msg.get('role', 'user')}\n{msg.get('content', '')}<|im_end|>")
|
||
if add_generation_prompt:
|
||
parts.append("<|im_start|>assistant\n")
|
||
return "\n".join(parts)
|
||
|
||
|
||
def tokenize_context(tokenizer, text: str, device: str, max_context_len: int = 0):
|
||
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
|
||
input_ids = inputs["input_ids"]
|
||
attention_mask = inputs["attention_mask"]
|
||
if max_context_len and input_ids.size(1) > max_context_len:
|
||
input_ids = input_ids[:, -max_context_len:]
|
||
attention_mask = attention_mask[:, -max_context_len:]
|
||
return {"input_ids": input_ids.to(device), "attention_mask": attention_mask.to(device)}
|
||
|
||
|
||
# -------------------------- rollout engine --------------------------
|
||
|
||
class StableTorchRolloutEngine:
|
||
"""Safe torch rollout.
|
||
|
||
PPO fix: old_logps are computed under the raw policy distribution. Sampling may use
|
||
temperature/top-p/top-k/repetition penalty, but PPO behavior logprobs are taken from
|
||
the unmodified policy distribution.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
policy_model,
|
||
tokenizer,
|
||
device="cuda",
|
||
autocast_factory=None,
|
||
top_p=0.90,
|
||
top_k=50,
|
||
repetition_penalty=1.03,
|
||
use_cache=True,
|
||
):
|
||
self.policy_model = policy_model
|
||
self.tokenizer = tokenizer
|
||
self.device = device
|
||
self.autocast_factory = autocast_factory or (lambda: nullcontext())
|
||
self.top_p = float(top_p)
|
||
self.top_k = int(top_k)
|
||
self.repetition_penalty = float(repetition_penalty)
|
||
self.use_cache = bool(use_cache)
|
||
self.nan_batches = 0
|
||
self.bad_prob_batches = 0
|
||
self.cache_prefill_tokens = 0
|
||
self.cache_decode_tokens = 0
|
||
self.cache_fallbacks = 0
|
||
|
||
def update_policy(self, model):
|
||
self.policy_model = model
|
||
return True
|
||
|
||
@staticmethod
|
||
def _sanitize_logits(logits: torch.Tensor) -> torch.Tensor:
|
||
return torch.nan_to_num(logits.float(), nan=0.0, posinf=30.0, neginf=-30.0).clamp(-80.0, 80.0)
|
||
|
||
def _apply_repetition_penalty_for_sampling(self, sample_logits, input_ids):
|
||
if not self.repetition_penalty or self.repetition_penalty == 1.0:
|
||
return sample_logits
|
||
sample_logits = sample_logits.clone()
|
||
for i in range(input_ids.size(0)):
|
||
uniq = torch.unique(input_ids[i])
|
||
uniq = uniq[(uniq >= 0) & (uniq < sample_logits.size(-1))]
|
||
if uniq.numel() > 0:
|
||
positive = sample_logits[i, uniq] > 0
|
||
sample_logits[i, uniq] = torch.where(
|
||
positive,
|
||
sample_logits[i, uniq] / self.repetition_penalty,
|
||
sample_logits[i, uniq] * self.repetition_penalty,
|
||
)
|
||
return sample_logits
|
||
|
||
def _sample_next(self, logits, input_ids, temperature, do_sample, eos_token_id):
|
||
had_bad = (~torch.isfinite(logits)).any().item()
|
||
raw_logits = self._sanitize_logits(logits)
|
||
raw_log_probs = F.log_softmax(raw_logits, dim=-1)
|
||
sample_logits = self._apply_repetition_penalty_for_sampling(raw_logits, input_ids)
|
||
|
||
if (not do_sample) or temperature <= 1e-6:
|
||
next_token = torch.argmax(sample_logits, dim=-1, keepdim=True)
|
||
next_logp = raw_log_probs.gather(1, next_token).squeeze(1)
|
||
return next_token, next_logp, bool(had_bad), False
|
||
|
||
sample_logits = sample_logits / max(float(temperature), 1e-6)
|
||
vocab_size = sample_logits.size(-1)
|
||
top_k = min(max(self.top_k, 0), vocab_size)
|
||
if 0 < top_k < vocab_size:
|
||
threshold = torch.topk(sample_logits, top_k, dim=-1).values[..., -1, None]
|
||
sample_logits = sample_logits.masked_fill(sample_logits < threshold, -1e9)
|
||
|
||
if 0.0 < self.top_p < 1.0:
|
||
sorted_logits, sorted_indices = torch.sort(sample_logits, descending=True, dim=-1)
|
||
sorted_probs = torch.softmax(sorted_logits, dim=-1)
|
||
sorted_probs = torch.nan_to_num(sorted_probs, nan=0.0, posinf=0.0, neginf=0.0)
|
||
sorted_remove = torch.cumsum(sorted_probs, dim=-1) > self.top_p
|
||
sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
|
||
sorted_remove[..., 0] = False
|
||
remove = torch.zeros_like(sample_logits, dtype=torch.bool)
|
||
remove.scatter_(1, sorted_indices, sorted_remove)
|
||
sample_logits = sample_logits.masked_fill(remove, -1e9)
|
||
|
||
probs = torch.softmax(sample_logits, dim=-1)
|
||
probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0).clamp(min=0.0)
|
||
denom = probs.sum(dim=-1, keepdim=True)
|
||
bad_rows = (~torch.isfinite(denom)) | (denom <= 0)
|
||
if bad_rows.any():
|
||
probs = torch.where(bad_rows, torch.zeros_like(probs), probs)
|
||
fallback_id = int(eos_token_id) if eos_token_id is not None else 0
|
||
fallback_id = min(max(fallback_id, 0), probs.size(-1) - 1)
|
||
probs[bad_rows.squeeze(1), fallback_id] = 1.0
|
||
denom = probs.sum(dim=-1, keepdim=True)
|
||
probs = probs / denom.clamp(min=1e-12)
|
||
next_token = torch.multinomial(probs, num_samples=1)
|
||
next_logp = raw_log_probs.gather(1, next_token).squeeze(1)
|
||
next_logp = torch.nan_to_num(next_logp, nan=0.0, posinf=0.0, neginf=0.0)
|
||
return next_token, next_logp, bool(had_bad), bool(bad_rows.any().item())
|
||
|
||
@torch.no_grad()
|
||
def rollout(self, prompt_ids, attention_mask, num_generations, max_new_tokens, temperature=0.8):
|
||
model = unwrap_model(self.policy_model).eval()
|
||
eos_id = self.tokenizer.eos_token_id if self.tokenizer.eos_token_id is not None else 2
|
||
pad_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else eos_id
|
||
input_ids = prompt_ids.repeat_interleave(num_generations, dim=0).to(self.device)
|
||
attn = attention_mask.repeat_interleave(num_generations, dim=0).to(self.device) if attention_mask is not None else None
|
||
self.cache_prefill_tokens += int(input_ids.numel())
|
||
vocab_size = getattr(getattr(model, "config", None), "vocab_size", None)
|
||
if vocab_size is not None:
|
||
bad_ids = (input_ids < 0) | (input_ids >= int(vocab_size))
|
||
if bad_ids.any():
|
||
input_ids = input_ids.clone()
|
||
input_ids[bad_ids] = eos_id
|
||
|
||
generated_tokens, generated_logps = [], []
|
||
past_key_values = None
|
||
finished = torch.zeros(input_ids.size(0), dtype=torch.bool, device=input_ids.device)
|
||
do_sample = temperature > 1e-6
|
||
use_cache = bool(self.use_cache)
|
||
|
||
for _ in range(int(max_new_tokens)):
|
||
step_input = input_ids[:, -1:] if (use_cache and past_key_values is not None) else input_ids
|
||
try:
|
||
with self.autocast_factory():
|
||
outputs = model(
|
||
input_ids=step_input,
|
||
attention_mask=attn,
|
||
past_key_values=past_key_values if use_cache else None,
|
||
use_cache=use_cache,
|
||
)
|
||
except RuntimeError as exc:
|
||
msg = str(exc)
|
||
if use_cache and ("attention_mask" in msg or "size of tensor" in msg or "must match" in msg):
|
||
self.cache_fallbacks += 1
|
||
use_cache = False
|
||
past_key_values = None
|
||
with self.autocast_factory():
|
||
outputs = model(input_ids=input_ids, attention_mask=attn, past_key_values=None, use_cache=False)
|
||
else:
|
||
raise
|
||
|
||
if use_cache and past_key_values is not None:
|
||
self.cache_decode_tokens += int(step_input.numel())
|
||
|
||
next_token, next_logp, had_bad, bad_prob = self._sample_next(
|
||
outputs.logits[:, -1, :], input_ids, temperature, do_sample, eos_id
|
||
)
|
||
self.nan_batches += int(had_bad)
|
||
self.bad_prob_batches += int(bad_prob)
|
||
forced_eos = torch.full_like(next_token, int(eos_id))
|
||
next_token = torch.where(finished.unsqueeze(-1), forced_eos, next_token)
|
||
next_logp = torch.where(finished, torch.zeros_like(next_logp), next_logp)
|
||
input_ids = torch.cat([input_ids, next_token], dim=-1)
|
||
generated_tokens.append(next_token)
|
||
generated_logps.append(next_logp.unsqueeze(1))
|
||
if attn is not None:
|
||
attn = torch.cat([attn, attn.new_ones(attn.size(0), 1)], dim=-1)
|
||
past_key_values = getattr(outputs, "past_key_values", None) if use_cache else None
|
||
finished |= next_token.squeeze(-1).eq(int(eos_id))
|
||
if finished.all():
|
||
break
|
||
|
||
if generated_tokens:
|
||
completion_ids = torch.cat(generated_tokens, dim=1)
|
||
per_token_logps = torch.cat(generated_logps, dim=1)
|
||
else:
|
||
completion_ids = input_ids.new_full((input_ids.size(0), 0), int(pad_id))
|
||
per_token_logps = input_ids.new_empty((input_ids.size(0), 0), dtype=torch.float32)
|
||
output_ids = torch.cat([prompt_ids.repeat_interleave(num_generations, dim=0).to(self.device), completion_ids], dim=1)
|
||
completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
|
||
return RolloutResult(output_ids=output_ids, completion_ids=completion_ids, per_token_logps=per_token_logps, completions=completions)
|
||
|
||
|
||
# -------------------------- logprob / entropy helpers --------------------------
|
||
|
||
|
||
def compute_per_token_logps_safe(model, input_ids, attention_mask, autocast_factory=None):
|
||
raw = unwrap_model(model)
|
||
ctx_factory = autocast_factory or (lambda: nullcontext())
|
||
with ctx_factory():
|
||
res = raw(input_ids=input_ids, attention_mask=attention_mask)
|
||
logits = res.logits[:, :-1, :]
|
||
logits = torch.nan_to_num(logits.float(), nan=0.0, posinf=30.0, neginf=-30.0).clamp(-80.0, 80.0)
|
||
token_ids = input_ids[:, 1:].clamp(min=0, max=logits.size(-1) - 1)
|
||
return F.log_softmax(logits, dim=-1).gather(2, token_ids.unsqueeze(-1)).squeeze(-1)
|
||
|
||
|
||
def logits_entropy_safe(logits, mask):
|
||
with torch.no_grad():
|
||
logits = torch.nan_to_num(logits.float(), nan=0.0, posinf=30.0, neginf=-30.0).clamp(-80.0, 80.0)
|
||
logp = F.log_softmax(logits, dim=-1)
|
||
ent = -(torch.exp(logp) * logp).sum(dim=-1)
|
||
return float(((ent * mask).sum() / mask.sum().clamp(min=1.0)).item())
|
||
|
||
|
||
# -------------------------- safe math and robust extraction --------------------------
|
||
|
||
class SafeMathEvaluator(ast.NodeVisitor):
|
||
ALLOWED_BINOPS = {
|
||
ast.Add: lambda a, b: a + b,
|
||
ast.Sub: lambda a, b: a - b,
|
||
ast.Mult: lambda a, b: a * b,
|
||
ast.Div: lambda a, b: a / b,
|
||
ast.FloorDiv: lambda a, b: a // b,
|
||
ast.Mod: lambda a, b: a % b,
|
||
ast.Pow: lambda a, b: a ** b,
|
||
}
|
||
ALLOWED_UNARY = {ast.UAdd: lambda a: +a, ast.USub: lambda a: -a}
|
||
ALLOWED_NAMES = {"pi": math.pi, "e": math.e}
|
||
ALLOWED_FUNCS = {
|
||
"sqrt": math.sqrt,
|
||
"sin": math.sin,
|
||
"cos": math.cos,
|
||
"tan": math.tan,
|
||
"log": math.log,
|
||
"exp": math.exp,
|
||
"abs": abs,
|
||
"round": round,
|
||
}
|
||
|
||
def visit_Expression(self, node):
|
||
return self.visit(node.body)
|
||
|
||
def visit_Constant(self, node):
|
||
if isinstance(node.value, (int, float)):
|
||
return float(node.value)
|
||
raise ValueError("invalid constant")
|
||
|
||
def visit_Num(self, node):
|
||
return float(node.n)
|
||
|
||
def visit_Name(self, node):
|
||
if node.id in self.ALLOWED_NAMES:
|
||
return float(self.ALLOWED_NAMES[node.id])
|
||
raise ValueError(f"invalid name: {node.id}")
|
||
|
||
def visit_Call(self, node):
|
||
if not isinstance(node.func, ast.Name) or node.func.id not in self.ALLOWED_FUNCS:
|
||
raise ValueError("invalid function")
|
||
if node.keywords:
|
||
raise ValueError("keyword args not allowed")
|
||
if len(node.args) > 3:
|
||
raise ValueError("too many args")
|
||
return float(self.ALLOWED_FUNCS[node.func.id](*[self.visit(arg) for arg in node.args]))
|
||
|
||
def visit_BinOp(self, node):
|
||
op = self.ALLOWED_BINOPS.get(type(node.op))
|
||
if op is None:
|
||
raise ValueError("invalid operator")
|
||
return float(op(self.visit(node.left), self.visit(node.right)))
|
||
|
||
def visit_UnaryOp(self, node):
|
||
op = self.ALLOWED_UNARY.get(type(node.op))
|
||
if op is None:
|
||
raise ValueError("invalid unary operator")
|
||
return float(op(self.visit(node.operand)))
|
||
|
||
def generic_visit(self, node):
|
||
raise ValueError(f"unsupported node: {type(node).__name__}")
|
||
|
||
|
||
def _build_fullwidth_translation() -> Dict[int, str]:
|
||
mapping = {ord(src): dst for src, dst in {
|
||
"＋": "+",
|
||
"－": "-",
|
||
"−": "-",
|
||
"–": "-",
|
||
"—": "-",
|
||
"﹣": "-",
|
||
"×": "*",
|
||
"✕": "*",
|
||
"＊": "*",
|
||
"÷": "/",
|
||
"／": "/",
|
||
"（": "(",
|
||
"）": ")",
|
||
"［": "[",
|
||
"］": "]",
|
||
"｛": "{",
|
||
"｝": "}",
|
||
"，": ",",
|
||
"。": ".",
|
||
"．": ".",
|
||
"：": ":",
|
||
"；": ";",
|
||
"％": "%",
|
||
"＾": "^",
|
||
" ": " ",
|
||
}.items()}
|
||
for src, dst in zip("０１２３４５６７８９", "0123456789"):
|
||
mapping[ord(src)] = dst
|
||
return mapping
|
||
|
||
|
||
_FULLWIDTH_TRANS = _build_fullwidth_translation()
|
||
|
||
|
||
def canonicalize_expression(expr: str) -> str:
|
||
s = str(expr).translate(_FULLWIDTH_TRANS)
|
||
s = s.replace("^", "**")
|
||
s = re.sub(r"\s+", "", s)
|
||
return s
|
||
|
||
|
||
def _strip_expr_edges(expr: str) -> str:
|
||
s = str(expr).translate(_FULLWIDTH_TRANS).strip()
|
||
s = s.strip(" \t\r\n:;,.?？!！、，。；：")
|
||
# Trim unmatched right parens that appear from casual text while preserving valid groups.
|
||
while s and s.count(")") > s.count("("):
|
||
s = s[:-1].rstrip()
|
||
return s
|
||
|
||
|
||
def safe_eval_math(expr: str) -> float:
|
||
expr = canonicalize_expression(expr)
|
||
if len(expr) > 256:
|
||
raise ValueError("expression too long")
|
||
# Avoid interpreting dates as subtraction chains.
|
||
if re.fullmatch(r"\d{4}-\d{1,2}-\d{1,2}", expr):
|
||
raise ValueError("date-like string is not an arithmetic expression")
|
||
return SafeMathEvaluator().visit(ast.parse(expr, mode="eval"))
|
||
|
||
|
||
def format_number(value: float) -> str:
|
||
value = float(value)
|
||
if abs(value - round(value)) < 1e-9:
|
||
return str(int(round(value)))
|
||
return f"{value:.10f}".rstrip("0").rstrip(".")
|
||
|
||
|
||
_NUM = r"[-+]?\d+(?:\.\d+)?"
|
||
_FUNC = r"(?:sqrt|sin|cos|tan|log|exp|abs|round)"
|
||
# Non-nested parentheses are enough for the synthetic data and most simple tool-use data.
|
||
_ATOM = rf"(?:{_FUNC}\s*\([^()]+\)|\([^()]+\)|{_NUM})"
|
||
_OP = r"(?:\*\*|//|[+\-*/×÷^%])"
|
||
# ASCII boundary only. Do not use \w boundary: in Python, Chinese characters are \w and would hide expressions after “计算”.
|
||
ARITH_EXPR_RE = re.compile(rf"(?<![A-Za-z0-9_]){_ATOM}(?:\s*{_OP}\s*{_ATOM})+(?![A-Za-z0-9_])", re.IGNORECASE)
|
||
|
||
|
||
def _candidate_has_real_operator(expr: str) -> bool:
|
||
s = canonicalize_expression(expr)
|
||
# Remove leading unary sign before checking for a binary operator.
|
||
s2 = re.sub(r"(?<=^)[+-]", "", s)
|
||
return any(op in s2 for op in ("**", "//", "+", "-", "*", "/", "%"))
|
||
|
||
|
||
def extract_arithmetic_expressions(text: str, max_exprs: int = 8) -> List[str]:
|
||
"""Extract simple arithmetic expressions from mixed Chinese/English text.
|
||
|
||
v3 used a boundary based on \\w and therefore missed expressions immediately after
|
||
Chinese words, e.g. “请计算884+645*18” -> only “645*18”. It also truncated
|
||
parenthesized expressions, e.g. “(289+532)*11” -> “(289+532)”. Both errors poison
|
||
BC targets and reward diagnostics. This extractor keeps complete arithmetic spans
|
||
and validates every candidate with SafeMathEvaluator.
|
||
"""
|
||
raw = str(text).translate(_FULLWIDTH_TRANS)
|
||
candidates: List[Tuple[int, int, str]] = []
|
||
for m in ARITH_EXPR_RE.finditer(raw):
|
||
expr = _strip_expr_edges(m.group(0))
|
||
if not expr or not _candidate_has_real_operator(expr):
|
||
continue
|
||
try:
|
||
safe_eval_math(expr)
|
||
except Exception:
|
||
continue
|
||
candidates.append((m.start(), m.end(), expr))
|
||
|
||
# Prefer longer overlapping spans to avoid keeping “289+532” when “(289+532)*11” exists.
|
||
candidates.sort(key=lambda x: (x[0], -(x[1] - x[0])))
|
||
selected: List[Tuple[int, int, str]] = []
|
||
for start, end, expr in candidates:
|
||
overlaps = [i for i, (s, e, _) in enumerate(selected) if not (end <= s or start >= e)]
|
||
if overlaps:
|
||
# Keep the candidate covering more characters.
|
||
keep = True
|
||
for i in overlaps:
|
||
s, e, _ = selected[i]
|
||
if (e - s) >= (end - start):
|
||
keep = False
|
||
break
|
||
if not keep:
|
||
continue
|
||
selected = [x for j, x in enumerate(selected) if j not in overlaps]
|
||
selected.append((start, end, expr))
|
||
selected.sort(key=lambda x: x[0])
|
||
if len(selected) >= max_exprs:
|
||
break
|
||
|
||
out: List[str] = []
|
||
seen = set()
|
||
for _, _, expr in selected:
|
||
key = canonicalize_expression(expr)
|
||
if key not in seen:
|
||
out.append(expr)
|
||
seen.add(key)
|
||
if len(out) >= max_exprs:
|
||
break
|
||
return out
|
||
|
||
|
||
# -------------------------- tool env --------------------------
|
||
|
||
class RealToolEnv:
|
||
def __init__(self, time_mode="fixed", active_tools: Optional[List[str]] = None, allow_calculator_alias: bool = False):
|
||
self.time_mode = time_mode
|
||
self.active_tools = active_tools or ["calculate_math", "unit_converter", "get_current_time"]
|
||
self.allow_calculator_alias = bool(allow_calculator_alias)
|
||
self.time_data = {
|
||
"Asia/Shanghai": "2025-03-07 14:30:00",
|
||
"America/New_York": "2025-03-07 01:30:00",
|
||
"Europe/London": "2025-03-07 06:30:00",
|
||
"Asia/Tokyo": "2025-03-07 15:30:00",
|
||
"Europe/Paris": "2025-03-07 07:30:00",
|
||
"Australia/Sydney": "2025-03-07 17:30:00",
|
||
}
|
||
self.unit_data = {
|
||
"km_miles": 0.621371,
|
||
"miles_km": 1.60934,
|
||
"kg_pounds": 2.20462,
|
||
"pounds_kg": 0.453592,
|
||
"meters_feet": 3.28084,
|
||
"feet_meters": 0.3048,
|
||
}
|
||
|
||
def tool_specs(self):
|
||
spec_map = {
|
||
"calculate_math": CALCULATE_MATH_TOOL_SPEC,
|
||
"unit_converter": UNIT_CONVERTER_TOOL_SPEC,
|
||
"get_current_time": TIME_TOOL_SPEC,
|
||
}
|
||
out = []
|
||
for name in self.active_tools:
|
||
if name in spec_map:
|
||
out.append(clone_tool_spec(spec_map[name]))
|
||
if self.allow_calculator_alias and "calculate_math" in self.active_tools:
|
||
out.append(clone_tool_spec(CALCULATOR_ALIAS_TOOL_SPEC))
|
||
return out
|
||
|
||
def supported_names(self):
|
||
names = set(supported_tool_names(self.tool_specs())) | {"python_math"}
|
||
if self.allow_calculator_alias:
|
||
names.add("calculator")
|
||
return names
|
||
|
||
def execute(self, name, args):
|
||
if name in ("calculate_math", "calculator", "python_math"):
|
||
expr = str(args.get("expression", args.get("query", ""))).strip()
|
||
if not expr:
|
||
raise ValueError("missing expression")
|
||
value = safe_eval_math(expr)
|
||
return {"value": value, "result": format_number(value)}
|
||
if name == "unit_converter":
|
||
value = float(args["value"])
|
||
from_unit = str(args["from_unit"]).lower().strip()
|
||
to_unit = str(args["to_unit"]).lower().strip()
|
||
key = f"{from_unit}_{to_unit}"
|
||
if key == "celsius_fahrenheit":
|
||
result = value * 1.8 + 32.0
|
||
elif key == "fahrenheit_celsius":
|
||
result = (value - 32.0) / 1.8
|
||
elif key in self.unit_data:
|
||
result = value * self.unit_data[key]
|
||
else:
|
||
raise ValueError(f"unsupported conversion: {key}")
|
||
return {"value": value, "from_unit": from_unit, "to_unit": to_unit, "result": format_number(round(result, 6))}
|
||
if name == "get_current_time":
|
||
timezone = str(args.get("timezone", "Asia/Shanghai") or "Asia/Shanghai")
|
||
if self.time_mode == "live":
|
||
return {"timezone": timezone, "datetime": time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()), "mode": "live_utc"}
|
||
return {"timezone": timezone, "datetime": self.time_data.get(timezone, self.time_data["Asia/Shanghai"]), "mode": "fixed"}
|
||
raise ValueError(f"unknown tool: {name}")
|
||
|
||
|
||
# -------------------------- schema/data helpers --------------------------
|
||
|
||
def normalize_tool_schema(raw_tools):
|
||
if raw_tools is None or raw_tools == "":
|
||
return None
|
||
if isinstance(raw_tools, str):
|
||
try:
|
||
raw_tools = json.loads(raw_tools)
|
||
except Exception:
|
||
return None
|
||
if isinstance(raw_tools, dict):
|
||
raw_tools = [raw_tools]
|
||
if not isinstance(raw_tools, list):
|
||
return None
|
||
tools = []
|
||
for item in raw_tools:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
if item.get("type") == "function" and isinstance(item.get("function"), dict):
|
||
fn = dict(item["function"])
|
||
elif isinstance(item.get("function"), dict):
|
||
fn = dict(item["function"])
|
||
elif "name" in item:
|
||
fn = {
|
||
"name": item.get("name"),
|
||
"description": item.get("description", ""),
|
||
"parameters": item.get("parameters", {"type": "object", "properties": {}}),
|
||
}
|
||
else:
|
||
continue
|
||
if isinstance(fn.get("parameters"), str):
|
||
try:
|
||
fn["parameters"] = json.loads(fn["parameters"])
|
||
except Exception:
|
||
fn["parameters"] = {"type": "object", "properties": {}}
|
||
if fn.get("name"):
|
||
tools.append({"type": "function", "function": fn})
|
||
return tools or None
|
||
|
||
|
||
def supported_tool_names(tool_specs):
|
||
if not tool_specs:
|
||
return []
|
||
out = []
|
||
for tool in tool_specs:
|
||
try:
|
||
name = tool.get("function", {}).get("name")
|
||
if name:
|
||
out.append(name)
|
||
except Exception:
|
||
pass
|
||
return out
|
||
|
||
|
||
def clone_tool_spec(tool):
|
||
return json.loads(json.dumps(tool, ensure_ascii=False))
|
||
|
||
|
||
def tool_name_set(tool_specs):
|
||
return set(supported_tool_names(tool_specs))
|
||
|
||
|
||
def merge_tool_specs(primary, extra, keep_alias=False, allowed_names: Optional[set] = None):
|
||
merged = []
|
||
seen = set()
|
||
for src in (primary or []), (extra or []):
|
||
for tool in src:
|
||
name = tool.get("function", {}).get("name") if isinstance(tool, dict) else None
|
||
if not name or name in seen:
|
||
continue
|
||
if name == "calculator" and not keep_alias:
|
||
continue
|
||
if allowed_names is not None and name not in allowed_names:
|
||
continue
|
||
merged.append(clone_tool_spec(tool))
|
||
seen.add(name)
|
||
return merged or None
|
||
|
||
|
||
def resolve_tool_specs(sample_tools, default_tools, force_all_tools=False, keep_alias=False, allowed_names: Optional[set] = None):
|
||
sample_only = merge_tool_specs(sample_tools, None, keep_alias=keep_alias, allowed_names=allowed_names)
|
||
default_only = merge_tool_specs(default_tools, None, keep_alias=keep_alias, allowed_names=allowed_names)
|
||
if force_all_tools:
|
||
return merge_tool_specs(sample_only, default_only, keep_alias=keep_alias, allowed_names=allowed_names)
|
||
return sample_only or default_only
|
||
|
||
|
||
def append_calculator_alias(tools):
|
||
if not tools:
|
||
return tools
|
||
names = tool_name_set(tools)
|
||
if "calculate_math" in names and "calculator" not in names:
|
||
tools = list(tools) + [clone_tool_spec(CALCULATOR_ALIAS_TOOL_SPEC)]
|
||
return tools
|
||
|
||
|
||
def prompt_text_from_messages(messages: List[Dict[str, Any]], roles: Tuple[str, ...] = ("system", "user")) -> str:
|
||
role_set = set(roles)
|
||
return "\n".join(str(m.get("content", "")) for m in messages if m.get("role") in role_set)
|
||
|
||
|
||
def task_text_from_messages(messages: Optional[List[Dict[str, Any]]]) -> str:
|
||
if not messages:
|
||
return ""
|
||
user_text = prompt_text_from_messages(messages, roles=("user",))
|
||
if user_text.strip():
|
||
return user_text
|
||
return prompt_text_from_messages(messages, roles=("system", "user"))
|
||
|
||
|
||
def task_text_from_prompt(prompt: str) -> str:
|
||
raw = re.sub(r"<\|tools\|>.*?<\|/tools\|>", "", str(prompt), flags=re.DOTALL)
|
||
matches = re.findall(r"<\|im_start\|>(system|user|assistant|tool)\n(.*?)<\|im_end\|>", raw, flags=re.DOTALL)
|
||
if matches:
|
||
user_parts = [content.strip() for role, content in matches if role == "user" and content.strip()]
|
||
if user_parts:
|
||
return "\n".join(user_parts)
|
||
su_parts = [content.strip() for role, content in matches if role in ("system", "user") and content.strip()]
|
||
if su_parts:
|
||
return "\n".join(su_parts)
|
||
return raw.strip()
|
||
|
||
|
||
def parse_tool_calls(text):
|
||
calls = []
|
||
raw_text = str(text)
|
||
|
||
for match in re.finditer(r"<tool_call>(.*?)</tool_call>", raw_text, flags=re.DOTALL):
|
||
raw = match.group(1).strip()
|
||
try:
|
||
obj = json.loads(raw)
|
||
name = obj.get("name") or obj.get("tool_name") or obj.get("function", {}).get("name")
|
||
args = obj.get("arguments", obj.get("args", {}))
|
||
if isinstance(args, str):
|
||
try:
|
||
args = json.loads(args)
|
||
except Exception:
|
||
args = {"expression": args}
|
||
calls.append({"name": name or "", "arguments": args if isinstance(args, dict) else {}, "canonical": True})
|
||
except Exception:
|
||
calls.append({"name": "", "arguments": {}, "parse_error": raw[:512], "canonical": False})
|
||
|
||
if calls:
|
||
return calls
|
||
|
||
# Tolerant exploration parser: calculate_math("1+1") or calculator(expression="(1+2)*3").
|
||
for name, body in extract_function_style_tool_calls(raw_text, ("calculate_math", "calculator", "python_math")):
|
||
expr = ""
|
||
q = re.search(r"""expression\s*=\s*(['"])(.*?)\1""", body, flags=re.DOTALL)
|
||
if q:
|
||
expr = q.group(2)
|
||
else:
|
||
q = re.search(r"""(['"])(.*?)\1""", body, flags=re.DOTALL)
|
||
if q:
|
||
expr = q.group(2)
|
||
expr = str(expr).strip()
|
||
if expr:
|
||
calls.append({"name": name, "arguments": {"expression": expr}, "canonical": False})
|
||
return calls
|
||
|
||
|
||
def extract_function_style_tool_calls(text: str, names: Tuple[str, ...]) -> List[Tuple[str, str]]:
|
||
out: List[Tuple[str, str]] = []
|
||
raw = str(text)
|
||
cursor = 0
|
||
while cursor < len(raw):
|
||
match = re.search(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(", raw[cursor:])
|
||
if not match:
|
||
break
|
||
name = match.group(1)
|
||
if name not in names:
|
||
cursor += match.end()
|
||
continue
|
||
open_idx = cursor + match.end() - 1
|
||
close_idx = find_matching_paren(raw, open_idx)
|
||
if close_idx < 0:
|
||
cursor = open_idx + 1
|
||
continue
|
||
out.append((name, raw[open_idx + 1:close_idx]))
|
||
cursor = close_idx + 1
|
||
return out
|
||
|
||
|
||
def find_matching_paren(text: str, open_idx: int) -> int:
|
||
if open_idx < 0 or open_idx >= len(text) or text[open_idx] != "(":
|
||
return -1
|
||
depth = 0
|
||
quote = ""
|
||
escaped = False
|
||
for idx in range(open_idx, len(text)):
|
||
ch = text[idx]
|
||
if quote:
|
||
if escaped:
|
||
escaped = False
|
||
continue
|
||
if ch == "\\":
|
||
escaped = True
|
||
continue
|
||
if ch == quote:
|
||
quote = ""
|
||
continue
|
||
if ch in ("'", '"'):
|
||
quote = ch
|
||
continue
|
||
if ch == "(":
|
||
depth += 1
|
||
continue
|
||
if ch == ")":
|
||
depth -= 1
|
||
if depth == 0:
|
||
return idx
|
||
return -1
|
||
|
||
|
||
class PromptPoolDataset(Dataset):
|
||
def __init__(
|
||
self,
|
||
jsonl_path,
|
||
max_samples=0,
|
||
filter_mode="all",
|
||
require_gt=True,
|
||
allowed_tools=None,
|
||
inject_answer_instruction=True,
|
||
augment_calculator_alias=False,
|
||
default_tools=None,
|
||
force_all_tools=True,
|
||
):
|
||
self.jsonl_path = jsonl_path
|
||
self.max_samples = int(max_samples)
|
||
self.filter_mode = filter_mode
|
||
self.require_gt = bool(require_gt)
|
||
self.allowed_tools = set(allowed_tools or [])
|
||
self.inject_answer_instruction = bool(inject_answer_instruction)
|
||
self.augment_calculator_alias = bool(augment_calculator_alias)
|
||
self.default_tools = normalize_tool_schema(default_tools) or []
|
||
self.force_all_tools = bool(force_all_tools)
|
||
self.samples = []
|
||
self.stats = {
|
||
"loaded": 0,
|
||
"kept": 0,
|
||
"bad_json": 0,
|
||
"dropped": 0,
|
||
"skip_no_gt": 0,
|
||
"skip_unsupported_tools": 0,
|
||
"skip_filter": 0,
|
||
"augmented_calculator_alias": 0,
|
||
"forced_default_tools": 0,
|
||
"defaulted_tools": 0,
|
||
}
|
||
self._load()
|
||
|
||
def _maybe_inject_instruction(self, messages, tool_names):
|
||
if not self.inject_answer_instruction or "calculate_math" not in tool_names:
|
||
return messages
|
||
messages = [dict(m) for m in messages]
|
||
for msg in messages:
|
||
if msg.get("role") == "system":
|
||
content = str(msg.get("content", "")).strip()
|
||
if "calculate_math" not in content or "<answer>" not in content:
|
||
msg["content"] = (content + "\n" + MATH_AGENT_SYSTEM_INSTRUCTION).strip()
|
||
return messages
|
||
return [{"role": "system", "content": MATH_AGENT_SYSTEM_INSTRUCTION}] + messages
|
||
|
||
def _normalize(self, raw):
|
||
tools = normalize_tool_schema(raw.get("tools")) or normalize_tool_schema(raw.get("functions"))
|
||
gt = raw.get("gt", raw.get("answer", raw.get("answers", raw.get("label", []))))
|
||
if gt is None:
|
||
gt = []
|
||
if not isinstance(gt, list):
|
||
gt = [gt]
|
||
gt = [x for x in gt if str(x).strip()]
|
||
|
||
messages = raw.get("messages")
|
||
if not messages and isinstance(raw.get("conversations"), list):
|
||
messages = []
|
||
for msg in raw["conversations"]:
|
||
if not isinstance(msg, dict):
|
||
continue
|
||
role = msg.get("role", msg.get("from", ""))
|
||
role = {"human": "user", "gpt": "assistant"}.get(role, role)
|
||
content = msg.get("content", msg.get("value", ""))
|
||
if msg.get("tools") is not None:
|
||
tools = normalize_tool_schema(msg.get("tools")) or tools
|
||
if msg.get("functions") is not None:
|
||
tools = normalize_tool_schema(msg.get("functions")) or tools
|
||
messages.append({"role": role, "content": str(content)})
|
||
if messages and messages[-1].get("role") == "assistant" and not str(messages[-1].get("content", "")).strip():
|
||
messages = messages[:-1]
|
||
|
||
if not messages:
|
||
prompt = raw.get("prompt") or raw.get("question") or raw.get("query")
|
||
if prompt:
|
||
messages = [{"role": "user", "content": str(prompt)}]
|
||
if not isinstance(messages, list) or not messages:
|
||
return None
|
||
|
||
norm_messages = []
|
||
for msg in messages:
|
||
if not isinstance(msg, dict):
|
||
continue
|
||
role = msg.get("role", "user")
|
||
role = {"human": "user", "gpt": "assistant"}.get(role, role)
|
||
if role not in ("system", "user", "assistant", "tool"):
|
||
role = "user"
|
||
norm_messages.append({"role": role, "content": str(msg.get("content", msg.get("value", "")))})
|
||
if not norm_messages:
|
||
return None
|
||
|
||
allowed_names = set(self.allowed_tools) if self.allowed_tools else None
|
||
sample_tool_count = len(tool_name_set(merge_tool_specs(tools, None, keep_alias=self.augment_calculator_alias, allowed_names=allowed_names)))
|
||
default_tool_count = len(tool_name_set(merge_tool_specs(self.default_tools, None, keep_alias=self.augment_calculator_alias, allowed_names=allowed_names)))
|
||
tools = resolve_tool_specs(
|
||
tools,
|
||
self.default_tools,
|
||
force_all_tools=self.force_all_tools,
|
||
keep_alias=self.augment_calculator_alias,
|
||
allowed_names=allowed_names,
|
||
)
|
||
after_n = len(tool_name_set(tools))
|
||
if sample_tool_count <= 0 and default_tool_count > 0 and after_n > 0:
|
||
self.stats["defaulted_tools"] += 1
|
||
if self.force_all_tools and after_n > sample_tool_count:
|
||
self.stats["forced_default_tools"] += 1
|
||
|
||
if self.augment_calculator_alias:
|
||
before_names = tool_name_set(tools)
|
||
tools = append_calculator_alias(tools)
|
||
after_names = tool_name_set(tools)
|
||
if "calculator" not in before_names and "calculator" in after_names:
|
||
self.stats["augmented_calculator_alias"] += 1
|
||
|
||
if self.require_gt and not gt:
|
||
self.stats["skip_no_gt"] += 1
|
||
return None
|
||
tool_names = supported_tool_names(tools)
|
||
if self.allowed_tools and any(name not in self.allowed_tools for name in tool_names):
|
||
self.stats["skip_unsupported_tools"] += 1
|
||
return None
|
||
if self.filter_mode == "math_only":
|
||
text = task_text_from_messages(norm_messages)
|
||
has_math_expr = bool(extract_arithmetic_expressions(text, max_exprs=1))
|
||
if not has_math_expr:
|
||
self.stats["skip_filter"] += 1
|
||
return None
|
||
norm_messages = self._maybe_inject_instruction(norm_messages, tool_names)
|
||
return {"messages": norm_messages, "tools": tools, "gt": gt, "raw": raw}
|
||
|
||
def _load(self):
|
||
with open(self.jsonl_path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
if self.max_samples > 0 and len(self.samples) >= self.max_samples:
|
||
break
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
self.stats["loaded"] += 1
|
||
try:
|
||
raw = json.loads(line)
|
||
except Exception:
|
||
self.stats["bad_json"] += 1
|
||
continue
|
||
sample = self._normalize(raw)
|
||
if sample is None:
|
||
self.stats["dropped"] += 1
|
||
continue
|
||
self.samples.append(sample)
|
||
self.stats["kept"] += 1
|
||
if not self.samples:
|
||
raise ValueError(f"No valid samples found in {self.jsonl_path}; stats={json.dumps(self.stats, ensure_ascii=False)}")
|
||
|
||
def __len__(self):
|
||
return len(self.samples)
|
||
|
||
def __getitem__(self, idx):
|
||
return self.samples[idx]
|
||
|
||
|
||
class PromptPoolSubset(Dataset):
|
||
def __init__(self, samples):
|
||
self.samples = samples
|
||
self.stats = {"kept": len(samples)}
|
||
|
||
def __len__(self):
|
||
return len(self.samples)
|
||
|
||
def __getitem__(self, idx):
|
||
return self.samples[idx]
|
||
|
||
|
||
def split_prompt_pool(dataset, eval_ratio, eval_max_samples, seed=42):
|
||
if eval_ratio <= 0 or len(dataset) < 10:
|
||
return dataset, PromptPoolSubset([])
|
||
indices = list(range(len(dataset)))
|
||
random.Random(seed).shuffle(indices)
|
||
eval_n = min(max(1, int(len(dataset) * eval_ratio)), int(eval_max_samples), len(dataset) - 1)
|
||
eval_ids = set(indices[:eval_n])
|
||
train_samples = [s for i, s in enumerate(dataset.samples) if i not in eval_ids]
|
||
eval_samples = [s for i, s in enumerate(dataset.samples) if i in eval_ids]
|
||
dataset.samples = train_samples
|
||
dataset.stats["kept"] = len(train_samples)
|
||
return dataset, PromptPoolSubset(eval_samples)
|
||
|
||
|
||
def build_fixed_eval_indices(eval_ds, eval_samples, seed=1234):
|
||
n = min(int(eval_samples), len(eval_ds))
|
||
if n <= 0:
|
||
return []
|
||
indices = list(range(len(eval_ds)))
|
||
random.Random(seed).shuffle(indices)
|
||
return indices[:n]
|
||
|
||
|
||
def shard_prompt_pool(dataset, world_size, rank):
|
||
if world_size <= 1 or dataset is None or len(dataset) == 0:
|
||
return dataset
|
||
samples = [sample for idx, sample in enumerate(dataset.samples) if idx % world_size == rank]
|
||
if not samples:
|
||
samples = [dataset.samples[rank % len(dataset.samples)]]
|
||
return PromptPoolSubset(samples)
|
||
|
||
|
||
# -------------------------- answer and reward helpers --------------------------
|
||
|
||
def strip_answer(text):
|
||
text = str(text)
|
||
text = re.sub(r"<\|im_start\|>[^\n]*\n?", "", text)
|
||
text = re.sub(r"<\|im_end\|>", "", text)
|
||
if "</think>" in text:
|
||
text = text.split("</think>", 1)[-1]
|
||
return text.strip()
|
||
|
||
|
||
def remove_tool_blocks(text):
|
||
cleaned = re.sub(r"<tool_call>.*?</tool_call>", "", str(text), flags=re.DOTALL)
|
||
cleaned = re.sub(r"<tool_response>.*?</tool_response>", "", cleaned, flags=re.DOTALL)
|
||
return strip_answer(cleaned).strip()
|
||
|
||
|
||
def answer_tag_stats(text):
|
||
clean = remove_tool_blocks(text)
|
||
open_n = len(re.findall(r"<answer>", clean, flags=re.IGNORECASE))
|
||
close_n = len(re.findall(r"</answer>", clean, flags=re.IGNORECASE))
|
||
matches = re.findall(r"<answer>(.*?)</answer>", clean, flags=re.DOTALL | re.IGNORECASE)
|
||
has = 1.0 if matches else 0.0
|
||
exactly_one = 1.0 if len(matches) == 1 and open_n == 1 and close_n == 1 else 0.0
|
||
content = strip_answer(matches[-1]) if matches else clean
|
||
nonempty = 1.0 if content.strip() else 0.0
|
||
malformed = abs(open_n - close_n) + max(0, open_n - len(matches)) + max(0, close_n - len(matches))
|
||
return {
|
||
"has_answer_tag": has,
|
||
"exactly_one_answer_tag": exactly_one,
|
||
"answer_tag_nonempty": nonempty,
|
||
"malformed_answer_tags": float(malformed),
|
||
"answer_content": content,
|
||
"clean_text": clean,
|
||
}
|
||
|
||
|
||
def extract_answer_region(text):
|
||
return answer_tag_stats(text)["answer_content"]
|
||
|
||
|
||
def numeric_candidates(text):
|
||
nums = []
|
||
for match in re.finditer(r"[-+]?\d+(?:,\d{3})*(?:\.\d+)?|[-+]?\d+(?:\.\d+)?", str(text)):
|
||
try:
|
||
nums.append(float(match.group(0).replace(",", "")))
|
||
except Exception:
|
||
pass
|
||
return nums
|
||
|
||
|
||
def normalize_text(x):
|
||
return re.sub(r"\s+", " ", str(x).lower().strip().replace(",", ""))
|
||
|
||
|
||
def _as_float_or_none(x):
|
||
try:
|
||
s = str(x).strip().replace(",", "")
|
||
if not re.fullmatch(r"[-+]?\d+(?:\.\d+)?", s):
|
||
return None
|
||
return float(s)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def numeric_gt_items(gt_items):
|
||
vals = []
|
||
for gt in gt_items or []:
|
||
f = _as_float_or_none(gt)
|
||
if f is None:
|
||
return []
|
||
vals.append(f)
|
||
return vals
|
||
|
||
|
||
def num_close(a, b):
|
||
return abs(float(a) - float(b)) <= max(1e-6, abs(float(b)) * 1e-6)
|
||
|
||
|
||
def ordered_numeric_metrics(answer_region, gt_items):
|
||
pred_nums = numeric_candidates(answer_region)
|
||
gt_nums = numeric_gt_items(gt_items)
|
||
if not gt_nums:
|
||
return {
|
||
"pred_nums": pred_nums,
|
||
"gt_nums": gt_nums,
|
||
"ordered_match_count": 0.0,
|
||
"ordered_cov": 0.0,
|
||
"strict_exact_match": 0.0,
|
||
"order_ok": 0.0,
|
||
"extra_numbers": float(len(pred_nums)),
|
||
"missing_numbers": 0.0,
|
||
}
|
||
ordered_match = 0
|
||
for i, gt in enumerate(gt_nums):
|
||
if i < len(pred_nums) and num_close(pred_nums[i], gt):
|
||
ordered_match += 1
|
||
strict_exact = 1.0 if len(pred_nums) == len(gt_nums) and ordered_match == len(gt_nums) else 0.0
|
||
order_ok = 1.0 if ordered_match == min(len(pred_nums), len(gt_nums)) and len(pred_nums) >= len(gt_nums) else 0.0
|
||
return {
|
||
"pred_nums": pred_nums,
|
||
"gt_nums": gt_nums,
|
||
"ordered_match_count": float(ordered_match),
|
||
"ordered_cov": float(ordered_match / max(len(gt_nums), 1)),
|
||
"strict_exact_match": float(strict_exact),
|
||
"order_ok": float(order_ok),
|
||
"extra_numbers": float(max(0, len(pred_nums) - len(gt_nums))),
|
||
"missing_numbers": float(max(0, len(gt_nums) - len(pred_nums))),
|
||
}
|
||
|
||
|
||
def gt_matches_text(text, gt_items):
|
||
answer_text = extract_answer_region(text)
|
||
pred_nums = numeric_candidates(answer_text)
|
||
used = [False] * len(pred_nums)
|
||
norm = normalize_text(answer_text)
|
||
matched = []
|
||
for gt in gt_items:
|
||
gt_s = normalize_text(gt)
|
||
if not gt_s:
|
||
continue
|
||
gt_f = _as_float_or_none(gt)
|
||
ok = False
|
||
if gt_f is not None:
|
||
for i, pred in enumerate(pred_nums):
|
||
if used[i]:
|
||
continue
|
||
if num_close(pred, gt_f):
|
||
used[i] = True
|
||
ok = True
|
||
break
|
||
else:
|
||
ok = gt_s in norm
|
||
if ok:
|
||
matched.append(gt)
|
||
return len(matched), matched
|
||
|
||
|
||
def rep_penalty(text):
|
||
toks = re.findall(r"\w+|[^\w\s]", str(text).lower())
|
||
if len(toks) < 16:
|
||
return 0.0
|
||
tri = [tuple(toks[i:i + 3]) for i in range(len(toks) - 2)]
|
||
repeat_rate = 1.0 - len(set(tri)) / max(len(tri), 1)
|
||
return min(0.8, max(0.0, repeat_rate - 0.25) * 2.0)
|
||
|
||
|
||
def flatten_json_scalars(obj):
|
||
out = []
|
||
if isinstance(obj, dict):
|
||
for v in obj.values():
|
||
out.extend(flatten_json_scalars(v))
|
||
elif isinstance(obj, list):
|
||
for v in obj:
|
||
out.extend(flatten_json_scalars(v))
|
||
elif obj is not None:
|
||
s = str(obj).strip()
|
||
if s:
|
||
out.append(s)
|
||
return out
|
||
|
||
|
||
def final_refs_tool_result(final_answer, tool_events):
|
||
answer_region = extract_answer_region(final_answer)
|
||
lowered = answer_region.lower()
|
||
nums = numeric_candidates(answer_region)
|
||
for event in tool_events:
|
||
if not event.get("ok"):
|
||
continue
|
||
for value in flatten_json_scalars(event.get("result", {})):
|
||
if value.lower() in lowered:
|
||
return True
|
||
v = value.replace(",", "")
|
||
if re.fullmatch(r"[-+]?\d+(?:\.\d+)?", v):
|
||
try:
|
||
target = float(v)
|
||
if any(num_close(n, target) for n in nums):
|
||
return True
|
||
except Exception:
|
||
pass
|
||
return False
|
||
|
||
|
||
def repeated_tool_call_penalty(tool_events):
|
||
keys = [(e.get("name", ""), json.dumps(e.get("args", {}), sort_keys=True, ensure_ascii=False)) for e in tool_events]
|
||
return min(0.8, 0.2 * (len(keys) - len(set(keys)))) if keys else 0.0
|
||
|
||
|
||
def math_task_like(task_text, tools=None):
|
||
return bool(extract_arithmetic_expressions(task_text, max_exprs=1))
|
||
|
||
|
||
def expected_math_call_count(task_text, gt_items):
|
||
expr_n = len(extract_arithmetic_expressions(task_text, max_exprs=16))
|
||
gt_n = len(gt_items or [])
|
||
return max(expr_n, gt_n, 1)
|
||
|
||
|
||
def tool_result_float(event: Dict[str, Any]) -> Optional[float]:
|
||
result = event.get("result", {}) or {}
|
||
for key in ("value", "result"):
|
||
if isinstance(result, dict) and key in result:
|
||
f = _as_float_or_none(result.get(key))
|
||
if f is not None:
|
||
return f
|
||
args = event.get("args", {}) or {}
|
||
expr = args.get("expression", args.get("query", "")) if isinstance(args, dict) else ""
|
||
if expr:
|
||
try:
|
||
return float(safe_eval_math(expr))
|
||
except Exception:
|
||
return None
|
||
return None
|
||
|
||
|
||
def math_tool_value_metrics(task_text: str, gt_items: List[Any], tool_events: List[Dict[str, Any]]):
|
||
exprs = extract_arithmetic_expressions(task_text, max_exprs=16)
|
||
expr_vals: List[float] = []
|
||
for expr in exprs:
|
||
try:
|
||
expr_vals.append(float(safe_eval_math(expr)))
|
||
except Exception:
|
||
pass
|
||
gt_nums = numeric_gt_items(gt_items)
|
||
# If GT length matches expression length, use GT for comparison to tolerate formatting/rounding policies.
|
||
expected_vals = gt_nums if gt_nums and len(gt_nums) == len(expr_vals) else expr_vals
|
||
expected_n = len(expected_vals)
|
||
math_events = [e for e in tool_events if e.get("name") in ("calculate_math", "calculator", "python_math")]
|
||
ok_events = [e for e in math_events if e.get("ok")]
|
||
call_vals = [tool_result_float(e) for e in ok_events]
|
||
call_vals = [v for v in call_vals if v is not None]
|
||
|
||
ordered_match = 0
|
||
for i, exp_v in enumerate(expected_vals):
|
||
if i < len(call_vals) and num_close(call_vals[i], exp_v):
|
||
ordered_match += 1
|
||
used = [False] * len(call_vals)
|
||
set_match = 0
|
||
for exp_v in expected_vals:
|
||
for j, got_v in enumerate(call_vals):
|
||
if used[j]:
|
||
continue
|
||
if num_close(got_v, exp_v):
|
||
used[j] = True
|
||
set_match += 1
|
||
break
|
||
over_calls = max(0, len(math_events) - expected_n) if expected_n else len(math_events)
|
||
under_calls = max(0, expected_n - len(math_events))
|
||
return {
|
||
"expected_exprs": exprs,
|
||
"expected_values": expected_vals,
|
||
"tool_values": call_vals,
|
||
"tool_value_order_match_count": float(ordered_match),
|
||
"tool_value_order_cov": float(ordered_match / max(expected_n, 1)) if expected_n else 0.0,
|
||
"tool_value_set_cov": float(set_match / max(expected_n, 1)) if expected_n else 0.0,
|
||
"tool_value_exact": float(expected_n > 0 and len(call_vals) == expected_n and ordered_match == expected_n),
|
||
"tool_over_calls": float(over_calls),
|
||
"tool_under_calls": float(under_calls),
|
||
}
|
||
|
||
|
||
def math_tool_expr_metrics(task_text: str, tool_events: List[Dict[str, Any]]):
|
||
expected = [canonicalize_expression(x) for x in extract_arithmetic_expressions(task_text, max_exprs=16)]
|
||
called = []
|
||
for event in tool_events:
|
||
if event.get("name") not in ("calculate_math", "calculator", "python_math"):
|
||
continue
|
||
args = event.get("args", {}) or {}
|
||
expr = args.get("expression", args.get("query", "")) if isinstance(args, dict) else ""
|
||
expr = canonicalize_expression(expr) if expr else ""
|
||
if expr:
|
||
called.append(expr)
|
||
|
||
ordered_match = 0
|
||
for i, exp in enumerate(expected):
|
||
if i < len(called) and called[i] == exp:
|
||
ordered_match += 1
|
||
used = [False] * len(called)
|
||
set_match = 0
|
||
for exp in expected:
|
||
for j, got in enumerate(called):
|
||
if used[j]:
|
||
continue
|
||
if got == exp:
|
||
used[j] = True
|
||
set_match += 1
|
||
break
|
||
expected_n = len(expected)
|
||
return {
|
||
"expected_expr_canon": expected,
|
||
"tool_exprs": called,
|
||
"tool_expr_order_match_count": float(ordered_match),
|
||
"tool_expr_order_cov": float(ordered_match / max(expected_n, 1)) if expected_n else 0.0,
|
||
"tool_expr_set_cov": float(set_match / max(expected_n, 1)) if expected_n else 0.0,
|
||
"tool_expr_exact": float(expected_n > 0 and len(called) == expected_n and ordered_match == expected_n),
|
||
}
|
||
|
||
|
||
def format_strength(args):
|
||
warmup = max(int(getattr(args, "format_warmup_steps", 0)), 0)
|
||
step = max(int(getattr(args, "current_step", 0)), 0)
|
||
if warmup <= 0:
|
||
return 1.0
|
||
return min(1.0, step / max(warmup, 1))
|
||
|
||
|
||
def scheduled_reward_cap(args, initial_attr: str, final_attr: str) -> float:
|
||
final_value = float(getattr(args, final_attr, 0.0))
|
||
initial_value = float(getattr(args, initial_attr, final_value))
|
||
warmup = max(int(getattr(args, "reward_cap_warmup_steps", 0)), 0)
|
||
step = max(int(getattr(args, "current_step", 0)), 0)
|
||
if warmup <= 0:
|
||
return final_value
|
||
strength = min(1.0, step / max(warmup, 1))
|
||
return initial_value + strength * (final_value - initial_value)
|
||
|
||
|
||
def score_one_trajectory(args, prompt, completion, gt_items, tools, tool_events, unfinished, reward_model=None, messages=None):
|
||
task_text = task_text_from_messages(messages) if messages else task_text_from_prompt(prompt)
|
||
tag = answer_tag_stats(completion)
|
||
final_answer = remove_tool_blocks(completion)
|
||
answer_region = tag["answer_content"]
|
||
matched_count, _ = gt_matches_text(completion, gt_items)
|
||
gt_total = max(len(gt_items), 1)
|
||
unordered_cov = matched_count / gt_total if gt_items else 0.0
|
||
|
||
nmet = ordered_numeric_metrics(answer_region, gt_items)
|
||
has_numeric_gt = len(nmet["gt_nums"]) > 0
|
||
answer_cov = nmet["ordered_cov"] if has_numeric_gt and bool(getattr(args, "strict_order_reward", 1)) else unordered_cov
|
||
exact_match = nmet["strict_exact_match"] if has_numeric_gt and bool(getattr(args, "strict_exact_reward", 1)) else (1.0 if gt_items and unordered_cov >= 1.0 else 0.0)
|
||
|
||
total_calls = len(tool_events)
|
||
success_calls = sum(1 for x in tool_events if x.get("ok"))
|
||
failed_calls = total_calls - success_calls
|
||
tool_success_rate = success_calls / max(total_calls, 1) if total_calls else 0.0
|
||
refs_tool = 1.0 if final_refs_tool_result(completion, tool_events) else 0.0
|
||
answer_tag = tag["has_answer_tag"]
|
||
exact_answer_tag = tag["exactly_one_answer_tag"]
|
||
malformed_tags = float(abs(str(completion).count("<tool_call>") - str(completion).count("</tool_call>"))) + tag["malformed_answer_tags"]
|
||
is_math = math_task_like(task_text, tools)
|
||
expected_math_calls = expected_math_call_count(task_text, gt_items) if is_math else 0
|
||
|
||
preferred_math_tool_calls = sum(1 for x in tool_events if x.get("name") == "calculate_math")
|
||
alias_math_tool_calls = sum(1 for x in tool_events if x.get("name") in ("calculator", "python_math"))
|
||
math_tool_calls = preferred_math_tool_calls + alias_math_tool_calls
|
||
non_math_tool_calls = max(0, total_calls - math_tool_calls) if is_math else 0
|
||
successful_math_calls = sum(1 for x in tool_events if x.get("ok") and x.get("name") in ("calculate_math", "calculator", "python_math"))
|
||
successful_preferred_math_calls = sum(1 for x in tool_events if x.get("ok") and x.get("name") == "calculate_math")
|
||
math_call_coverage = min(1.0, successful_math_calls / max(expected_math_calls, 1)) if is_math else 0.0
|
||
preferred_math_rate = successful_preferred_math_calls / max(successful_math_calls, 1) if is_math and successful_math_calls else 0.0
|
||
correct_tool_choice = 1.0 if (is_math and total_calls > 0 and non_math_tool_calls == 0 and preferred_math_tool_calls > 0) else (1.0 if not is_math else 0.0)
|
||
tool_count_exact = 1.0 if (is_math and math_tool_calls == expected_math_calls and non_math_tool_calls == 0) else (1.0 if not is_math else 0.0)
|
||
canonical_tool_rate = safe_mean([1.0 if e.get("canonical", False) else 0.0 for e in tool_events], 0.0) if tool_events else 0.0
|
||
tvm = math_tool_value_metrics(task_text, gt_items, tool_events) if is_math else {
|
||
"tool_value_order_cov": 0.0,
|
||
"tool_value_set_cov": 0.0,
|
||
"tool_value_exact": 0.0,
|
||
"tool_over_calls": 0.0,
|
||
"tool_under_calls": 0.0,
|
||
"tool_value_order_match_count": 0.0,
|
||
}
|
||
tem = math_tool_expr_metrics(task_text, tool_events) if is_math else {
|
||
"tool_expr_order_cov": 0.0,
|
||
"tool_expr_set_cov": 0.0,
|
||
"tool_expr_exact": 0.0,
|
||
"tool_expr_order_match_count": 0.0,
|
||
}
|
||
answer_len = len(answer_region)
|
||
rep = rep_penalty(final_answer)
|
||
repeat_tool_pen = repeated_tool_call_penalty(tool_events)
|
||
|
||
reward = float(getattr(args, "reward_base", -0.10))
|
||
reward += float(getattr(args, "reward_answer_cov_weight", 8.0)) * answer_cov
|
||
reward += float(getattr(args, "reward_exact_match_weight", 8.0)) * exact_match
|
||
|
||
if is_math:
|
||
if total_calls > 0:
|
||
reward += float(getattr(args, "reward_tool_success_weight", 0.60)) * tool_success_rate
|
||
reward += float(getattr(args, "reward_math_call_cov_weight", 0.60)) * math_call_coverage
|
||
reward += float(getattr(args, "reward_tool_expr_order_weight", 1.40)) * float(tem["tool_expr_order_cov"])
|
||
reward += float(getattr(args, "reward_tool_expr_set_weight", 0.50)) * float(tem["tool_expr_set_cov"])
|
||
reward += float(getattr(args, "reward_tool_expr_exact_weight", 0.60)) * float(tem["tool_expr_exact"])
|
||
reward += float(getattr(args, "reward_tool_value_order_weight", 2.40)) * float(tvm["tool_value_order_cov"])
|
||
reward += float(getattr(args, "reward_tool_value_set_weight", 0.80)) * float(tvm["tool_value_set_cov"])
|
||
reward += float(getattr(args, "reward_tool_value_exact_weight", 0.80)) * float(tvm["tool_value_exact"])
|
||
reward += float(getattr(args, "reward_preferred_math_tool_weight", 0.70)) * preferred_math_rate
|
||
reward += float(getattr(args, "reward_correct_tool_choice_weight", 0.70)) * correct_tool_choice
|
||
reward += float(getattr(args, "reward_tool_count_exact_weight", 0.30)) * tool_count_exact
|
||
reward += float(getattr(args, "reward_canonical_tool_weight", 0.25)) * canonical_tool_rate
|
||
if answer_cov > 0.0:
|
||
reward += float(getattr(args, "reward_final_refs_tool_weight", 0.45)) * refs_tool
|
||
reward -= float(getattr(args, "reward_tool_fail_penalty", 0.55)) * min(failed_calls, 4)
|
||
reward -= float(getattr(args, "reward_wrong_tool_penalty", 0.55)) * min(non_math_tool_calls, 4)
|
||
reward -= float(getattr(args, "reward_tool_over_call_penalty", 0.25)) * min(float(tvm["tool_over_calls"]), 6.0)
|
||
reward -= float(getattr(args, "reward_tool_under_call_penalty", 0.35)) * min(float(tvm["tool_under_calls"]), 6.0)
|
||
else:
|
||
reward -= float(getattr(args, "reward_no_tool_penalty", 1.50))
|
||
if alias_math_tool_calls > 0 and preferred_math_tool_calls <= 0:
|
||
reward -= float(getattr(args, "reward_alias_only_penalty", 0.35))
|
||
|
||
strength = format_strength(args)
|
||
missing_format_penalty = float(getattr(args, "format_missing_penalty_initial", 0.30)) + strength * (
|
||
float(getattr(args, "format_missing_penalty_final", 1.50)) - float(getattr(args, "format_missing_penalty_initial", 0.30))
|
||
)
|
||
reward += float(getattr(args, "reward_answer_tag_weight", 0.90)) * answer_tag
|
||
reward += float(getattr(args, "reward_exact_answer_tag_weight", 0.80)) * exact_answer_tag
|
||
if answer_tag <= 0.0 and is_math:
|
||
reward -= missing_format_penalty
|
||
if getattr(args, "format_enforce_mode", "semi_hard") in ("semi_hard", "hard"):
|
||
reward = min(reward, float(getattr(args, "no_answer_tag_reward_cap", 2.0)))
|
||
if getattr(args, "format_enforce_mode", "semi_hard") == "hard" and strength >= 1.0:
|
||
answer_cov = 0.0
|
||
exact_match = 0.0
|
||
reward = min(reward, 0.0)
|
||
reward += 0.10 if malformed_tags == 0 else -0.50 * malformed_tags
|
||
reward += 0.10 if 1 <= answer_len <= int(getattr(args, "max_reward_answer_len", 96)) else -0.30
|
||
|
||
extra_numbers = nmet["extra_numbers"] if has_numeric_gt else 0.0
|
||
missing_numbers = nmet["missing_numbers"] if has_numeric_gt else 0.0
|
||
order_ok = nmet["order_ok"] if has_numeric_gt else exact_match
|
||
if has_numeric_gt:
|
||
reward -= float(getattr(args, "extra_number_penalty", 0.55)) * min(extra_numbers, 6.0)
|
||
reward -= float(getattr(args, "missing_number_penalty", 0.70)) * min(missing_numbers, 6.0)
|
||
if order_ok <= 0.0 and unordered_cov > 0.0:
|
||
reward -= float(getattr(args, "order_mismatch_penalty", 1.20))
|
||
|
||
reward -= repeat_tool_pen
|
||
reward -= 0.80 if unfinished else 0.0
|
||
reward -= rep
|
||
|
||
if gt_items and exact_match < 1.0:
|
||
reward = min(reward, scheduled_reward_cap(args, "non_exact_reward_cap_initial", "non_exact_reward_cap"))
|
||
if gt_items and answer_cov <= 0.0:
|
||
reward = min(reward, scheduled_reward_cap(args, "zero_answer_cov_reward_cap_initial", "zero_answer_cov_reward_cap"))
|
||
if reward_model is not None:
|
||
try:
|
||
pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
|
||
messages = [{"role": role, "content": content.strip()} for role, content in re.findall(pattern, prompt, re.DOTALL)]
|
||
reward += max(min(float(reward_model.get_score(messages, final_answer)), 1.0), -1.0)
|
||
except Exception:
|
||
pass
|
||
raw_reward = float(reward)
|
||
reward = max(min(raw_reward, 12.0), -12.0)
|
||
|
||
return reward, {
|
||
"reward": reward,
|
||
"raw_reward": raw_reward,
|
||
"clipped": float(abs(raw_reward - reward) > 1e-6),
|
||
"answer_cov": float(answer_cov),
|
||
"unordered_answer_cov": float(unordered_cov),
|
||
"exact_match": float(exact_match),
|
||
"strict_exact_match": float(nmet["strict_exact_match"] if has_numeric_gt else exact_match),
|
||
"matched_count": float(matched_count),
|
||
"ordered_match_count": float(nmet["ordered_match_count"]),
|
||
"extra_numbers": float(extra_numbers),
|
||
"missing_numbers": float(missing_numbers),
|
||
"order_ok": float(order_ok),
|
||
"tool_calls": float(total_calls),
|
||
"expected_math_calls": float(expected_math_calls),
|
||
"math_tool_calls": float(math_tool_calls),
|
||
"preferred_math_tool_calls": float(preferred_math_tool_calls),
|
||
"alias_math_tool_calls": float(alias_math_tool_calls),
|
||
"non_math_tool_calls": float(non_math_tool_calls),
|
||
"tool_success_calls": float(success_calls),
|
||
"tool_failed_calls": float(failed_calls),
|
||
"tool_success_rate": float(tool_success_rate),
|
||
"math_call_coverage": float(math_call_coverage),
|
||
"tool_expr_order_cov": float(tem["tool_expr_order_cov"]),
|
||
"tool_expr_set_cov": float(tem["tool_expr_set_cov"]),
|
||
"tool_expr_exact": float(tem["tool_expr_exact"]),
|
||
"tool_expr_order_match_count": float(tem["tool_expr_order_match_count"]),
|
||
"tool_value_order_cov": float(tvm["tool_value_order_cov"]),
|
||
"tool_value_set_cov": float(tvm["tool_value_set_cov"]),
|
||
"tool_value_exact": float(tvm["tool_value_exact"]),
|
||
"tool_value_order_match_count": float(tvm["tool_value_order_match_count"]),
|
||
"tool_over_calls": float(tvm["tool_over_calls"]),
|
||
"tool_under_calls": float(tvm["tool_under_calls"]),
|
||
"preferred_math_rate": float(preferred_math_rate),
|
||
"correct_tool_choice": float(correct_tool_choice),
|
||
"tool_count_exact": float(tool_count_exact),
|
||
"canonical_tool_rate": float(canonical_tool_rate),
|
||
"final_refs_tool": float(refs_tool),
|
||
"answer_tag": float(answer_tag),
|
||
"exact_answer_tag": float(exact_answer_tag),
|
||
"answer_tag_nonempty": float(tag["answer_tag_nonempty"]),
|
||
"malformed_tags": float(malformed_tags),
|
||
"answer_len": float(answer_len),
|
||
"unfinished": float(bool(unfinished)),
|
||
"rep_penalty": float(rep),
|
||
"repeat_tool_penalty": float(repeat_tool_pen),
|
||
"format_strength": float(strength),
|
||
}
|
||
|
||
|
||
def calculate_rewards(args, prompts, completions, gt_batch, tools_batch, tool_events_batch, unfinished_batch, num_generations, device, reward_model=None, return_details=False, messages_batch=None):
|
||
rewards_list, details = [], []
|
||
for idx, completion in enumerate(completions):
|
||
sample_idx = idx // max(num_generations, 1)
|
||
reward, detail = score_one_trajectory(
|
||
args,
|
||
prompts[sample_idx],
|
||
completion,
|
||
gt_batch[sample_idx],
|
||
tools_batch[sample_idx],
|
||
tool_events_batch[idx],
|
||
unfinished_batch[idx],
|
||
reward_model,
|
||
messages_batch[sample_idx] if messages_batch is not None else None,
|
||
)
|
||
rewards_list.append(reward)
|
||
details.append(detail)
|
||
rewards = torch.tensor(rewards_list, dtype=torch.float32, device=device)
|
||
return (rewards, details) if return_details else rewards
|
||
|
||
|
||
# -------------------------- rollout data structures --------------------------
|
||
|
||
@dataclass
|
||
class Trajectory:
|
||
prompt: str
|
||
messages: List[Dict[str, Any]]
|
||
tools: Optional[List[Dict[str, Any]]]
|
||
gt: List[Any]
|
||
completion: str
|
||
context: str
|
||
prompt_ids: List[int]
|
||
response_ids: List[int]
|
||
response_mask: List[int]
|
||
response_old_logps: List[float]
|
||
response_versions: List[int]
|
||
turn_outputs: List[str]
|
||
tool_events: List[Dict[str, Any]]
|
||
unfinished: bool
|
||
policy_version: int
|
||
data_index: int
|
||
rollout_time: float
|
||
|
||
|
||
@dataclass
|
||
class RolloutGroup:
|
||
prompt: str
|
||
data_index: int
|
||
policy_version: int
|
||
trajectories: List[Trajectory]
|
||
created_at: float
|
||
|
||
|
||
class PolicyVersion:
|
||
def __init__(self, initial=0):
|
||
self._version = int(initial)
|
||
self._lock = threading.Lock()
|
||
|
||
def get(self):
|
||
with self._lock:
|
||
return self._version
|
||
|
||
def bump(self):
|
||
with self._lock:
|
||
self._version += 1
|
||
return self._version
|
||
|
||
def set(self, value):
|
||
with self._lock:
|
||
self._version = int(value)
|
||
|
||
|
||
def execute_tool_call(env, valid_names, tool_call):
|
||
name = tool_call.get("name", "") or ""
|
||
raw_args = tool_call.get("arguments", tool_call.get("args", {}))
|
||
if isinstance(raw_args, str):
|
||
try:
|
||
raw_args = json.loads(raw_args)
|
||
except Exception:
|
||
raw_args = {"expression": raw_args}
|
||
if not isinstance(raw_args, dict):
|
||
raw_args = {}
|
||
if not name:
|
||
return {"ok": False, "name": name, "args": raw_args, "result": {"error": "bad tool call"}, "canonical": bool(tool_call.get("canonical", False))}
|
||
if valid_names and name not in valid_names:
|
||
return {
|
||
"ok": False,
|
||
"name": name,
|
||
"args": raw_args,
|
||
"result": {"error": f"tool not allowed: {name}"},
|
||
"canonical": bool(tool_call.get("canonical", False)),
|
||
}
|
||
try:
|
||
return {"ok": True, "name": name, "args": raw_args, "result": env.execute(name, raw_args), "canonical": bool(tool_call.get("canonical", False))}
|
||
except Exception as exc:
|
||
return {"ok": False, "name": name, "args": raw_args, "result": {"error": str(exc)[:200]}, "canonical": bool(tool_call.get("canonical", False))}
|
||
|
||
|
||
def suffix_delta(old_ids, new_ids):
|
||
if len(new_ids) >= len(old_ids) and new_ids[:len(old_ids)] == old_ids:
|
||
return new_ids[len(old_ids):]
|
||
common = 0
|
||
max_common = min(len(old_ids), len(new_ids))
|
||
while common < max_common and old_ids[common] == new_ids[common]:
|
||
common += 1
|
||
return new_ids[common:]
|
||
|
||
|
||
def rollout_single_online(
|
||
rollout_engine,
|
||
tokenizer,
|
||
env,
|
||
sample,
|
||
max_turns,
|
||
max_new_tokens,
|
||
max_context_len,
|
||
thinking_ratio,
|
||
device,
|
||
policy_version,
|
||
temperature,
|
||
):
|
||
messages = [dict(m) for m in sample["messages"]]
|
||
tools = sample["tools"] or env.tool_specs()
|
||
valid_names = set(supported_tool_names(tools)) or env.supported_names()
|
||
prompt = apply_chat_template_compat(tokenizer, messages, tools=tools, add_generation_prompt=True)
|
||
all_outputs, tool_events = [], []
|
||
prompt_ids = None
|
||
response_ids, response_mask, response_old_logps, response_versions = [], [], [], []
|
||
|
||
# Full untruncated rendered context used only for observation-delta tracking.
|
||
# The packed training sequence still uses prompt_ids + response_ids as before.
|
||
current_ids_full = None
|
||
final_context = prompt
|
||
unfinished = False
|
||
open_thinking = random.random() < thinking_ratio
|
||
rollout_start = time.time()
|
||
|
||
for turn in range(max_turns):
|
||
context = apply_chat_template_compat(tokenizer, messages, tools=tools, add_generation_prompt=True, open_thinking=open_thinking)
|
||
inputs = tokenize_context(tokenizer, context, device=device, max_context_len=max_context_len)
|
||
context_ids = inputs["input_ids"][0].tolist()
|
||
context_ids_full = tokenize_context(tokenizer, context, device="cpu", max_context_len=0)["input_ids"][0].tolist()
|
||
if prompt_ids is None:
|
||
prompt_ids = context_ids
|
||
# Source of truth for full rendered dialogue before the current assistant generation.
|
||
if current_ids_full is None:
|
||
current_ids_full = context_ids_full
|
||
elif current_ids_full != context_ids_full:
|
||
# This should rarely happen. If the chat template re-renders differently,
|
||
# use the freshly rendered context to avoid a broken suffix_delta.
|
||
current_ids_full = context_ids_full
|
||
rr = rollout_engine.rollout(
|
||
inputs["input_ids"],
|
||
inputs["attention_mask"],
|
||
num_generations=1,
|
||
max_new_tokens=max_new_tokens,
|
||
temperature=temperature,
|
||
)
|
||
gen_ids_raw = rr.completion_ids[0].tolist()
|
||
gen_logps_raw = rr.per_token_logps[0].tolist()
|
||
pairs = [(int(tok), finite_float(logp)) for tok, logp in zip(gen_ids_raw, gen_logps_raw) if tok != tokenizer.pad_token_id and tok != tokenizer.eos_token_id]
|
||
gen_ids = [x for x, _ in pairs]
|
||
gen_logps = [x for _, x in pairs]
|
||
gen_text = rr.completions[0]
|
||
all_outputs.append(gen_text)
|
||
response_ids.extend(gen_ids)
|
||
response_mask.extend([1] * len(gen_ids))
|
||
response_old_logps.extend(gen_logps)
|
||
response_versions.extend([policy_version] * len(gen_ids))
|
||
assert_rollout_alignment(response_ids, response_mask, response_old_logps, response_versions)
|
||
|
||
# Critical fix:
|
||
# current_ids_full must include the assistant generation before computing
|
||
# the observation delta. Otherwise gen_ids will be duplicated as mask=0 obs.
|
||
current_ids_after_gen_full = list(current_ids_full) + list(gen_ids)
|
||
|
||
final_context = context + gen_text
|
||
calls = parse_tool_calls(gen_text)
|
||
if not calls:
|
||
break
|
||
|
||
unfinished = turn == max_turns - 1
|
||
messages.append({"role": "assistant", "content": gen_text})
|
||
for tool_call in calls:
|
||
event = execute_tool_call(env, valid_names, tool_call)
|
||
tool_events.append(event)
|
||
messages.append({"role": "tool", "content": json.dumps(event["result"], ensure_ascii=False, separators=(",", ":"))[:1024]})
|
||
observe_context = apply_chat_template_compat(tokenizer, messages, tools=tools, add_generation_prompt=not unfinished, open_thinking=open_thinking)
|
||
observe_ids_full = tokenize_context(tokenizer, observe_context, device="cpu", max_context_len=0)["input_ids"][0].tolist()
|
||
|
||
obs_delta = suffix_delta(current_ids_after_gen_full, observe_ids_full)
|
||
|
||
# Extra safety: if suffix tracking ever fails and the generated text appears
|
||
# again at the front of obs_delta, strip it instead of corrupting PPO context.
|
||
if gen_ids and len(obs_delta) >= len(gen_ids) and obs_delta[:len(gen_ids)] == gen_ids:
|
||
obs_delta = obs_delta[len(gen_ids):]
|
||
|
||
current_ids_full = observe_ids_full
|
||
if obs_delta:
|
||
response_ids.extend(obs_delta)
|
||
response_mask.extend([0] * len(obs_delta))
|
||
response_old_logps.extend([0.0] * len(obs_delta))
|
||
response_versions.extend([-1] * len(obs_delta))
|
||
assert_rollout_alignment(response_ids, response_mask, response_old_logps, response_versions)
|
||
final_context = observe_context
|
||
if unfinished:
|
||
break
|
||
|
||
# Score the whole assistant trajectory. v3 only scored all_outputs[-1], which can lose
|
||
# an answer or format signal in multi-turn/tool-call trajectories.
|
||
completion = "\n".join(x for x in all_outputs if x is not None)
|
||
assert_rollout_alignment(response_ids, response_mask, response_old_logps, response_versions)
|
||
return Trajectory(
|
||
prompt,
|
||
sample["messages"],
|
||
tools,
|
||
sample["gt"],
|
||
completion,
|
||
final_context,
|
||
prompt_ids or [],
|
||
response_ids,
|
||
response_mask,
|
||
response_old_logps,
|
||
response_versions,
|
||
list(all_outputs),
|
||
tool_events,
|
||
unfinished,
|
||
policy_version,
|
||
-1,
|
||
time.time() - rollout_start,
|
||
)
|
||
|
||
|
||
def rollout_group_online(
|
||
rollout_engine,
|
||
tokenizer,
|
||
env,
|
||
sample,
|
||
data_index,
|
||
num_generations,
|
||
max_turns,
|
||
max_new_tokens,
|
||
max_context_len,
|
||
thinking_ratio,
|
||
device,
|
||
policy_version,
|
||
temperature,
|
||
):
|
||
trajectories = []
|
||
for _ in range(num_generations):
|
||
traj = rollout_single_online(
|
||
rollout_engine,
|
||
tokenizer,
|
||
env,
|
||
sample,
|
||
max_turns,
|
||
max_new_tokens,
|
||
max_context_len,
|
||
thinking_ratio,
|
||
device,
|
||
policy_version,
|
||
temperature,
|
||
)
|
||
traj.data_index = data_index
|
||
trajectories.append(traj)
|
||
return RolloutGroup(trajectories[0].prompt if trajectories else "", data_index, policy_version, trajectories, time.time())
|
||
|
||
|
||
class AsyncRolloutManager:
|
||
def __init__(
|
||
self,
|
||
dataset,
|
||
tokenizer,
|
||
rollout_engine,
|
||
env,
|
||
version,
|
||
batch_size,
|
||
num_generations,
|
||
max_turns,
|
||
max_new_tokens,
|
||
max_context_len,
|
||
thinking_ratio,
|
||
temperature,
|
||
device,
|
||
max_version_gap,
|
||
max_queue_groups,
|
||
num_workers,
|
||
engine_lock,
|
||
seed,
|
||
):
|
||
self.dataset, self.tokenizer, self.rollout_engine, self.env, self.version = dataset, tokenizer, rollout_engine, env, version
|
||
self.batch_size, self.num_generations, self.max_turns, self.max_new_tokens = int(batch_size), int(num_generations), int(max_turns), int(max_new_tokens)
|
||
self.max_context_len, self.thinking_ratio, self.temperature, self.device = int(max_context_len), float(thinking_ratio), float(temperature), device
|
||
self.max_version_gap = max(0, int(max_version_gap))
|
||
self.max_pending_groups = max(1, int(max_queue_groups))
|
||
self.num_workers = max(0, int(num_workers))
|
||
self.engine_lock, self.seed = engine_lock, int(seed)
|
||
self.buffer = queue.Queue(maxsize=self.max_pending_groups)
|
||
self.stop_event = threading.Event()
|
||
self.threads = []
|
||
self.index_lock = threading.Lock()
|
||
self.next_index = 0
|
||
self.stats_lock = threading.Lock()
|
||
self.rollout_time_hist = deque(maxlen=512)
|
||
self.train_wait_hist = deque(maxlen=512)
|
||
self.rollout_lock_wait_hist = deque(maxlen=512)
|
||
self.queue_hist = deque(maxlen=512)
|
||
self.stats = {
|
||
"produced": 0.0,
|
||
"accepted": 0.0,
|
||
"generated": 0.0,
|
||
"stale_drop": 0.0,
|
||
"errors": 0.0,
|
||
"queue_full": 0.0,
|
||
"rollout_time_sum": 0.0,
|
||
"rollout_lock_wait_sum": 0.0,
|
||
"train_wait_sum": 0.0,
|
||
"train_batches": 0.0,
|
||
"rollout_turns": 0.0,
|
||
"rollout_tool_calls": 0.0,
|
||
"rollout_tokens": 0.0,
|
||
"last_train_wait": 0.0,
|
||
"last_rollout_time": 0.0,
|
||
"last_rollout_lock_wait": 0.0,
|
||
}
|
||
|
||
def _next_sample(self, rng):
|
||
with self.index_lock:
|
||
if rng.random() < 0.50:
|
||
idx = self.next_index % len(self.dataset)
|
||
self.next_index += 1
|
||
else:
|
||
idx = rng.randrange(len(self.dataset))
|
||
return idx, self.dataset[idx]
|
||
|
||
def _inc(self, key, val=1.0):
|
||
with self.stats_lock:
|
||
self.stats[key] = self.stats.get(key, 0.0) + float(val)
|
||
|
||
def _set(self, key, val):
|
||
with self.stats_lock:
|
||
self.stats[key] = float(val)
|
||
|
||
def _push_hist(self, hist: deque, value: float):
|
||
with self.stats_lock:
|
||
hist.append(float(value))
|
||
|
||
def _producer(self, worker_id):
|
||
rank = dist.get_rank() if dist.is_initialized() else 0
|
||
rng = random.Random(self.seed + 1009 * (worker_id + 1) + 9173 * rank)
|
||
while not self.stop_event.is_set():
|
||
try:
|
||
if self.buffer.qsize() >= self.max_pending_groups:
|
||
self._inc("queue_full")
|
||
time.sleep(0.05)
|
||
continue
|
||
data_index, sample = self._next_sample(rng)
|
||
policy_version = self.version.get()
|
||
lock_wait_start = time.time()
|
||
with self.engine_lock:
|
||
lock_wait = time.time() - lock_wait_start
|
||
t0 = time.time()
|
||
group = rollout_group_online(
|
||
self.rollout_engine,
|
||
self.tokenizer,
|
||
self.env,
|
||
sample,
|
||
data_index,
|
||
self.num_generations,
|
||
self.max_turns,
|
||
self.max_new_tokens,
|
||
self.max_context_len,
|
||
self.thinking_ratio,
|
||
self.device,
|
||
policy_version,
|
||
self.temperature,
|
||
)
|
||
rollout_sec = time.time() - t0
|
||
self._inc("produced")
|
||
self._inc("generated", len(group.trajectories))
|
||
self._inc("rollout_time_sum", rollout_sec)
|
||
self._inc("rollout_lock_wait_sum", lock_wait)
|
||
self._inc("rollout_turns", sum(len(t.turn_outputs) for t in group.trajectories))
|
||
self._inc("rollout_tool_calls", sum(len(t.tool_events) for t in group.trajectories))
|
||
self._inc("rollout_tokens", sum(sum(t.response_mask) for t in group.trajectories))
|
||
self._set("last_rollout_time", rollout_sec)
|
||
self._set("last_rollout_lock_wait", lock_wait)
|
||
self._push_hist(self.rollout_time_hist, rollout_sec)
|
||
self._push_hist(self.rollout_lock_wait_hist, lock_wait)
|
||
self._push_hist(self.queue_hist, self.buffer.qsize())
|
||
if self.version.get() - group.policy_version > self.max_version_gap:
|
||
self._inc("stale_drop")
|
||
continue
|
||
self.buffer.put(group, timeout=1.0)
|
||
self._inc("accepted")
|
||
except queue.Full:
|
||
self._inc("queue_full")
|
||
except Exception as exc:
|
||
self._inc("errors")
|
||
if is_main_process():
|
||
Logger(f"[WARN] rollout worker {worker_id} error: {repr(exc)}")
|
||
time.sleep(0.2)
|
||
|
||
def start(self):
|
||
for worker_id in range(self.num_workers):
|
||
thread = threading.Thread(target=self._producer, args=(worker_id,), daemon=False)
|
||
thread.start()
|
||
self.threads.append(thread)
|
||
|
||
def get_batch(self, timeout=None):
|
||
groups = []
|
||
start = time.time()
|
||
while len(groups) < self.batch_size and not self.stop_event.is_set():
|
||
try:
|
||
group = self.buffer.get(timeout=timeout if timeout is not None else 60.0)
|
||
except queue.Empty:
|
||
break
|
||
if self.version.get() - group.policy_version > self.max_version_gap:
|
||
self._inc("stale_drop")
|
||
continue
|
||
groups.append(group)
|
||
wait_sec = time.time() - start
|
||
self._set("last_train_wait", wait_sec)
|
||
self._inc("train_wait_sum", wait_sec)
|
||
self._inc("train_batches")
|
||
self._push_hist(self.train_wait_hist, wait_sec)
|
||
self._push_hist(self.queue_hist, self.buffer.qsize())
|
||
return groups
|
||
|
||
def get_one(self, timeout=None):
|
||
start = time.time()
|
||
while not self.stop_event.is_set():
|
||
try:
|
||
group = self.buffer.get(timeout=timeout if timeout is not None else 60.0)
|
||
except queue.Empty:
|
||
return None
|
||
if self.version.get() - group.policy_version > self.max_version_gap:
|
||
self._inc("stale_drop")
|
||
continue
|
||
wait_sec = time.time() - start
|
||
self._set("last_train_wait", wait_sec)
|
||
self._inc("train_wait_sum", wait_sec)
|
||
self._inc("train_batches")
|
||
self._push_hist(self.train_wait_hist, wait_sec)
|
||
self._push_hist(self.queue_hist, self.buffer.qsize())
|
||
return group
|
||
return None
|
||
|
||
def wait_until_ready(self, min_groups: int, timeout_sec: float = 180.0):
|
||
target = min(max(int(min_groups), 0), self.max_pending_groups)
|
||
if target <= 0:
|
||
return True
|
||
start = time.time()
|
||
while not self.stop_event.is_set():
|
||
if self.buffer.qsize() >= target:
|
||
return True
|
||
if (time.time() - start) >= max(float(timeout_sec), 0.0):
|
||
return False
|
||
time.sleep(0.05)
|
||
return False
|
||
|
||
def snapshot_stats(self):
|
||
with self.stats_lock:
|
||
out = dict(self.stats)
|
||
rollout_hist = list(self.rollout_time_hist)
|
||
train_hist = list(self.train_wait_hist)
|
||
lock_hist = list(self.rollout_lock_wait_hist)
|
||
queue_hist = list(self.queue_hist)
|
||
produced = max(out.get("produced", 0.0), 1.0)
|
||
generated = max(out.get("generated", 0.0), 1.0)
|
||
train_batches = max(out.get("train_batches", 0.0), 1.0)
|
||
out["queued"] = float(self.buffer.qsize())
|
||
out["queue_capacity"] = float(self.max_pending_groups)
|
||
out["queue_utilization"] = out["queued"] / max(out["queue_capacity"], 1.0)
|
||
out["online_group_accept_rate"] = out.get("accepted", 0.0) / produced
|
||
out["stale_drop_rate"] = out.get("stale_drop", 0.0) / produced
|
||
out["avg_rollout_sec"] = out.get("rollout_time_sum", 0.0) / produced
|
||
out["avg_train_wait_sec"] = out.get("train_wait_sum", 0.0) / train_batches
|
||
out["avg_tokens_per_traj"] = out.get("rollout_tokens", 0.0) / generated
|
||
out["avg_turns_per_traj"] = out.get("rollout_turns", 0.0) / generated
|
||
out["avg_tool_calls_per_traj"] = out.get("rollout_tool_calls", 0.0) / generated
|
||
out["rollout_sec_p50"] = percentile(rollout_hist, 0.50)
|
||
out["rollout_sec_p90"] = percentile(rollout_hist, 0.90)
|
||
out["rollout_sec_p99"] = percentile(rollout_hist, 0.99)
|
||
out["train_wait_sec_p50"] = percentile(train_hist, 0.50)
|
||
out["train_wait_sec_p90"] = percentile(train_hist, 0.90)
|
||
out["train_wait_sec_p99"] = percentile(train_hist, 0.99)
|
||
out["rollout_lock_wait_sec_p90"] = percentile(lock_hist, 0.90)
|
||
out["queue_size_p50"] = percentile(queue_hist, 0.50)
|
||
out["queue_size_p90"] = percentile(queue_hist, 0.90)
|
||
return out
|
||
|
||
def stop(self):
|
||
self.stop_event.set()
|
||
for thread in self.threads:
|
||
thread.join(timeout=30.0)
|
||
|
||
|
||
def sync_rollout_policy(args, rollout_engine, rollout_model, model, engine_lock, version):
|
||
with engine_lock:
|
||
if args.rollout_engine == "torch":
|
||
ok = copy_policy_weights(rollout_model, model)
|
||
if ok:
|
||
rollout_engine.update_policy(rollout_model)
|
||
else:
|
||
ok = bool(rollout_engine.update_policy(model))
|
||
return version.bump() if ok else version.get()
|
||
|
||
|
||
def manager_pop_one_group(manager, timeout=None):
|
||
return manager.get_one(timeout=timeout)
|
||
|
||
|
||
def sync_rollout_policy_distributed(args, rollout_engine, rollout_model, model, engine_lock, version, src_rank=0):
|
||
raw_model = unwrap_model(model)
|
||
state_tensors = [tensor for tensor in raw_model.state_dict().values() if torch.is_tensor(tensor)]
|
||
ok = True
|
||
if dist.get_rank() == src_rank and model_has_nonfinite(model):
|
||
ok = False
|
||
ok_tensor = torch.tensor([1 if ok else 0], device=args.device, dtype=torch.int64)
|
||
dist.broadcast(ok_tensor, src=src_rank)
|
||
ok = bool(ok_tensor.item() > 0)
|
||
if ok:
|
||
for tensor in state_tensors:
|
||
dist.broadcast(tensor.data, src=src_rank)
|
||
with engine_lock:
|
||
if args.rollout_engine == "torch":
|
||
if rollout_model is not raw_model:
|
||
ok = ok and copy_policy_weights(rollout_model, raw_model)
|
||
if ok:
|
||
rollout_engine.update_policy(rollout_model)
|
||
else:
|
||
ok = ok and bool(rollout_engine.update_policy(raw_model))
|
||
version_value = version.bump() if (dist.get_rank() == src_rank and ok) else version.get()
|
||
version_tensor = torch.tensor([int(version_value)], device=args.device, dtype=torch.int64)
|
||
dist.broadcast(version_tensor, src=src_rank)
|
||
version.set(int(version_tensor.item()))
|
||
return version.get()
|
||
|
||
|
||
def distributed_fetch_groups(rank, world_size, manager, pending_groups, target_groups, fetch_timeout_sec=180.0, producer_ranks=None, src_rank=0, min_return_groups=None):
|
||
start = time.time()
|
||
target_groups = max(0, int(target_groups))
|
||
if min_return_groups is None:
|
||
min_return_groups = target_groups
|
||
min_return_groups = max(1, int(min_return_groups)) if target_groups > 0 else 0
|
||
if target_groups > 0:
|
||
min_return_groups = min(min_return_groups, target_groups)
|
||
producers = max(1, int(producer_ranks) if producer_ranks is not None else int(world_size))
|
||
while True:
|
||
if rank == src_rank:
|
||
available = len(pending_groups)
|
||
need = max(0, target_groups - available)
|
||
if available >= min_return_groups:
|
||
rounds = 0
|
||
else:
|
||
rounds = 0 if need <= 0 else math.ceil(need / producers)
|
||
if need > 0 and available < min_return_groups and (time.time() - start) >= max(float(fetch_timeout_sec), 0.0):
|
||
rounds = -1
|
||
ctrl = [int(rounds)]
|
||
else:
|
||
ctrl = [0]
|
||
dist.broadcast_object_list(ctrl, src=src_rank)
|
||
rounds = int(ctrl[0])
|
||
if rounds < 0:
|
||
return None if rank == src_rank else False
|
||
if rounds == 0:
|
||
break
|
||
|
||
for _ in range(rounds):
|
||
remain = max(0.0, float(fetch_timeout_sec) - (time.time() - start))
|
||
local_group = manager_pop_one_group(manager, timeout=min(max(remain, 0.2), 30.0)) if (remain > 0.0 and rank != src_rank) else None
|
||
if rank == src_rank:
|
||
gathered = [None for _ in range(world_size)]
|
||
dist.gather_object(local_group, object_gather_list=gathered, dst=src_rank)
|
||
for item in gathered:
|
||
if item is not None:
|
||
pending_groups.append(item)
|
||
else:
|
||
dist.gather_object(local_group, dst=src_rank)
|
||
|
||
if rank == src_rank:
|
||
if len(pending_groups) < min_return_groups:
|
||
return None
|
||
out = [pending_groups.pop() for _ in range(min(target_groups, len(pending_groups)))]
|
||
return out if out else None
|
||
return True
|
||
|
||
|
||
# -------------------------- packing / PPO loss --------------------------
|
||
|
||
def pack_groups(groups, tokenizer, max_total_len, device, num_generations, strict_group_packing=True):
|
||
prompts, messages_batch, gt_batch, tools_batch, versions = [], [], [], [], []
|
||
trajectory_versions = []
|
||
completions, tool_events_batch, unfinished_batch, packed = [], [], [], []
|
||
dropped_groups = 0
|
||
for group in groups:
|
||
if len(group.trajectories) != num_generations:
|
||
dropped_groups += 1
|
||
continue
|
||
group_packed, group_completions, group_tool_events, group_unfinished = [], [], [], []
|
||
for traj in group.trajectories:
|
||
ids = traj.prompt_ids + traj.response_ids
|
||
mask = [0] * len(traj.prompt_ids) + traj.response_mask
|
||
old_logps = [0.0] * max(len(traj.prompt_ids) - 1, 0) + traj.response_old_logps
|
||
token_versions = [-1] * max(len(traj.prompt_ids) - 1, 0) + traj.response_versions
|
||
if len(ids) > max_total_len:
|
||
ids = ids[-max_total_len:]
|
||
mask = mask[-max_total_len:]
|
||
old_logps = old_logps[-(len(ids) - 1):]
|
||
token_versions = token_versions[-(len(ids) - 1):]
|
||
if len(ids) < 2 or sum(mask) <= 0 or len(old_logps) != len(ids) - 1 or len(token_versions) != len(ids) - 1:
|
||
group_packed = []
|
||
break
|
||
group_packed.append((ids, mask, old_logps, token_versions))
|
||
group_completions.append(traj.completion)
|
||
group_tool_events.append(traj.tool_events)
|
||
group_unfinished.append(traj.unfinished)
|
||
if strict_group_packing and len(group_packed) != num_generations:
|
||
dropped_groups += 1
|
||
continue
|
||
if not group_packed:
|
||
dropped_groups += 1
|
||
continue
|
||
prompts.append(group.prompt)
|
||
messages_batch.append(group.trajectories[0].messages)
|
||
gt_batch.append(group.trajectories[0].gt)
|
||
tools_batch.append(group.trajectories[0].tools)
|
||
versions.append(group.policy_version)
|
||
trajectory_versions.extend([group.policy_version] * len(group_packed))
|
||
packed.extend(group_packed)
|
||
completions.extend(group_completions)
|
||
tool_events_batch.extend(group_tool_events)
|
||
unfinished_batch.extend(group_unfinished)
|
||
if not packed:
|
||
raise RuntimeError("empty packed batch")
|
||
if len(packed) % num_generations != 0:
|
||
raise RuntimeError(f"packed trajectories not divisible by num_generations: {len(packed)} vs {num_generations}")
|
||
seq_lens = torch.tensor([len(ids) for ids, _, _, _ in packed], device=device, dtype=torch.long)
|
||
max_len = int(seq_lens.max().item())
|
||
pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
|
||
input_ids = torch.tensor([ids + [pad_id] * (max_len - len(ids)) for ids, _, _, _ in packed], device=device, dtype=torch.long)
|
||
full_response_masks = torch.tensor([mask + [0] * (max_len - len(mask)) for _, mask, _, _ in packed], device=device, dtype=torch.float32)
|
||
old_per_token_logps = torch.tensor([old_lp + [0.0] * ((max_len - 1) - len(old_lp)) for _, _, old_lp, _ in packed], device=device, dtype=torch.float32)
|
||
per_token_versions = torch.tensor([v + [-1] * ((max_len - 1) - len(v)) for _, _, _, v in packed], device=device, dtype=torch.long)
|
||
return {
|
||
"input_ids": input_ids,
|
||
"seq_lens": seq_lens,
|
||
"full_response_masks": full_response_masks,
|
||
"old_per_token_logps": old_per_token_logps,
|
||
"per_token_versions": per_token_versions,
|
||
"prompts": prompts,
|
||
"messages_batch": messages_batch,
|
||
"gt_batch": gt_batch,
|
||
"tools_batch": tools_batch,
|
||
"versions": versions,
|
||
"trajectory_versions": trajectory_versions,
|
||
"completions": completions,
|
||
"tool_events_batch": tool_events_batch,
|
||
"unfinished_batch": unfinished_batch,
|
||
"dropped_groups": dropped_groups,
|
||
}
|
||
|
||
|
||
def resolve_prox_logps(old_logps, current_logps_detached, token_versions, current_version, method="loglinear"):
|
||
old_logps = torch.nan_to_num(old_logps, nan=0.0, posinf=0.0, neginf=0.0)
|
||
current_logps_detached = torch.nan_to_num(current_logps_detached, nan=0.0, posinf=0.0, neginf=0.0)
|
||
if method == "behavior":
|
||
return old_logps
|
||
valid = token_versions >= 0
|
||
v_behav = token_versions.float()
|
||
v_theta = torch.full_like(v_behav, float(current_version))
|
||
v_prox = torch.full_like(v_behav, float(max(current_version - 1, 0)))
|
||
alpha = ((v_prox - v_behav) / (v_theta - v_behav).clamp(min=1.0)).clamp(min=0.0, max=1.0)
|
||
prox = old_logps + alpha * (current_logps_detached - old_logps)
|
||
return torch.where(valid, prox, old_logps)
|
||
|
||
|
||
def masked_stats(values, mask, prefix):
|
||
with torch.no_grad():
|
||
selected = values[mask.bool()]
|
||
if selected.numel() == 0:
|
||
return {f"{prefix}_mean": 0.0, f"{prefix}_p05": 0.0, f"{prefix}_p50": 0.0, f"{prefix}_p95": 0.0}
|
||
selected = selected.float().detach()
|
||
return {
|
||
f"{prefix}_mean": float(selected.mean().item()),
|
||
f"{prefix}_p05": float(torch.quantile(selected, 0.05).item()),
|
||
f"{prefix}_p50": float(torch.quantile(selected, 0.50).item()),
|
||
f"{prefix}_p95": float(torch.quantile(selected, 0.95).item()),
|
||
}
|
||
|
||
|
||
def build_advantages(args, rewards, num_generations):
|
||
grouped_rewards = rewards.view(-1, num_generations)
|
||
group_mean = grouped_rewards.mean(dim=1).repeat_interleave(num_generations)
|
||
group_std_raw = grouped_rewards.std(dim=1, unbiased=False)
|
||
group_std = group_std_raw.repeat_interleave(num_generations)
|
||
group_adv = (rewards - group_mean) / group_std.clamp(min=max(float(args.adv_min_std), 1e-4))
|
||
baseline = float(getattr(args, "reward_baseline", 0.0))
|
||
scale = max(float(getattr(args, "reward_scale", 3.0)), 1e-3)
|
||
baseline_adv = (rewards - baseline) / scale
|
||
zero_group = (group_std_raw < args.adv_min_std).repeat_interleave(num_generations)
|
||
zero_std_mode = str(getattr(args, "zero_std_advantage_mode", "baseline"))
|
||
if zero_std_mode == "zero":
|
||
zero_std_adv = torch.zeros_like(baseline_adv)
|
||
elif zero_std_mode == "positive_baseline":
|
||
zero_std_adv = baseline_adv.clamp(min=0.0)
|
||
else:
|
||
zero_std_adv = baseline_adv
|
||
if args.advantage_mode == "group":
|
||
advantages = group_adv
|
||
elif args.advantage_mode == "baseline":
|
||
advantages = baseline_adv
|
||
else:
|
||
advantages = torch.where(zero_group, zero_std_adv, group_adv)
|
||
if args.adv_clip > 0:
|
||
advantages = advantages.clamp(-args.adv_clip, args.adv_clip)
|
||
return advantages, {
|
||
"reward/group_std": float(group_std_raw.mean().item()),
|
||
"reward/zero_std_group_ratio": float((group_std_raw < args.adv_min_std).float().mean().item()),
|
||
"reward/nonzero_adv_ratio": float((advantages.abs() > 1e-6).float().mean().item()),
|
||
"adv/negative_ratio": float((advantages < -1e-6).float().mean().item()),
|
||
"adv/positive_ratio": float((advantages > 1e-6).float().mean().item()),
|
||
"adv/mean": float(advantages.mean().item()),
|
||
"adv/std": float(advantages.std(unbiased=False).item()),
|
||
}
|
||
|
||
|
||
# -------------------------- BC auxiliary loss --------------------------
|
||
|
||
def canonical_tool_call(name: str, args: Dict[str, Any]) -> str:
|
||
return "<tool_call>" + json.dumps({"name": name, "arguments": args}, ensure_ascii=False, separators=(",", ":")) + "</tool_call>"
|
||
|
||
|
||
def build_math_bc_examples(args, tokenizer, env, messages_batch, tools_batch, gt_batch, device):
|
||
examples = []
|
||
max_examples = int(getattr(args, "bc_max_examples", 16))
|
||
for messages, tools, gt_items in zip(messages_batch, tools_batch, gt_batch):
|
||
if len(examples) >= max_examples:
|
||
break
|
||
tools = tools or env.tool_specs()
|
||
text = task_text_from_messages(messages)
|
||
exprs = extract_arithmetic_expressions(text, max_exprs=max(1, int(getattr(args, "bc_max_exprs", 8))))
|
||
if not exprs:
|
||
continue
|
||
call_texts, tool_results = [], []
|
||
for expr in exprs:
|
||
try:
|
||
result = env.execute("calculate_math", {"expression": expr})
|
||
except Exception:
|
||
continue
|
||
call_texts.append(canonical_tool_call("calculate_math", {"expression": expr}))
|
||
tool_results.append(result)
|
||
if not call_texts:
|
||
continue
|
||
|
||
gt_clean = [str(x).strip() for x in (gt_items or []) if str(x).strip()]
|
||
# If GT count agrees with extracted expression count, final-answer BC uses GT; otherwise use tool results.
|
||
if gt_clean and len(gt_clean) == len(tool_results):
|
||
answer_values = gt_clean
|
||
else:
|
||
answer_values = [str(x.get("result", "")) for x in tool_results]
|
||
answer_text = "<answer>" + ", ".join(answer_values) + "</answer>"
|
||
|
||
prefix1 = apply_chat_template_compat(tokenizer, messages, tools=tools, add_generation_prompt=True)
|
||
examples.append((prefix1, "\n".join(call_texts)))
|
||
|
||
hist = [dict(m) for m in messages] + [{"role": "assistant", "content": "\n".join(call_texts)}]
|
||
for result in tool_results:
|
||
hist.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False, separators=(",", ":"))})
|
||
prefix2 = apply_chat_template_compat(tokenizer, hist, tools=tools, add_generation_prompt=True)
|
||
examples.append((prefix2, answer_text))
|
||
return examples[:max_examples]
|
||
|
||
|
||
def bc_coef_at_step(args) -> float:
|
||
base = float(getattr(args, "bc_coef", 0.15))
|
||
if base <= 0:
|
||
return 0.0
|
||
step = int(getattr(args, "current_step", 0))
|
||
hold = int(getattr(args, "bc_hold_steps", 1200))
|
||
decay = int(getattr(args, "bc_decay_steps", 4000))
|
||
min_coef = max(0.0, float(getattr(args, "bc_min_coef", 0.03)))
|
||
stop = int(getattr(args, "bc_stop_steps", 0))
|
||
if stop > 0 and step > stop:
|
||
return 0.0
|
||
if step <= hold:
|
||
return base
|
||
if decay <= 0:
|
||
return max(min_coef, 0.0)
|
||
frac = max(0.0, 1.0 - (step - hold) / max(decay, 1))
|
||
return max(base * frac, min_coef)
|
||
|
||
|
||
def compute_bc_aux_loss(args, batch, model, tokenizer, env, autocast_factory):
|
||
coef = bc_coef_at_step(args)
|
||
if coef <= 0:
|
||
zero = torch.tensor(0.0, device=args.device)
|
||
return zero, {"bc/loss": 0.0, "bc/coef": 0.0, "bc/examples": 0.0}
|
||
examples = build_math_bc_examples(args, tokenizer, env, batch["messages_batch"], batch["tools_batch"], batch["gt_batch"], args.device)
|
||
if not examples:
|
||
zero = torch.tensor(0.0, device=args.device)
|
||
return zero, {"bc/loss": 0.0, "bc/coef": float(coef), "bc/examples": 0.0}
|
||
|
||
pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
|
||
eos_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else pad_id
|
||
max_len = int(getattr(args, "bc_max_len", min(args.max_total_len, 1536)))
|
||
all_ids, all_labels = [], []
|
||
for prefix, target in examples:
|
||
prefix_ids = tokenizer(prefix, add_special_tokens=False)["input_ids"]
|
||
target_ids = tokenizer(target, add_special_tokens=False)["input_ids"]
|
||
if eos_id is not None:
|
||
target_ids = target_ids + [int(eos_id)]
|
||
if len(target_ids) >= max_len:
|
||
ids = target_ids[-max_len:]
|
||
labels = ids[:]
|
||
else:
|
||
keep_prefix = max(0, max_len - len(target_ids))
|
||
kept_prefix_ids = prefix_ids[-keep_prefix:] if keep_prefix > 0 else []
|
||
ids = kept_prefix_ids + target_ids
|
||
labels = [-100] * len(kept_prefix_ids) + target_ids
|
||
if len(ids) < 2 or all(x == -100 for x in labels[1:]):
|
||
continue
|
||
all_ids.append(ids)
|
||
all_labels.append(labels)
|
||
|
||
if not all_ids:
|
||
zero = torch.tensor(0.0, device=args.device)
|
||
return zero, {"bc/loss": 0.0, "bc/coef": float(coef), "bc/examples": 0.0}
|
||
|
||
ml = max(len(x) for x in all_ids)
|
||
input_ids = torch.tensor([x + [pad_id] * (ml - len(x)) for x in all_ids], device=args.device, dtype=torch.long)
|
||
labels = torch.tensor([x + [-100] * (ml - len(x)) for x in all_labels], device=args.device, dtype=torch.long)
|
||
attn = (input_ids != pad_id).long()
|
||
with autocast_factory():
|
||
res = unwrap_model(model)(input_ids=input_ids, attention_mask=attn)
|
||
logits = res.logits[:, :-1, :]
|
||
labels_shift = labels[:, 1:]
|
||
loss = F.cross_entropy(logits.float().reshape(-1, logits.size(-1)), labels_shift.reshape(-1), ignore_index=-100)
|
||
return loss, {"bc/loss": float(loss.detach().item()), "bc/coef": float(coef), "bc/examples": float(len(all_ids))}
|
||
|
||
|
||
def build_dashboard_metrics(detail, prefix="train"):
|
||
return {
|
||
f"{prefix}/accuracy_exact_match": float(detail.get("exact_match", 0.0)),
|
||
f"{prefix}/accuracy_answer_cov": float(detail.get("answer_cov", 0.0)),
|
||
f"{prefix}/format_answer_tag_rate": float(detail.get("answer_tag", 0.0)),
|
||
f"{prefix}/format_exact_answer_tag_rate": float(detail.get("exact_answer_tag", 0.0)),
|
||
f"{prefix}/tool_success_rate": float(detail.get("tool_success_rate", 0.0)),
|
||
f"{prefix}/tool_correct_choice_rate": float(detail.get("correct_tool_choice", 0.0)),
|
||
f"{prefix}/tool_math_call_coverage": float(detail.get("math_call_coverage", 0.0)),
|
||
f"{prefix}/tool_expr_order_cov": float(detail.get("tool_expr_order_cov", 0.0)),
|
||
f"{prefix}/tool_expr_set_cov": float(detail.get("tool_expr_set_cov", 0.0)),
|
||
f"{prefix}/tool_expr_exact": float(detail.get("tool_expr_exact", 0.0)),
|
||
f"{prefix}/tool_value_order_cov": float(detail.get("tool_value_order_cov", 0.0)),
|
||
f"{prefix}/tool_preferred_math_rate": float(detail.get("preferred_math_rate", 0.0)),
|
||
f"{prefix}/tool_canonical_rate": float(detail.get("canonical_tool_rate", 0.0)),
|
||
}
|
||
|
||
|
||
def compute_async_loss(args, batch, model, ref_model, reward_model, tokenizer, autocast_factory, lm_config, env):
|
||
input_ids = batch["input_ids"]
|
||
pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
|
||
full_mask = (input_ids != pad_id).long()
|
||
full_response_masks = batch["full_response_masks"]
|
||
old_per_token_logps = torch.nan_to_num(batch["old_per_token_logps"], nan=0.0, posinf=0.0, neginf=0.0)
|
||
per_token_versions = batch["per_token_versions"]
|
||
|
||
with autocast_factory():
|
||
res = unwrap_model(model)(input_ids=input_ids, attention_mask=full_mask)
|
||
aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
|
||
logits = res.logits[:, :-1, :]
|
||
logits = torch.nan_to_num(logits.float(), nan=0.0, posinf=30.0, neginf=-30.0).clamp(-80.0, 80.0)
|
||
target_ids = input_ids[:, 1:].clamp(min=0, max=logits.size(-1) - 1)
|
||
per_token_logps = F.log_softmax(logits, dim=-1).gather(2, target_ids.unsqueeze(-1)).squeeze(-1)
|
||
with torch.no_grad():
|
||
ref_per_token_logps = compute_per_token_logps_safe(ref_model, input_ids, full_mask, autocast_factory=None)
|
||
|
||
completion_mask = full_response_masks[:, 1:]
|
||
is_eos = (input_ids[:, 1:] == tokenizer.eos_token_id) & completion_mask.bool()
|
||
eos_idx = torch.full((completion_mask.size(0),), completion_mask.size(1) - 1, device=args.device, dtype=torch.long)
|
||
has_eos = is_eos.any(dim=1)
|
||
eos_idx[has_eos] = is_eos.int().argmax(dim=1)[has_eos]
|
||
pos = torch.arange(completion_mask.size(1), device=args.device).unsqueeze(0)
|
||
completion_mask = completion_mask * (pos <= eos_idx.unsqueeze(1)).float()
|
||
trajectory_versions = batch.get("trajectory_versions", [])
|
||
if trajectory_versions and len(trajectory_versions) == completion_mask.size(0):
|
||
row_versions = torch.tensor(trajectory_versions, device=args.device, dtype=torch.long)
|
||
row_staleness = (torch.full_like(row_versions, int(args.current_version)) - row_versions).clamp(min=0)
|
||
else:
|
||
row_staleness = torch.zeros(completion_mask.size(0), device=args.device, dtype=torch.long)
|
||
stale_discount = min(max(float(getattr(args, "async_stale_discount", 0.75)), 0.0), 1.0)
|
||
row_weights = torch.pow(torch.full_like(row_staleness, stale_discount, dtype=torch.float32), row_staleness.float()) if stale_discount < 1.0 else torch.ones_like(row_staleness, dtype=torch.float32)
|
||
|
||
rewards, reward_details = calculate_rewards(
|
||
args,
|
||
batch["prompts"],
|
||
batch["completions"],
|
||
batch["gt_batch"],
|
||
batch["tools_batch"],
|
||
batch["tool_events_batch"],
|
||
batch["unfinished_batch"],
|
||
args.num_generations,
|
||
device=args.device,
|
||
reward_model=reward_model,
|
||
return_details=True,
|
||
messages_batch=batch["messages_batch"],
|
||
)
|
||
advantages, adv_metrics = build_advantages(args, rewards, args.num_generations)
|
||
ref_delta = (ref_per_token_logps - per_token_logps).clamp(-20, 20)
|
||
per_token_kl = torch.exp(ref_delta) - ref_delta - 1
|
||
|
||
if args.use_decoupled_ppo:
|
||
prox_logps = resolve_prox_logps(old_per_token_logps, per_token_logps.detach(), per_token_versions, args.current_version, args.prox_logp_method)
|
||
ratio = torch.exp((per_token_logps - prox_logps).clamp(-10, 10))
|
||
behav_weight = torch.exp((prox_logps - old_per_token_logps).clamp(-10, 10)).detach()
|
||
if args.behav_imp_weight_cap > 0:
|
||
behav_cap_frac = ((behav_weight > args.behav_imp_weight_cap).float() * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
|
||
behav_weight = behav_weight.clamp(max=args.behav_imp_weight_cap)
|
||
else:
|
||
behav_cap_frac = torch.tensor(0.0, device=args.device)
|
||
clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
|
||
surrogate = behav_weight * torch.min(ratio * advantages.unsqueeze(1), clipped_ratio * advantages.unsqueeze(1))
|
||
m2_mask_frac = torch.tensor(0.0, device=args.device)
|
||
if args.m2_threshold > 0:
|
||
m2 = (old_per_token_logps - prox_logps).pow(2)
|
||
keep = (m2 <= args.m2_threshold).float()
|
||
m2_mask_frac = ((1.0 - keep) * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
|
||
completion_mask = completion_mask * keep
|
||
else:
|
||
ratio = torch.exp((per_token_logps - old_per_token_logps).clamp(-10, 10))
|
||
clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
|
||
surrogate = torch.min(ratio * advantages.unsqueeze(1), clipped_ratio * advantages.unsqueeze(1))
|
||
behav_weight = torch.ones_like(ratio)
|
||
behav_cap_frac = torch.tensor(0.0, device=args.device)
|
||
m2_mask_frac = torch.tensor(0.0, device=args.device)
|
||
|
||
per_token_loss = -(surrogate - args.beta * per_token_kl)
|
||
token_counts = completion_mask.sum(dim=1)
|
||
valid_rows = token_counts > 0
|
||
if getattr(args, "loss_reduction", "token") == "token":
|
||
weighted_mask = completion_mask * row_weights.unsqueeze(1)
|
||
policy_loss = (per_token_loss * weighted_mask).sum() / weighted_mask.sum().clamp(min=1.0)
|
||
else:
|
||
row_loss = (per_token_loss * completion_mask).sum(dim=1) / token_counts.clamp(min=1.0)
|
||
policy_loss = (row_loss[valid_rows] * row_weights[valid_rows]).sum() / row_weights[valid_rows].sum().clamp(min=1.0) if valid_rows.any() else per_token_loss.sum() * 0.0
|
||
|
||
bc_loss, bc_metrics = compute_bc_aux_loss(args, batch, model, tokenizer, env, autocast_factory)
|
||
bc_coef = float(bc_metrics.get("bc/coef", 0.0))
|
||
loss = policy_loss + aux_loss + bc_coef * bc_loss
|
||
|
||
with torch.no_grad():
|
||
mask_denom = completion_mask.sum().clamp(min=1.0)
|
||
ratio_clip_frac = ((((ratio < 1 - args.epsilon) | (ratio > 1 + args.epsilon)).float() * completion_mask).sum() / mask_denom).item()
|
||
approx_kl = (((per_token_logps - old_per_token_logps) * completion_mask).sum() / mask_denom).item()
|
||
old_current_abs_delta = (((per_token_logps - old_per_token_logps).abs() * completion_mask).sum() / mask_denom).item()
|
||
version_gaps = torch.tensor([args.current_version - v for v in batch["versions"]], device=args.device, dtype=torch.float32) if batch["versions"] else torch.zeros(1, device=args.device)
|
||
valid_version_mask = (per_token_versions >= 0) & completion_mask.bool()
|
||
token_stale = torch.where(valid_version_mask, torch.full_like(per_token_versions, int(args.current_version)) - per_token_versions, torch.zeros_like(per_token_versions)).float()
|
||
reward_scalar = {k: safe_mean([d[k] for d in reward_details]) for k in reward_details[0].keys()} if reward_details else {}
|
||
metrics = {
|
||
"loss/total": float(loss.detach().item()),
|
||
"loss/policy": float(policy_loss.detach().item()),
|
||
"loss/aux": float(aux_loss.detach().item()),
|
||
"reward/mean": float(rewards.mean().item()),
|
||
"reward/min": float(rewards.min().item()),
|
||
"reward/max": float(rewards.max().item()),
|
||
"reward/baseline": float(getattr(args, "reward_baseline", 0.0)),
|
||
"length/avg_completion_tokens": float(token_counts.float().mean().item()),
|
||
"policy/kl_ref_k3": float(((per_token_kl * completion_mask).sum() / mask_denom).item()),
|
||
"policy/approx_kl_old": float(approx_kl),
|
||
"policy/approx_kl_old_abs": float(abs(approx_kl)),
|
||
"policy/old_current_abs_delta": float(old_current_abs_delta),
|
||
"policy/entropy": float(logits_entropy_safe(logits, completion_mask)),
|
||
"policy/ratio_clip_frac": float(ratio_clip_frac),
|
||
"policy/behav_imp_weight_cap_frac": float(behav_cap_frac.item()),
|
||
"policy/m2_mask_frac": float(m2_mask_frac.item()),
|
||
"policy/valid_completion_rows": float(valid_rows.float().mean().item()),
|
||
"async/row_staleness_mean": float(row_staleness.float().mean().item()),
|
||
"async/row_weight_mean": float(row_weights.mean().item()),
|
||
"async/staleness_group_mean": float(version_gaps.mean().item()),
|
||
"async/staleness_group_max": float(version_gaps.max().item()),
|
||
"pack/dropped_groups": float(batch.get("dropped_groups", 0)),
|
||
**adv_metrics,
|
||
**masked_stats(ratio.detach(), completion_mask, "policy/ratio"),
|
||
**masked_stats(behav_weight.detach(), completion_mask, "policy/behav_weight"),
|
||
**masked_stats(token_stale.detach(), valid_version_mask.float(), "async/token_staleness"),
|
||
**bc_metrics,
|
||
}
|
||
for k, v in reward_scalar.items():
|
||
metrics[f"reward_detail/{k}"] = float(v)
|
||
metrics.update(build_dashboard_metrics(reward_scalar, prefix="train"))
|
||
return loss, metrics, {"rewards": [float(x) for x in rewards.detach().cpu().tolist()], "reward_details": reward_details}
|
||
|
||
|
||
@torch.no_grad()
|
||
def run_agent_eval(args, eval_ds, rollout_engine, tokenizer, env, engine_lock, policy_version, reward_model=None):
|
||
if eval_ds is None or len(eval_ds) == 0 or args.eval_samples <= 0:
|
||
return {}
|
||
indices = list(getattr(args, "eval_fixed_indices", []))
|
||
if not indices:
|
||
indices = build_fixed_eval_indices(eval_ds, args.eval_samples, seed=getattr(args, "eval_seed", 1234))
|
||
indices = indices[: min(args.eval_samples, len(indices), len(eval_ds))]
|
||
prompts, completions, gt_batch, tools_batch, events_batch, unfinished_batch, messages_batch = [], [], [], [], [], [], []
|
||
with engine_lock:
|
||
for idx in indices:
|
||
sample = eval_ds[idx]
|
||
traj = rollout_single_online(
|
||
rollout_engine,
|
||
tokenizer,
|
||
env,
|
||
sample,
|
||
args.max_turns,
|
||
args.eval_max_gen_len if args.eval_max_gen_len > 0 else args.max_gen_len,
|
||
args.max_context_len,
|
||
0.0,
|
||
args.device,
|
||
policy_version,
|
||
args.eval_temperature,
|
||
)
|
||
prompts.append(traj.prompt)
|
||
messages_batch.append(traj.messages)
|
||
completions.append(traj.completion)
|
||
gt_batch.append(traj.gt)
|
||
tools_batch.append(traj.tools)
|
||
events_batch.append(traj.tool_events)
|
||
unfinished_batch.append(traj.unfinished)
|
||
rewards, details = calculate_rewards(
|
||
args,
|
||
prompts,
|
||
completions,
|
||
gt_batch,
|
||
tools_batch,
|
||
events_batch,
|
||
unfinished_batch,
|
||
1,
|
||
args.device,
|
||
reward_model,
|
||
True,
|
||
messages_batch=messages_batch,
|
||
)
|
||
scalar = {k: safe_mean([d[k] for d in details]) for k in details[0].keys()} if details else {}
|
||
out = {
|
||
"eval/reward": float(rewards.mean().item()),
|
||
"eval/exact_match": float(scalar.get("exact_match", 0.0)),
|
||
"eval/answer_cov": float(scalar.get("answer_cov", 0.0)),
|
||
"eval/strict_exact_match": float(scalar.get("strict_exact_match", 0.0)),
|
||
"eval/tool_success_rate": float(scalar.get("tool_success_rate", 0.0)),
|
||
"eval/tool_correct_choice_rate": float(scalar.get("correct_tool_choice", 0.0)),
|
||
"eval/math_call_coverage": float(scalar.get("math_call_coverage", 0.0)),
|
||
"eval/tool_expr_order_cov": float(scalar.get("tool_expr_order_cov", 0.0)),
|
||
"eval/tool_expr_set_cov": float(scalar.get("tool_expr_set_cov", 0.0)),
|
||
"eval/tool_expr_exact": float(scalar.get("tool_expr_exact", 0.0)),
|
||
"eval/tool_value_order_cov": float(scalar.get("tool_value_order_cov", 0.0)),
|
||
"eval/tool_value_set_cov": float(scalar.get("tool_value_set_cov", 0.0)),
|
||
"eval/tool_value_exact": float(scalar.get("tool_value_exact", 0.0)),
|
||
"eval/preferred_math_rate": float(scalar.get("preferred_math_rate", 0.0)),
|
||
"eval/tool_canonical_rate": float(scalar.get("canonical_tool_rate", 0.0)),
|
||
"eval/non_math_tool_calls": float(scalar.get("non_math_tool_calls", 0.0)),
|
||
"eval/final_refs_tool": float(scalar.get("final_refs_tool", 0.0)),
|
||
"eval/answer_tag": float(scalar.get("answer_tag", 0.0)),
|
||
"eval/exact_answer_tag": float(scalar.get("exact_answer_tag", 0.0)),
|
||
"eval/extra_numbers": float(scalar.get("extra_numbers", 0.0)),
|
||
"eval/missing_numbers": float(scalar.get("missing_numbers", 0.0)),
|
||
"eval/order_ok": float(scalar.get("order_ok", 0.0)),
|
||
"eval/avg_answer_len": float(scalar.get("answer_len", 0.0)),
|
||
"eval/unfinished_rate": float(scalar.get("unfinished", 0.0)),
|
||
"eval/reward_clipped_rate": float(scalar.get("clipped", 0.0)),
|
||
}
|
||
out.update(build_dashboard_metrics(scalar, prefix="eval_dashboard"))
|
||
return out
|
||
|
||
|
||
def maybe_log_debug_sample(step, groups, diagnostics, debug_interval):
|
||
if not is_main_process() or debug_interval <= 0 or step % debug_interval != 0 or not groups:
|
||
return
|
||
rewards = diagnostics.get("rewards", [])
|
||
details = diagnostics.get("reward_details", [])
|
||
group = groups[0]
|
||
Logger(f"[DEBUG] step={step}, data_index={group.data_index}, policy_version={group.policy_version}")
|
||
Logger("-" * 100)
|
||
Logger(group.prompt[:1200])
|
||
for i, traj in enumerate(group.trajectories[:4]):
|
||
reward = rewards[i] if i < len(rewards) else None
|
||
detail = details[i] if i < len(details) else {}
|
||
Logger(
|
||
f"[DEBUG] traj={i}, reward={reward}, exact={detail.get('exact_match')}, cov={detail.get('answer_cov')}, "
|
||
f"tag={detail.get('answer_tag')}, tool_val={detail.get('tool_value_order_cov')}, "
|
||
f"tool_choice={detail.get('correct_tool_choice')}, tool_success={detail.get('tool_success_rate')}, "
|
||
f"unfinished={traj.unfinished}, events={len(traj.tool_events)}"
|
||
)
|
||
for turn_id, turn_text in enumerate(traj.turn_outputs):
|
||
Logger(f"[DEBUG] assistant turn={turn_id}\n{turn_text[:1200]}")
|
||
if traj.tool_events:
|
||
Logger(f"[DEBUG] tool_events={json.dumps(traj.tool_events[:8], ensure_ascii=False)}")
|
||
Logger("-" * 100)
|
||
|
||
|
||
# -------------------------- CLI / main --------------------------
|
||
|
||
def build_args():
|
||
p = argparse.ArgumentParser(description="MiniMind stable async agent RL trainer - tool-use optimized v5 for 4-GPU torch")
|
||
p.add_argument("--save_dir", type=str, default="../out")
|
||
p.add_argument("--save_weight", type=str, default="agentic_async_rl_tooluse_v5_fix")
|
||
p.add_argument("--epochs", type=int, default=1)
|
||
p.add_argument("--total_steps", type=int, default=0)
|
||
p.add_argument("--scale_total_steps_by_world_size", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--batch_size", type=int, default=2)
|
||
p.add_argument("--learning_rate", type=float, default=2e-7)
|
||
p.add_argument("--weight_decay", type=float, default=0.01)
|
||
p.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||
p.add_argument("--dtype", type=str, default="bfloat16", choices=["bfloat16", "float16", "float32"])
|
||
p.add_argument("--accumulation_steps", type=int, default=2)
|
||
p.add_argument("--grad_clip", type=float, default=0.5)
|
||
p.add_argument("--log_interval", type=int, default=1)
|
||
p.add_argument("--save_interval", type=int, default=100)
|
||
p.add_argument("--seed", type=int, default=42)
|
||
p.add_argument("--hidden_size", type=int, default=768)
|
||
p.add_argument("--num_hidden_layers", type=int, default=8)
|
||
p.add_argument("--use_moe", type=int, default=0, choices=[0, 1])
|
||
p.add_argument("--max_seq_len", type=int, default=1024)
|
||
p.add_argument("--max_context_len", type=int, default=1024)
|
||
p.add_argument("--max_gen_len", type=int, default=192)
|
||
p.add_argument("--eval_max_gen_len", type=int, default=192)
|
||
p.add_argument("--max_total_len", type=int, default=1536)
|
||
p.add_argument("--data_path", type=str, default="/workspace/minimind/dataset/agent_rl_math.jsonl")
|
||
p.add_argument("--max_samples", type=int, default=0)
|
||
p.add_argument("--from_weight", type=str, default="full_sft")
|
||
p.add_argument("--from_resume", type=int, default=0, choices=[0, 1])
|
||
p.add_argument("--ref_from_rl_checkpoint", type=int, default=0, choices=[0, 1])
|
||
p.add_argument("--use_compile", type=int, default=0, choices=[0, 1])
|
||
p.add_argument("--num_generations", type=int, default=4)
|
||
p.add_argument("--max_turns", type=int, default=3)
|
||
p.add_argument("--thinking_ratio", type=float, default=0.0)
|
||
p.add_argument("--temperature", type=float, default=0.7)
|
||
p.add_argument("--top_p", type=float, default=0.88)
|
||
p.add_argument("--top_k", type=int, default=40)
|
||
p.add_argument("--repetition_penalty", type=float, default=1.03)
|
||
p.add_argument("--torch_rollout_use_cache", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--filter_mode", type=str, default="math_only", choices=["all", "math_only"])
|
||
p.add_argument("--require_gt", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--time_mode", type=str, default="fixed", choices=["fixed", "live"])
|
||
p.add_argument("--tool_subset", type=str, default="calculate_math,unit_converter,get_current_time")
|
||
p.add_argument("--inject_answer_instruction", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--augment_calculator_alias", type=int, default=0, choices=[0, 1])
|
||
p.add_argument("--force_all_tools", type=int, default=0, choices=[0, 1])
|
||
|
||
p.add_argument("--epsilon", type=float, default=0.20)
|
||
p.add_argument("--beta", type=float, default=0.03)
|
||
p.add_argument("--use_decoupled_ppo", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--prox_logp_method", type=str, default="loglinear", choices=["loglinear", "behavior"])
|
||
p.add_argument("--behav_imp_weight_cap", type=float, default=1.5)
|
||
p.add_argument("--m2_threshold", type=float, default=0.0)
|
||
p.add_argument("--advantage_mode", type=str, default="mixed", choices=["group", "baseline", "mixed"])
|
||
p.add_argument("--adv_min_std", type=float, default=0.10)
|
||
p.add_argument("--zero_std_advantage_mode", type=str, default="baseline", choices=["zero", "baseline", "positive_baseline"])
|
||
p.add_argument("--adv_clip", type=float, default=3.0)
|
||
p.add_argument("--reward_ema_decay", type=float, default=0.98)
|
||
p.add_argument("--reward_scale", type=float, default=3.0)
|
||
p.add_argument("--beta_auto_tune", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--beta_control_metric", type=str, default="old_current_abs_delta", choices=["old_current_abs_delta", "approx_kl_old", "kl_ref"])
|
||
p.add_argument("--beta_update_interval", type=int, default=4)
|
||
p.add_argument("--beta_ema_decay", type=float, default=0.90)
|
||
p.add_argument("--kl_target", type=float, default=0.06)
|
||
p.add_argument("--beta_min", type=float, default=0.01)
|
||
p.add_argument("--beta_max", type=float, default=0.08)
|
||
p.add_argument("--hard_kl_stop_mult", type=float, default=8.0)
|
||
p.add_argument("--loss_reduction", type=str, default="sequence", choices=["token", "sequence"])
|
||
|
||
p.add_argument("--bc_coef", type=float, default=0.16)
|
||
p.add_argument("--bc_min_coef", type=float, default=0.03)
|
||
p.add_argument("--bc_stop_steps", type=int, default=0)
|
||
p.add_argument("--bc_hold_steps", type=int, default=1200)
|
||
p.add_argument("--bc_decay_steps", type=int, default=4000)
|
||
p.add_argument("--bc_max_examples", type=int, default=24)
|
||
p.add_argument("--bc_max_exprs", type=int, default=8)
|
||
p.add_argument("--bc_max_len", type=int, default=1536)
|
||
|
||
p.add_argument("--format_enforce_mode", type=str, default="semi_hard", choices=["soft", "semi_hard", "hard"])
|
||
p.add_argument("--format_warmup_steps", type=int, default=100)
|
||
p.add_argument("--format_missing_penalty_initial", type=float, default=0.30)
|
||
p.add_argument("--format_missing_penalty_final", type=float, default=1.50)
|
||
p.add_argument("--no_answer_tag_reward_cap", type=float, default=2.0)
|
||
p.add_argument("--strict_order_reward", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--strict_exact_reward", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--reward_base", type=float, default=-0.10)
|
||
p.add_argument("--reward_answer_cov_weight", type=float, default=9.0)
|
||
p.add_argument("--reward_exact_match_weight", type=float, default=9.0)
|
||
p.add_argument("--reward_answer_tag_weight", type=float, default=0.90)
|
||
p.add_argument("--reward_exact_answer_tag_weight", type=float, default=0.80)
|
||
p.add_argument("--reward_tool_success_weight", type=float, default=0.60)
|
||
p.add_argument("--reward_final_refs_tool_weight", type=float, default=0.45)
|
||
p.add_argument("--reward_tool_fail_penalty", type=float, default=0.55)
|
||
p.add_argument("--reward_no_tool_penalty", type=float, default=1.50)
|
||
p.add_argument("--reward_math_call_cov_weight", type=float, default=0.60)
|
||
p.add_argument("--reward_tool_expr_order_weight", type=float, default=0.60)
|
||
p.add_argument("--reward_tool_expr_set_weight", type=float, default=0.20)
|
||
p.add_argument("--reward_tool_expr_exact_weight", type=float, default=0.20)
|
||
p.add_argument("--reward_tool_value_order_weight", type=float, default=2.40)
|
||
p.add_argument("--reward_tool_value_set_weight", type=float, default=0.80)
|
||
p.add_argument("--reward_tool_value_exact_weight", type=float, default=0.80)
|
||
p.add_argument("--reward_tool_over_call_penalty", type=float, default=0.25)
|
||
p.add_argument("--reward_tool_under_call_penalty", type=float, default=0.35)
|
||
p.add_argument("--reward_preferred_math_tool_weight", type=float, default=0.70)
|
||
p.add_argument("--reward_correct_tool_choice_weight", type=float, default=0.70)
|
||
p.add_argument("--reward_tool_count_exact_weight", type=float, default=0.30)
|
||
p.add_argument("--reward_canonical_tool_weight", type=float, default=0.25)
|
||
p.add_argument("--reward_wrong_tool_penalty", type=float, default=0.55)
|
||
p.add_argument("--reward_alias_only_penalty", type=float, default=0.35)
|
||
p.add_argument("--extra_number_penalty", type=float, default=0.55)
|
||
p.add_argument("--missing_number_penalty", type=float, default=0.70)
|
||
p.add_argument("--order_mismatch_penalty", type=float, default=1.20)
|
||
p.add_argument("--reward_cap_warmup_steps", type=int, default=300)
|
||
p.add_argument("--non_exact_reward_cap_initial", type=float, default=4.0)
|
||
p.add_argument("--non_exact_reward_cap", type=float, default=3.5)
|
||
p.add_argument("--zero_answer_cov_reward_cap_initial", type=float, default=0.5)
|
||
p.add_argument("--zero_answer_cov_reward_cap", type=float, default=0.0)
|
||
p.add_argument("--max_reward_answer_len", type=int, default=96)
|
||
|
||
p.add_argument("--rollout_engine", type=str, default="torch", choices=["torch", "sglang"])
|
||
p.add_argument("--multi_gpu_mode", type=str, default="trainer_rollout", choices=["trainer_rollout", "ddp"])
|
||
p.add_argument("--sglang_base_url", type=str, default="http://localhost:8998")
|
||
p.add_argument("--sglang_model_path", type=str, default="../model")
|
||
p.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agentic_async")
|
||
p.add_argument("--async_workers", type=int, default=1)
|
||
p.add_argument("--max_version_gap", type=int, default=-1)
|
||
p.add_argument("--max_head_offpolicyness", type=int, default=1)
|
||
p.add_argument("--max_queue_groups", type=int, default=0)
|
||
p.add_argument("--global_prefetch_groups", type=int, default=0)
|
||
p.add_argument("--prefill_queue_groups", type=int, default=0)
|
||
p.add_argument("--prefill_timeout_sec", type=float, default=180.0)
|
||
p.add_argument("--fetch_timeout_sec", type=float, default=180.0)
|
||
p.add_argument("--policy_sync_interval", type=int, default=1)
|
||
p.add_argument("--async_stale_discount", type=float, default=0.65)
|
||
p.add_argument("--strict_group_packing", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--eval_ratio", type=float, default=0.02)
|
||
p.add_argument("--eval_max_pool", type=int, default=256)
|
||
p.add_argument("--eval_interval", type=int, default=100)
|
||
p.add_argument("--eval_samples", type=int, default=32)
|
||
p.add_argument("--eval_seed", type=int, default=1234)
|
||
p.add_argument("--eval_temperature", type=float, default=0.0)
|
||
p.add_argument("--eval_before_train", type=int, default=1, choices=[0, 1])
|
||
p.add_argument("--use_reward_model", type=int, default=0, choices=[0, 1])
|
||
p.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward")
|
||
p.add_argument("--use_wandb", default=True, type=str2bool)
|
||
p.add_argument("--wandb_project", type=str, default="MiniMind-Agentic-Async-RL")
|
||
p.add_argument("--debug_interval", type=int, default=0)
|
||
return p.parse_args()
|
||
|
||
|
||
def build_autocast_factory(device_type, dtype):
|
||
if device_type == "cpu" or dtype == torch.float32:
|
||
return lambda: nullcontext()
|
||
return lambda: torch.cuda.amp.autocast(dtype=dtype)
|
||
|
||
|
||
def maybe_step_optimizer(args, model, optimizer, scheduler):
|
||
grad_norm = 0.0
|
||
if args.grad_clip > 0:
|
||
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||
grad_norm = float(norm.detach().item()) if torch.is_tensor(norm) else float(norm)
|
||
bad_grad = any(p.grad is not None and not torch.isfinite(p.grad).all() for p in model.parameters())
|
||
if bad_grad or not math.isfinite(grad_norm):
|
||
Logger("[WARN] non-finite gradient; skip optimizer step")
|
||
optimizer.zero_grad(set_to_none=True)
|
||
return False, grad_norm
|
||
optimizer.step()
|
||
scheduler.step()
|
||
optimizer.zero_grad(set_to_none=True)
|
||
return True, grad_norm
|
||
|
||
|
||
def beta_control_metric_key(args) -> str:
|
||
metric = getattr(args, "beta_control_metric", "old_current_abs_delta")
|
||
if metric == "kl_ref":
|
||
return "policy/kl_ref_k3"
|
||
if metric == "approx_kl_old":
|
||
return "policy/approx_kl_old_abs"
|
||
return "policy/old_current_abs_delta"
|
||
|
||
|
||
def maybe_update_beta(args, metrics: Dict[str, Any], opt_step: int) -> float:
|
||
metric_key = beta_control_metric_key(args)
|
||
observed = finite_float(metrics.get(metric_key, 0.0), 0.0)
|
||
prev_ema = getattr(args, "beta_ema", observed)
|
||
decay = min(max(float(getattr(args, "beta_ema_decay", 0.90)), 0.0), 0.999)
|
||
ema = observed if not math.isfinite(prev_ema) else decay * prev_ema + (1.0 - decay) * observed
|
||
args.beta_ema = float(ema)
|
||
if not getattr(args, "beta_auto_tune", False):
|
||
return float(ema)
|
||
|
||
interval = max(1, int(getattr(args, "beta_update_interval", 4)))
|
||
if opt_step % interval != 0:
|
||
return float(ema)
|
||
|
||
target = max(finite_float(getattr(args, "kl_target", 0.08), 0.08), 1e-6)
|
||
upper = target * 1.25
|
||
lower = target * 0.80
|
||
max_mult = 1.12
|
||
min_mult = 1.0 / max_mult
|
||
if ema > upper:
|
||
ratio = min(max(ema / target, 1.0), max_mult * max_mult)
|
||
args.beta = min(args.beta * max(min(math.sqrt(ratio), max_mult), 1.0), args.beta_max)
|
||
elif ema < lower:
|
||
ratio = min(max(target / max(ema, 1e-6), 1.0), max_mult * max_mult)
|
||
args.beta = max(args.beta * min(max(1.0 / math.sqrt(ratio), min_mult), 1.0), args.beta_min)
|
||
return float(ema)
|
||
|
||
|
||
def split_prefetched_groups(fetched_groups, pending_groups, batch_size: int, current_version: int, max_version_gap: int):
|
||
if not fetched_groups:
|
||
return None, 0
|
||
batch_size = int(batch_size)
|
||
current_version = int(current_version)
|
||
max_version_gap = max(0, int(max_version_gap))
|
||
stale_drops = 0
|
||
batch = []
|
||
remainder = []
|
||
|
||
def keep_group(group) -> bool:
|
||
nonlocal stale_drops
|
||
if group is None:
|
||
return False
|
||
if current_version - int(group.policy_version) > max_version_gap:
|
||
stale_drops += 1
|
||
return False
|
||
return True
|
||
|
||
for group in fetched_groups:
|
||
if not keep_group(group):
|
||
continue
|
||
if len(batch) < batch_size:
|
||
batch.append(group)
|
||
else:
|
||
remainder.append(group)
|
||
|
||
while len(batch) < batch_size and pending_groups:
|
||
group = pending_groups.pop()
|
||
if keep_group(group):
|
||
batch.append(group)
|
||
|
||
if remainder:
|
||
pending_groups.extend(reversed(remainder))
|
||
return (batch if len(batch) >= batch_size else None), stale_drops
|
||
|
||
|
||
def parse_tool_subset(text: str) -> List[str]:
|
||
allowed = {"calculate_math", "unit_converter", "get_current_time"}
|
||
names = [x.strip() for x in str(text).split(",") if x.strip()]
|
||
names = [x for x in names if x in allowed]
|
||
return names or ["calculate_math", "unit_converter", "get_current_time"]
|
||
|
||
|
||
def main():
|
||
args = build_args()
|
||
local_rank = init_distributed_mode_compat()
|
||
if dist.is_initialized():
|
||
args.device = f"cuda:{local_rank}"
|
||
rank = dist.get_rank() if dist.is_initialized() else 0
|
||
setup_seed(args.seed + rank)
|
||
use_default_max_version_gap = args.max_version_gap < 0
|
||
args.use_decoupled_ppo = bool(args.use_decoupled_ppo)
|
||
args.reward_baseline = 0.0
|
||
args.beta_ema = finite_float(getattr(args, "kl_target", 0.08), 0.08)
|
||
args.current_step = 0
|
||
os.makedirs(args.save_dir, exist_ok=True)
|
||
|
||
lm_config = MiniMindConfig(
|
||
hidden_size=args.hidden_size,
|
||
num_hidden_layers=args.num_hidden_layers,
|
||
max_seq_len=max(args.max_seq_len + args.max_gen_len, args.max_total_len, args.max_context_len + args.max_gen_len),
|
||
use_moe=bool(args.use_moe),
|
||
)
|
||
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir="../checkpoints") if args.from_resume == 1 else None
|
||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 if args.dtype == "float16" else torch.float32
|
||
autocast_factory = build_autocast_factory(device_type, dtype)
|
||
|
||
wandb = None
|
||
if args.use_wandb and is_main_process():
|
||
try:
|
||
import swanlab as wandb
|
||
wandb_id = ckp_data.get("wandb_id") if ckp_data else None
|
||
wandb.init(project=args.wandb_project, name=f"ToolUseV4-E{args.epochs}-B{args.batch_size}-G{args.num_generations}", id=wandb_id, resume="must" if wandb_id else None)
|
||
except Exception as exc:
|
||
Logger(f"[WARN] swanlab disabled: {repr(exc)}")
|
||
wandb = None
|
||
|
||
train_on_this_rank = is_trainer_process(args)
|
||
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
|
||
ensure_tokenizer_ids(tokenizer, lm_config)
|
||
ref_model = None
|
||
if train_on_this_rank:
|
||
ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
|
||
ref_model = ref_model.eval().requires_grad_(False)
|
||
rollout_model = None
|
||
if args.rollout_engine == "torch":
|
||
rollout_model, _ = init_model(lm_config, args.from_weight, device=args.device)
|
||
rollout_model = rollout_model.eval().requires_grad_(False)
|
||
|
||
if ckp_data:
|
||
model.load_state_dict(ckp_data["model"], strict=False)
|
||
if ref_model is not None and args.ref_from_rl_checkpoint:
|
||
ref_model.load_state_dict(ckp_data["model"], strict=False)
|
||
Logger("[WARN] ref_model loaded from RL checkpoint because --ref_from_rl_checkpoint=1")
|
||
if rollout_model is not None:
|
||
rollout_model.load_state_dict(ckp_data["model"], strict=False)
|
||
args.reward_baseline = float(ckp_data.get("reward_baseline", 0.0))
|
||
|
||
if args.use_compile == 1 and train_on_this_rank:
|
||
model = torch.compile(model)
|
||
Logger("torch.compile enabled")
|
||
if dist.is_initialized() and not is_trainer_rollout_mode(args):
|
||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||
elif not train_on_this_rank:
|
||
model = model.eval().requires_grad_(False)
|
||
|
||
reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16) if args.use_reward_model and train_on_this_rank else None
|
||
if reward_model is not None:
|
||
Logger(f"Loaded reward model from {args.reward_model_path}")
|
||
|
||
active_tools = parse_tool_subset(args.tool_subset)
|
||
env = RealToolEnv(time_mode=args.time_mode, active_tools=active_tools, allow_calculator_alias=bool(args.augment_calculator_alias))
|
||
train_ds = PromptPoolDataset(
|
||
args.data_path,
|
||
max_samples=args.max_samples,
|
||
filter_mode=args.filter_mode,
|
||
require_gt=bool(args.require_gt),
|
||
allowed_tools=list(env.supported_names()),
|
||
inject_answer_instruction=bool(args.inject_answer_instruction),
|
||
augment_calculator_alias=bool(args.augment_calculator_alias),
|
||
default_tools=env.tool_specs(),
|
||
force_all_tools=bool(args.force_all_tools),
|
||
)
|
||
train_ds, eval_ds = split_prompt_pool(train_ds, args.eval_ratio, args.eval_max_pool, seed=args.seed)
|
||
global_train_samples = len(train_ds)
|
||
world_size = ddp_world_size()
|
||
trainer_rollout_dist = dist.is_initialized() and is_trainer_rollout_mode(args) and world_size > 1
|
||
rollout_ranks = max(1, world_size - 1) if trainer_rollout_dist else max(world_size, 1)
|
||
trainer_produces_rollout = not (trainer_rollout_dist and rank == trainer_rank())
|
||
if use_default_max_version_gap:
|
||
if trainer_rollout_dist:
|
||
args.max_version_gap = max(int(args.max_head_offpolicyness), min(max(2, int(args.policy_sync_interval) * 2), 4))
|
||
else:
|
||
args.max_version_gap = int(args.max_head_offpolicyness)
|
||
if dist.is_initialized() and not trainer_rollout_dist:
|
||
train_ds = shard_prompt_pool(train_ds, world_size, rank)
|
||
args.eval_fixed_indices = build_fixed_eval_indices(eval_ds, args.eval_samples, seed=args.eval_seed)
|
||
if args.rollout_engine == "torch" and world_size > 1 and args.async_workers > 1 and not trainer_rollout_dist:
|
||
if is_main_process():
|
||
Logger(f"[WARN] torch rollout + DDP: async_workers {args.async_workers} -> 1 to avoid per-rank lock contention")
|
||
args.async_workers = 1
|
||
if is_main_process():
|
||
Logger(f"Train samples: global={global_train_samples}, local_rank0={len(train_ds)}, Eval samples: {len(eval_ds)}")
|
||
Logger(f"Dataset stats: {json.dumps(train_ds.stats, ensure_ascii=False)}")
|
||
Logger(f"Active tools: {','.join(supported_tool_names(env.tool_specs()))}")
|
||
|
||
if args.rollout_engine == "torch":
|
||
rollout_engine = StableTorchRolloutEngine(
|
||
rollout_model,
|
||
tokenizer,
|
||
device=args.device,
|
||
autocast_factory=autocast_factory,
|
||
top_p=args.top_p,
|
||
top_k=args.top_k,
|
||
repetition_penalty=args.repetition_penalty,
|
||
use_cache=bool(args.torch_rollout_use_cache),
|
||
)
|
||
else:
|
||
rollout_engine = create_rollout_engine(
|
||
engine_type=args.rollout_engine,
|
||
policy_model=model,
|
||
tokenizer=tokenizer,
|
||
device=args.device,
|
||
autocast_ctx=autocast_factory(),
|
||
sglang_base_url=args.sglang_base_url,
|
||
sglang_model_path=args.sglang_model_path,
|
||
sglang_shared_path=args.sglang_shared_path,
|
||
)
|
||
|
||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, betas=(0.9, 0.95), weight_decay=args.weight_decay) if train_on_this_rank else None
|
||
if args.total_steps > 0:
|
||
if dist.is_initialized() and bool(args.scale_total_steps_by_world_size) and not trainer_rollout_dist:
|
||
total_steps = max(1, math.ceil(args.total_steps / max(ddp_world_size(), 1)))
|
||
if is_main_process():
|
||
Logger(f"[WARN] total_steps scaled by world_size: requested={args.total_steps}, effective={total_steps}")
|
||
else:
|
||
total_steps = args.total_steps
|
||
else:
|
||
total_steps = math.ceil(len(train_ds) / args.batch_size) * args.epochs
|
||
scheduler = CosineAnnealingLR(optimizer, T_max=max(1, math.ceil(total_steps / args.accumulation_steps)), eta_min=args.learning_rate / 10) if train_on_this_rank else None
|
||
start_step = int(ckp_data.get("step", 0)) if ckp_data else 0
|
||
opt_step = int(ckp_data.get("opt_step", 0)) if ckp_data else 0
|
||
if ckp_data and optimizer is not None:
|
||
optimizer.load_state_dict(ckp_data["optimizer"])
|
||
if scheduler is not None and ckp_data.get("scheduler") is not None:
|
||
scheduler.load_state_dict(ckp_data["scheduler"])
|
||
|
||
version = PolicyVersion(int(ckp_data.get("policy_version", 0)) if ckp_data else 0)
|
||
args.current_version = version.get()
|
||
engine_lock = threading.Lock()
|
||
if trainer_rollout_dist:
|
||
sync_rollout_policy_distributed(args, rollout_engine, rollout_model, model, engine_lock, version, src_rank=trainer_rank())
|
||
else:
|
||
sync_rollout_policy(args, rollout_engine, rollout_model, model, engine_lock, version)
|
||
args.current_version = version.get()
|
||
|
||
if args.eval_before_train:
|
||
if dist.is_initialized():
|
||
dist.barrier()
|
||
if is_main_process():
|
||
eval_metrics = run_agent_eval(args, eval_ds, rollout_engine, tokenizer, env, engine_lock, version.get(), reward_model)
|
||
if eval_metrics:
|
||
Logger(
|
||
f"[EVAL_INIT] R={eval_metrics['eval/reward']:.4f}, "
|
||
f"Exact={eval_metrics['eval/exact_match']:.3f}, Cov={eval_metrics['eval/answer_cov']:.3f}, "
|
||
f"Tag={eval_metrics['eval/answer_tag']:.3f}, ToolSucc={eval_metrics['eval/tool_success_rate']:.3f}, "
|
||
f"ToolChoice={eval_metrics['eval/tool_correct_choice_rate']:.3f}, ToolExpr={eval_metrics['eval/tool_expr_order_cov']:.3f}, ToolVal={eval_metrics['eval/tool_value_order_cov']:.3f}, "
|
||
f"Canon={eval_metrics['eval/tool_canonical_rate']:.3f}"
|
||
)
|
||
if wandb:
|
||
wandb.log(eval_metrics, step=start_step)
|
||
if dist.is_initialized():
|
||
dist.barrier()
|
||
|
||
if args.max_queue_groups > 0:
|
||
max_queue_groups = int(args.max_queue_groups)
|
||
elif trainer_rollout_dist:
|
||
max_queue_groups = max((args.max_version_gap + 1) * args.batch_size, args.batch_size * max(8, rollout_ranks * 4))
|
||
else:
|
||
max_queue_groups = (args.max_version_gap + 1) * args.batch_size
|
||
manager_workers = int(args.async_workers) if trainer_produces_rollout else 0
|
||
manager = AsyncRolloutManager(
|
||
train_ds,
|
||
tokenizer,
|
||
rollout_engine,
|
||
env,
|
||
version,
|
||
args.batch_size,
|
||
args.num_generations,
|
||
args.max_turns,
|
||
args.max_gen_len,
|
||
args.max_context_len,
|
||
args.thinking_ratio,
|
||
args.temperature,
|
||
args.device,
|
||
args.max_version_gap,
|
||
max_queue_groups,
|
||
manager_workers,
|
||
engine_lock,
|
||
args.seed,
|
||
)
|
||
stop_requested = {"flag": False}
|
||
|
||
def _sig_handler(signum, frame):
|
||
stop_requested["flag"] = True
|
||
Logger("received stop signal, will save and exit safely")
|
||
|
||
signal.signal(signal.SIGINT, _sig_handler)
|
||
signal.signal(signal.SIGTERM, _sig_handler)
|
||
manager.start()
|
||
global_queue_groups = max_queue_groups * (rollout_ranks if trainer_rollout_dist else 1)
|
||
pending_groups = deque()
|
||
prefill_capacity = global_queue_groups if trainer_rollout_dist else max_queue_groups
|
||
default_global_prefetch_groups = (
|
||
max(args.batch_size, min(global_queue_groups, args.batch_size * max(int(args.max_version_gap) + 1, rollout_ranks)))
|
||
if trainer_rollout_dist
|
||
else args.batch_size
|
||
)
|
||
prefill_default = max(args.batch_size, min(prefill_capacity, default_global_prefetch_groups if trainer_rollout_dist else args.batch_size * 2))
|
||
prefill_groups = min(prefill_capacity, args.prefill_queue_groups if args.prefill_queue_groups > 0 else prefill_default)
|
||
if trainer_rollout_dist:
|
||
prefetched = distributed_fetch_groups(
|
||
rank,
|
||
world_size,
|
||
manager,
|
||
pending_groups,
|
||
prefill_groups,
|
||
fetch_timeout_sec=args.prefill_timeout_sec,
|
||
producer_ranks=rollout_ranks,
|
||
src_rank=trainer_rank(),
|
||
min_return_groups=prefill_groups,
|
||
)
|
||
ready = prefetched is not None if rank == trainer_rank() else bool(prefetched)
|
||
if rank == trainer_rank() and prefetched is not None:
|
||
pending_groups.extend(prefetched)
|
||
else:
|
||
ready = manager.wait_until_ready(prefill_groups, timeout_sec=args.prefill_timeout_sec)
|
||
if is_main_process():
|
||
Logger(f"Queue prefill: target={prefill_groups}, ready={int(bool(ready))}")
|
||
Logger(
|
||
f"ToolUseV4 async rollout started: engine={args.rollout_engine}, mode={args.multi_gpu_mode}, workers={manager_workers}, rollout_ranks={rollout_ranks}, "
|
||
f"max_version_gap={args.max_version_gap}, queue(local/global)={max_queue_groups}/{global_queue_groups}, decoupled={int(args.use_decoupled_ppo)}, "
|
||
f"advantage={args.advantage_mode}, beta={args.beta}, beta_max={args.beta_max}, kl_target={args.kl_target}, "
|
||
f"bc_coef={args.bc_coef}, bc_min={args.bc_min_coef}, loss_reduction={args.loss_reduction}, cache={args.torch_rollout_use_cache}, "
|
||
f"tools={','.join(supported_tool_names(env.tool_specs()))}"
|
||
)
|
||
|
||
if optimizer is not None:
|
||
optimizer.zero_grad(set_to_none=True)
|
||
if train_on_this_rank:
|
||
model.train()
|
||
else:
|
||
model.eval().requires_grad_(False)
|
||
last_metrics = {}
|
||
step = start_step
|
||
try:
|
||
for step in range(start_step + 1, total_steps + 1):
|
||
if trainer_rollout_dist:
|
||
loop_flag = torch.tensor([1 if (rank == trainer_rank() and not stop_requested["flag"]) else 0], device=args.device, dtype=torch.int64)
|
||
dist.broadcast(loop_flag, src=trainer_rank())
|
||
if loop_flag.item() <= 0:
|
||
break
|
||
elif stop_requested["flag"]:
|
||
break
|
||
args.current_step = step
|
||
args.current_version = version.get()
|
||
diagnostics, metrics, groups = {}, {}, None
|
||
sync_needed = 0
|
||
|
||
if trainer_rollout_dist:
|
||
fetch_target = min(global_queue_groups, int(args.global_prefetch_groups) if args.global_prefetch_groups > 0 else default_global_prefetch_groups)
|
||
stale_prefetch_drops = 0
|
||
enough_groups = False
|
||
groups = None
|
||
fetch_wait_rounds = 0
|
||
stop_fetch = False
|
||
while True:
|
||
fetched = distributed_fetch_groups(
|
||
rank,
|
||
world_size,
|
||
manager,
|
||
pending_groups,
|
||
fetch_target,
|
||
fetch_timeout_sec=args.fetch_timeout_sec,
|
||
producer_ranks=rollout_ranks,
|
||
src_rank=trainer_rank(),
|
||
min_return_groups=args.batch_size,
|
||
)
|
||
if rank == trainer_rank():
|
||
groups, stale_drops_now = split_prefetched_groups(fetched, pending_groups, args.batch_size, version.get(), args.max_version_gap)
|
||
stale_prefetch_drops += stale_drops_now
|
||
enough_groups = groups is not None and len(groups) >= args.batch_size
|
||
else:
|
||
enough_groups = False
|
||
stale_drops_now = 0
|
||
enough_tensor = torch.tensor([1 if (rank == trainer_rank() and enough_groups) else 0], device=args.device, dtype=torch.int64)
|
||
dist.broadcast(enough_tensor, src=trainer_rank())
|
||
enough_groups = bool(enough_tensor.item() > 0)
|
||
if enough_groups:
|
||
break
|
||
fetch_ctrl = torch.tensor([1 if (rank == trainer_rank() and stop_requested["flag"]) else 0], device=args.device, dtype=torch.int64)
|
||
dist.broadcast(fetch_ctrl, src=trainer_rank())
|
||
stop_fetch = bool(fetch_ctrl.item() > 0)
|
||
if stop_fetch:
|
||
break
|
||
if rank == trainer_rank() and fetched is None:
|
||
fetch_wait_rounds += 1
|
||
if fetch_wait_rounds == 1 or fetch_wait_rounds % 10 == 0:
|
||
Logger("[WARN] rollout queue lagging behind trainer; continue waiting for enough fresh groups")
|
||
if stop_fetch:
|
||
break
|
||
|
||
if rank == trainer_rank():
|
||
pack_ok = True
|
||
try:
|
||
batch = pack_groups(groups, tokenizer, args.max_total_len, args.device, args.num_generations, bool(args.strict_group_packing))
|
||
except RuntimeError as exc:
|
||
pack_ok = False
|
||
Logger(f"[WARN] pack_groups failed at step={step}: {repr(exc)}")
|
||
batch = None
|
||
|
||
if pack_ok:
|
||
loss, metrics, diagnostics = compute_async_loss(args, batch, model, ref_model, reward_model, tokenizer, autocast_factory, lm_config, env)
|
||
metrics["train/world_size"] = float(world_size)
|
||
metrics["policy/kl_control"] = float(metrics.get(beta_control_metric_key(args), 0.0))
|
||
if not bool(torch.isfinite(loss).item()):
|
||
Logger(f"[WARN] non-finite loss at step={step}; skip update")
|
||
optimizer.zero_grad(set_to_none=True)
|
||
metrics = {}
|
||
else:
|
||
kl_now = float(metrics.get("policy/kl_ref_k3", 0.0))
|
||
kl_control_now = float(metrics.get("policy/kl_control", kl_now))
|
||
if args.hard_kl_stop_mult > 0 and kl_control_now > args.kl_target * args.hard_kl_stop_mult:
|
||
Logger(f"[WARN] KL too high at step={step}: control={kl_control_now:.4f}, ref={kl_now:.4f}; skip update")
|
||
optimizer.zero_grad(set_to_none=True)
|
||
else:
|
||
sync_now = (step % args.accumulation_steps == 0)
|
||
with nullcontext():
|
||
(loss / args.accumulation_steps).backward()
|
||
last_grad_norm = 0.0
|
||
if sync_now:
|
||
updated, last_grad_norm = maybe_step_optimizer(args, model, optimizer, scheduler)
|
||
if updated:
|
||
opt_step += 1
|
||
reward_now = float(metrics.get("reward/mean", 0.0))
|
||
args.reward_baseline = reward_now if opt_step == 1 and abs(args.reward_baseline) < 1e-12 else args.reward_ema_decay * args.reward_baseline + (1 - args.reward_ema_decay) * reward_now
|
||
metrics["train/beta_ema"] = maybe_update_beta(args, metrics, opt_step)
|
||
if opt_step % args.policy_sync_interval == 0:
|
||
sync_needed = 1
|
||
metrics["train/grad_norm"] = float(last_grad_norm)
|
||
metrics["train/beta"] = float(args.beta)
|
||
metrics["train/beta_ema"] = float(getattr(args, "beta_ema", 0.0))
|
||
metrics["train/lr"] = float(optimizer.param_groups[0]["lr"])
|
||
metrics["policy/version"] = float(version.get())
|
||
metrics["policy/max_version_gap"] = float(args.max_version_gap)
|
||
metrics["async/prefetch_stale_drop"] = float(stale_prefetch_drops)
|
||
if args.rollout_engine == "torch":
|
||
metrics["rollout/sanitized_logits_batches"] = float(getattr(rollout_engine, "nan_batches", 0))
|
||
metrics["rollout/bad_prob_batches"] = float(getattr(rollout_engine, "bad_prob_batches", 0))
|
||
last_metrics = metrics
|
||
|
||
sync_tensor = torch.tensor([int(sync_needed)], device=args.device, dtype=torch.int64)
|
||
dist.broadcast(sync_tensor, src=trainer_rank())
|
||
if int(sync_tensor.item()) > 0:
|
||
args.current_version = sync_rollout_policy_distributed(args, rollout_engine, rollout_model, model, engine_lock, version, src_rank=trainer_rank())
|
||
else:
|
||
args.current_version = version.get()
|
||
else:
|
||
groups = manager.get_batch(timeout=None)
|
||
enough_groups = ddp_all_true(len(groups) >= args.batch_size, args.device)
|
||
if not enough_groups:
|
||
Logger("Not enough rollout groups; stopping early.")
|
||
break
|
||
pack_ok = True
|
||
try:
|
||
batch = pack_groups(groups, tokenizer, args.max_total_len, args.device, args.num_generations, bool(args.strict_group_packing))
|
||
except RuntimeError as exc:
|
||
pack_ok = False
|
||
Logger(f"[WARN] pack_groups failed at step={step}: {repr(exc)}")
|
||
batch = None
|
||
if not ddp_all_true(pack_ok, args.device):
|
||
optimizer.zero_grad(set_to_none=True)
|
||
continue
|
||
loss, metrics, diagnostics = compute_async_loss(args, batch, model, ref_model, reward_model, tokenizer, autocast_factory, lm_config, env)
|
||
metrics["train/world_size"] = float(ddp_world_size())
|
||
metrics["policy/kl_control"] = float(metrics.get(beta_control_metric_key(args), 0.0))
|
||
if not ddp_all_true(bool(torch.isfinite(loss).item()), args.device):
|
||
Logger(f"[WARN] non-finite loss at step={step}; skip update")
|
||
optimizer.zero_grad(set_to_none=True)
|
||
continue
|
||
kl_now = ddp_mean_scalar(float(metrics.get("policy/kl_ref_k3", 0.0)), args.device)
|
||
kl_control_now = ddp_mean_scalar(float(metrics.get("policy/kl_control", kl_now)), args.device)
|
||
if args.hard_kl_stop_mult > 0 and kl_control_now > args.kl_target * args.hard_kl_stop_mult:
|
||
if is_main_process():
|
||
Logger(f"[WARN] KL too high at step={step}: control={kl_control_now:.4f}, ref={kl_now:.4f}; skip update")
|
||
optimizer.zero_grad(set_to_none=True)
|
||
continue
|
||
sync_now = (step % args.accumulation_steps == 0)
|
||
backward_ctx = model.no_sync() if isinstance(model, DistributedDataParallel) and not sync_now else nullcontext()
|
||
with backward_ctx:
|
||
(loss / args.accumulation_steps).backward()
|
||
last_grad_norm = 0.0
|
||
if sync_now:
|
||
updated, last_grad_norm = maybe_step_optimizer(args, model, optimizer, scheduler)
|
||
if updated:
|
||
opt_step += 1
|
||
reward_now = ddp_mean_scalar(float(metrics.get("reward/mean", 0.0)), args.device)
|
||
args.reward_baseline = reward_now if opt_step == 1 and abs(args.reward_baseline) < 1e-12 else args.reward_ema_decay * args.reward_baseline + (1 - args.reward_ema_decay) * reward_now
|
||
metrics["train/beta_ema"] = maybe_update_beta(args, metrics, opt_step)
|
||
if opt_step % args.policy_sync_interval == 0:
|
||
args.current_version = sync_rollout_policy(args, rollout_engine, rollout_model, model, engine_lock, version)
|
||
|
||
metrics["train/grad_norm"] = float(last_grad_norm)
|
||
metrics["train/beta"] = float(args.beta)
|
||
metrics["train/beta_ema"] = float(getattr(args, "beta_ema", 0.0))
|
||
metrics["train/lr"] = float(optimizer.param_groups[0]["lr"])
|
||
metrics["policy/version"] = float(version.get())
|
||
metrics["policy/max_version_gap"] = float(args.max_version_gap)
|
||
if args.rollout_engine == "torch":
|
||
metrics["rollout/sanitized_logits_batches"] = float(getattr(rollout_engine, "nan_batches", 0))
|
||
metrics["rollout/bad_prob_batches"] = float(getattr(rollout_engine, "bad_prob_batches", 0))
|
||
last_metrics = metrics
|
||
|
||
if step % args.log_interval == 0 or step == total_steps:
|
||
if trainer_rollout_dist:
|
||
local_stats = manager.snapshot_stats()
|
||
stats = ddp_average_metrics(local_stats, args.device)
|
||
local_queued_total = ddp_sum_scalar(float(local_stats.get("queued", 0.0)), args.device)
|
||
metrics_to_log = dict(last_metrics) if rank == trainer_rank() and last_metrics else {}
|
||
if args.rollout_engine == "torch":
|
||
cache_prefill = ddp_mean_scalar(float(getattr(rollout_engine, "cache_prefill_tokens", 0.0)), args.device)
|
||
cache_decode = ddp_mean_scalar(float(getattr(rollout_engine, "cache_decode_tokens", 0.0)), args.device)
|
||
cache_fallbacks = ddp_mean_scalar(float(getattr(rollout_engine, "cache_fallbacks", 0.0)), args.device)
|
||
if metrics_to_log:
|
||
metrics_to_log["rollout/cache_prefill_tokens"] = cache_prefill
|
||
metrics_to_log["rollout/cache_decode_tokens"] = cache_decode
|
||
metrics_to_log["rollout/cache_decode_ratio"] = cache_decode / max(cache_prefill + cache_decode, 1.0)
|
||
metrics_to_log["rollout/cache_fallbacks"] = cache_fallbacks
|
||
if rank == trainer_rank() and metrics_to_log:
|
||
stats["local_queued_total"] = float(local_queued_total)
|
||
stats["global_pending_groups"] = float(len(pending_groups))
|
||
stats["global_queue_capacity"] = float(global_queue_groups)
|
||
stats["global_available_groups"] = stats["local_queued_total"] + stats["global_pending_groups"]
|
||
stats["global_queue_utilization"] = stats["global_available_groups"] / max(stats["global_queue_capacity"], 1.0)
|
||
stats["rollout_ranks"] = float(rollout_ranks)
|
||
last_metrics = metrics_to_log
|
||
Logger(
|
||
f"Step:[{step}/{total_steps}] V:{version.get()} "
|
||
f"R:{metrics_to_log['reward/mean']:.3f}({metrics_to_log['reward/min']:.2f}/{metrics_to_log['reward/max']:.2f}) "
|
||
f"Acc(E/C):{metrics_to_log.get('reward_detail/exact_match', 0.0):.2f}/{metrics_to_log.get('reward_detail/answer_cov', 0.0):.2f} "
|
||
f"Fmt(T/ET/Mal):{metrics_to_log.get('reward_detail/answer_tag', 0.0):.2f}/{metrics_to_log.get('reward_detail/exact_answer_tag', 0.0):.2f}/{metrics_to_log.get('reward_detail/malformed_tags', 0.0):.2f} "
|
||
f"Tool(S/Cov/Expr/Val/Choice/Pref/Canon):{metrics_to_log.get('reward_detail/tool_success_rate', 0.0):.2f}/{metrics_to_log.get('reward_detail/math_call_coverage', 0.0):.2f}/{metrics_to_log.get('reward_detail/tool_expr_order_cov', 0.0):.2f}/{metrics_to_log.get('reward_detail/tool_value_order_cov', 0.0):.2f}/{metrics_to_log.get('reward_detail/correct_tool_choice', 0.0):.2f}/{metrics_to_log.get('reward_detail/preferred_math_rate', 0.0):.2f}/{metrics_to_log.get('reward_detail/canonical_tool_rate', 0.0):.2f} "
|
||
f"Num(+/-/Ord):{metrics_to_log.get('reward_detail/extra_numbers', 0.0):.1f}/{metrics_to_log.get('reward_detail/missing_numbers', 0.0):.1f}/{metrics_to_log.get('reward_detail/order_ok', 0.0):.2f} "
|
||
f"PPO(KL/Ent/Clip):{metrics_to_log['policy/kl_ref_k3']:.3f}/{metrics_to_log['policy/entropy']:.2f}/{metrics_to_log['policy/ratio_clip_frac']:.3f} "
|
||
f"BC:{metrics_to_log.get('bc/loss', 0.0):.3f}*{metrics_to_log.get('bc/coef', 0.0):.3f} "
|
||
f"Q(local/global):{stats['local_queued_total']:.0f}/{stats['global_available_groups']:.0f} Roll:{stats['last_rollout_time']:.2f}s Wait:{stats['last_train_wait']:.2f}s P90:{stats.get('train_wait_sec_p90', 0.0):.2f}s "
|
||
f"Tok:{stats['avg_tokens_per_traj']:.1f} LR:{optimizer.param_groups[0]['lr']:.2e} Beta:{args.beta:.3f}"
|
||
)
|
||
if wandb:
|
||
payload = {
|
||
**{k: v for k, v in metrics_to_log.items() if isinstance(v, (int, float)) and math.isfinite(float(v))},
|
||
**{f"async/{k}": v for k, v in stats.items() if isinstance(v, (int, float)) and math.isfinite(float(v))},
|
||
}
|
||
wandb.log(payload, step=step)
|
||
else:
|
||
metrics = ddp_average_metrics(metrics, args.device)
|
||
if args.rollout_engine == "torch":
|
||
cache_prefill = ddp_mean_scalar(float(getattr(rollout_engine, "cache_prefill_tokens", 0.0)), args.device)
|
||
cache_decode = ddp_mean_scalar(float(getattr(rollout_engine, "cache_decode_tokens", 0.0)), args.device)
|
||
cache_fallbacks = ddp_mean_scalar(float(getattr(rollout_engine, "cache_fallbacks", 0.0)), args.device)
|
||
metrics["rollout/cache_prefill_tokens"] = cache_prefill
|
||
metrics["rollout/cache_decode_tokens"] = cache_decode
|
||
metrics["rollout/cache_decode_ratio"] = cache_decode / max(cache_prefill + cache_decode, 1.0)
|
||
metrics["rollout/cache_fallbacks"] = cache_fallbacks
|
||
stats = ddp_average_metrics(manager.snapshot_stats(), args.device)
|
||
last_metrics = metrics
|
||
Logger(
|
||
f"Step:[{step}/{total_steps}] V:{version.get()} "
|
||
f"R:{metrics['reward/mean']:.3f}({metrics['reward/min']:.2f}/{metrics['reward/max']:.2f}) "
|
||
f"Acc(E/C):{metrics.get('reward_detail/exact_match', 0.0):.2f}/{metrics.get('reward_detail/answer_cov', 0.0):.2f} "
|
||
f"Fmt(T/ET/Mal):{metrics.get('reward_detail/answer_tag', 0.0):.2f}/{metrics.get('reward_detail/exact_answer_tag', 0.0):.2f}/{metrics.get('reward_detail/malformed_tags', 0.0):.2f} "
|
||
f"Tool(S/Cov/Expr/Val/Choice/Pref/Canon):{metrics.get('reward_detail/tool_success_rate', 0.0):.2f}/{metrics.get('reward_detail/math_call_coverage', 0.0):.2f}/{metrics.get('reward_detail/tool_expr_order_cov', 0.0):.2f}/{metrics.get('reward_detail/tool_value_order_cov', 0.0):.2f}/{metrics.get('reward_detail/correct_tool_choice', 0.0):.2f}/{metrics.get('reward_detail/preferred_math_rate', 0.0):.2f}/{metrics.get('reward_detail/canonical_tool_rate', 0.0):.2f} "
|
||
f"Num(+/-/Ord):{metrics.get('reward_detail/extra_numbers', 0.0):.1f}/{metrics.get('reward_detail/missing_numbers', 0.0):.1f}/{metrics.get('reward_detail/order_ok', 0.0):.2f} "
|
||
f"PPO(KL/Ent/Clip):{metrics['policy/kl_ref_k3']:.3f}/{metrics['policy/entropy']:.2f}/{metrics['policy/ratio_clip_frac']:.3f} "
|
||
f"BC:{metrics.get('bc/loss', 0.0):.3f}*{metrics.get('bc/coef', 0.0):.3f} "
|
||
f"Q:{stats['queued']:.0f}/{stats['queue_capacity']:.0f} Roll:{stats['last_rollout_time']:.2f}s Wait:{stats['last_train_wait']:.2f}s P90:{stats.get('train_wait_sec_p90', 0.0):.2f}s "
|
||
f"Tok:{stats['avg_tokens_per_traj']:.1f} LR:{optimizer.param_groups[0]['lr']:.2e} Beta:{args.beta:.3f}"
|
||
)
|
||
if wandb and is_main_process():
|
||
payload = {
|
||
**{k: v for k, v in metrics.items() if isinstance(v, (int, float)) and math.isfinite(float(v))},
|
||
**{f"async/{k}": v for k, v in stats.items() if isinstance(v, (int, float)) and math.isfinite(float(v))},
|
||
}
|
||
wandb.log(payload, step=step)
|
||
|
||
if not trainer_rollout_dist or rank == trainer_rank():
|
||
maybe_log_debug_sample(step, groups, diagnostics, args.debug_interval)
|
||
if args.eval_interval > 0 and (step % args.eval_interval == 0 or step == total_steps):
|
||
if dist.is_initialized():
|
||
dist.barrier()
|
||
if is_main_process():
|
||
eval_metrics = run_agent_eval(args, eval_ds, rollout_engine, tokenizer, env, engine_lock, version.get(), reward_model)
|
||
if eval_metrics:
|
||
Logger(
|
||
f"[EVAL] step={step}, R={eval_metrics['eval/reward']:.4f}, "
|
||
f"Exact={eval_metrics['eval/exact_match']:.3f}, Cov={eval_metrics['eval/answer_cov']:.3f}, "
|
||
f"Tag={eval_metrics['eval/answer_tag']:.3f}, ToolSucc={eval_metrics['eval/tool_success_rate']:.3f}, "
|
||
f"ToolChoice={eval_metrics['eval/tool_correct_choice_rate']:.3f}, ToolExpr={eval_metrics['eval/tool_expr_order_cov']:.3f}, ToolVal={eval_metrics['eval/tool_value_order_cov']:.3f}, "
|
||
f"Canon={eval_metrics['eval/tool_canonical_rate']:.3f}"
|
||
)
|
||
if wandb:
|
||
wandb.log(eval_metrics, step=step)
|
||
if dist.is_initialized():
|
||
dist.barrier()
|
||
if args.save_interval > 0 and step % args.save_interval == 0:
|
||
save_checkpoint(args, lm_config, model, optimizer, scheduler, epoch=0, step=step, opt_step=opt_step, version=version.get(), wandb=wandb)
|
||
if step % 20 == 0:
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
save_checkpoint(args, lm_config, model, optimizer, scheduler, epoch=0, step=step, opt_step=opt_step, version=version.get(), wandb=wandb)
|
||
if last_metrics and is_main_process():
|
||
Logger(f"Training finished. final_reward={last_metrics.get('reward/mean', 0.0):.4f}")
|
||
finally:
|
||
manager.stop()
|
||
if wandb and is_main_process():
|
||
try:
|
||
wandb.finish()
|
||
except Exception:
|
||
pass
|
||
if dist.is_initialized():
|
||
dist.destroy_process_group()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|