159 lines
7.1 KiB
Python
159 lines
7.1 KiB
Python
"""Student-owned tool trajectories and token-level on-policy distillation."""
|
|
|
|
import json
|
|
from contextlib import nullcontext
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from trainer.train_agent import execute_tool, parse_tool_calls
|
|
|
|
|
|
@dataclass
|
|
class AgentTrajectory:
|
|
input_ids: list
|
|
assistant_mask: list
|
|
turns: int
|
|
tool_calls: int
|
|
truncated: bool
|
|
|
|
|
|
def tool_observation(tokenizer, calls, tools, open_thinking=False):
|
|
"""Render only the environment response and the next assistant prefix."""
|
|
marker = "<|agent_opd_observation_boundary|>"
|
|
messages = [{"role": "assistant", "content": marker}]
|
|
allowed = {tool["function"]["name"] for tool in (tools or [])}
|
|
for call in calls:
|
|
name, arguments = call.get("name", ""), call.get("arguments", {})
|
|
if isinstance(arguments, str):
|
|
try:
|
|
arguments = json.loads(arguments)
|
|
except json.JSONDecodeError:
|
|
arguments = None
|
|
result = None
|
|
if isinstance(name, str) and name in allowed and isinstance(arguments, dict):
|
|
result = execute_tool(name, arguments)
|
|
if result is None:
|
|
result = {"error": "unknown tool or invalid arguments"}
|
|
messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)[:2048]})
|
|
rendered = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True,
|
|
tools=tools, open_thinking=open_thinking,
|
|
)
|
|
_, found, observation = rendered.partition(marker)
|
|
if not found:
|
|
raise ValueError("The chat template must preserve assistant content")
|
|
return tokenizer(observation, add_special_tokens=False)["input_ids"]
|
|
|
|
|
|
@torch.no_grad()
|
|
def collect_trajectory(student, tokenizer, messages, tools=None, *, max_turns=3,
|
|
max_new_tokens=256, max_total_len=2048, temperature=1.0,
|
|
open_thinking=False, autocast_ctx=None):
|
|
"""Sample each assistant turn, execute local tools, and preserve sampled IDs.
|
|
|
|
Only student-generated tokens (including each turn's first EOS) are targets.
|
|
Environment observations remain visible context with no direct loss.
|
|
"""
|
|
if min(max_turns, max_new_tokens, max_total_len) < 1 or temperature <= 0:
|
|
raise ValueError("Generation limits and temperature must be positive")
|
|
prompt = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True,
|
|
tools=tools, open_thinking=open_thinking,
|
|
)
|
|
ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
|
|
if not ids or len(ids) >= max_total_len:
|
|
raise ValueError("Prompt leaves no generation space; increase --max_total_len")
|
|
mask = [False] * len(ids)
|
|
model = getattr(student, "module", student)
|
|
device = next(model.parameters()).device
|
|
was_training = model.training
|
|
model.eval()
|
|
turns, call_count, truncated = 0, 0, False
|
|
try:
|
|
for turn in range(max_turns):
|
|
inputs = torch.tensor([ids], device=device)
|
|
with autocast_ctx if autocast_ctx is not None else nullcontext():
|
|
output = model.generate(
|
|
input_ids=inputs, attention_mask=torch.ones_like(inputs),
|
|
max_new_tokens=min(max_new_tokens, max_total_len - len(ids)),
|
|
do_sample=True, temperature=temperature, top_p=1.0, top_k=0,
|
|
repetition_penalty=1.0, num_beams=1, use_cache=True,
|
|
eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
new_ids = output[0, len(ids):].tolist()
|
|
if tokenizer.eos_token_id in new_ids:
|
|
new_ids = new_ids[:new_ids.index(tokenizer.eos_token_id) + 1]
|
|
if not new_ids:
|
|
truncated = True
|
|
break
|
|
ids.extend(new_ids)
|
|
mask.extend([True] * len(new_ids))
|
|
turns += 1
|
|
if new_ids[-1] != tokenizer.eos_token_id:
|
|
truncated = True
|
|
break
|
|
text = tokenizer.decode(new_ids, skip_special_tokens=True)
|
|
calls = [call for call in parse_tool_calls(text) if isinstance(call, dict)]
|
|
call_count += len(calls)
|
|
if not calls:
|
|
break
|
|
if turn + 1 == max_turns:
|
|
truncated = True
|
|
break
|
|
observation = tool_observation(tokenizer, calls, tools, open_thinking)
|
|
if observation[:1] == [tokenizer.eos_token_id]:
|
|
observation = observation[1:] # EOS is already in the sampled IDs.
|
|
if len(ids) + len(observation) >= max_total_len:
|
|
truncated = True
|
|
break
|
|
ids.extend(observation)
|
|
mask.extend([False] * len(observation))
|
|
finally:
|
|
model.train(was_training)
|
|
return AgentTrajectory(ids, mask, turns, call_count, truncated)
|
|
|
|
|
|
def pack_trajectories(trajectories, pad_token_id, device):
|
|
if not trajectories:
|
|
raise ValueError("At least one trajectory is required")
|
|
if any(len(t.input_ids) != len(t.assistant_mask) for t in trajectories):
|
|
raise ValueError("Trajectory tokens and masks must have the same length")
|
|
length = max(len(t.input_ids) for t in trajectories)
|
|
ids, attention, targets = [], [], []
|
|
for trajectory in trajectories:
|
|
size = len(trajectory.input_ids)
|
|
ids.append(trajectory.input_ids + [pad_token_id] * (length - size))
|
|
attention.append([1] * size + [0] * (length - size))
|
|
targets.append(trajectory.assistant_mask + [False] * (length - size))
|
|
return (
|
|
torch.tensor(ids, dtype=torch.long, device=device),
|
|
torch.tensor(attention, dtype=torch.long, device=device),
|
|
torch.tensor(targets, dtype=torch.bool, device=device)[:, 1:],
|
|
)
|
|
|
|
|
|
def distillation_loss(student_logits, teacher_logits, target_mask,
|
|
temperature=1.0, direction="forward_kl"):
|
|
"""Full-vocabulary KL averaged over assistant tokens on sampled prefixes.
|
|
|
|
Sampling is detached (the GKD semi-gradient); this is not a REINFORCE
|
|
estimator of trajectory-level reverse KL. Teacher gradients are blocked.
|
|
"""
|
|
if student_logits.shape != teacher_logits.shape or student_logits.ndim != 3:
|
|
raise ValueError("Student and teacher logits must share [batch, time, vocab]")
|
|
if target_mask.shape != student_logits.shape[:2]:
|
|
raise ValueError("Target mask must align with the shifted logits")
|
|
if temperature <= 0 or direction not in {"forward_kl", "reverse_kl"}:
|
|
raise ValueError("Invalid distillation temperature or KL direction")
|
|
if not target_mask.any():
|
|
return student_logits[:, :0, :].sum()
|
|
student_logps = F.log_softmax(student_logits[target_mask].float() / temperature, dim=-1)
|
|
teacher_logps = F.log_softmax(teacher_logits.detach()[target_mask].float() / temperature, dim=-1)
|
|
if direction == "forward_kl":
|
|
loss = F.kl_div(student_logps, teacher_logps, log_target=True, reduction="batchmean")
|
|
else:
|
|
loss = F.kl_div(teacher_logps, student_logps, log_target=True, reduction="batchmean")
|
|
return loss * temperature ** 2
|