From 3233d9872cdf81e0ab045f65f0383c37acb90f76 Mon Sep 17 00:00:00 2001
From: baihaochen <1291308544@qq.com>
Date: Fri, 12 Jun 2026 15:11:31 +0800
Subject: [PATCH] agentic
---
agentic/__init__.py | 21 +
agentic/data_analysis_env.py | 566 +++++++++++++++++++++++++
dataset/agentic_dataset.py | 40 ++
docs/agentic_data_analysis_project.md | 183 ++++++++
model/chat_template.jinja | 114 ++---
requirements.txt | 1 +
scripts/agentic_data_demo.py | 97 +++++
scripts/build_agentic_data_tasks.py | 488 ++++++++++++++++++++++
scripts/build_agentic_sft.py | 86 ++++
scripts/eval_agentic_data.py | 237 +++++++++++
scripts/run_general_sft_300m.sh | 78 ++++
scripts/validate_agentic_data.py | 69 +++
trainer/train_agent_deepspeed.py | 576 ++++++++++++++++++++++++++
trainer/train_full_sft_deepspeed.py | 38 +-
14 files changed, 2504 insertions(+), 90 deletions(-)
create mode 100644 agentic/__init__.py
create mode 100644 agentic/data_analysis_env.py
create mode 100644 dataset/agentic_dataset.py
create mode 100644 docs/agentic_data_analysis_project.md
create mode 100644 scripts/agentic_data_demo.py
create mode 100644 scripts/build_agentic_data_tasks.py
create mode 100644 scripts/build_agentic_sft.py
create mode 100644 scripts/eval_agentic_data.py
create mode 100755 scripts/run_general_sft_300m.sh
create mode 100644 scripts/validate_agentic_data.py
create mode 100644 trainer/train_agent_deepspeed.py
diff --git a/agentic/__init__.py b/agentic/__init__.py
new file mode 100644
index 0000000..7ddacba
--- /dev/null
+++ b/agentic/__init__.py
@@ -0,0 +1,21 @@
+from .data_analysis_env import (
+ AGENTIC_SYSTEM_PROMPT,
+ AgenticToolEnv,
+ average_agentic_metrics,
+ extract_final_answer,
+ format_agentic_user_prompt,
+ get_agentic_tools,
+ parse_tool_calls,
+ score_agentic_trajectory,
+)
+
+__all__ = [
+ "AGENTIC_SYSTEM_PROMPT",
+ "AgenticToolEnv",
+ "average_agentic_metrics",
+ "extract_final_answer",
+ "format_agentic_user_prompt",
+ "get_agentic_tools",
+ "parse_tool_calls",
+ "score_agentic_trajectory",
+]
diff --git a/agentic/data_analysis_env.py b/agentic/data_analysis_env.py
new file mode 100644
index 0000000..53fbb22
--- /dev/null
+++ b/agentic/data_analysis_env.py
@@ -0,0 +1,566 @@
+import ast
+import contextlib
+import io
+import json
+import math
+import multiprocessing as mp
+import os
+import re
+import statistics
+from collections import Counter
+from types import SimpleNamespace
+
+import numpy as np
+import pandas as pd
+
+
+AGENTIC_SYSTEM_PROMPT = (
+ "你是一个面向运营数据分析的 Agent。你需要根据用户问题选择合适工具，"
+ "优先检索指标口径，再用 Python 对给定 CSV 数据进行计算；最终答案必须简洁，"
+ "包含关键结论和必要数值。工具调用必须使用 包裹的 JSON。"
+)
+
+
+AGENTIC_TOOL_SCHEMAS = [
+ {
+ "type": "function",
+ "function": {
+ "name": "retriever",
+ "description": "检索指标口径、表字段说明或业务规则，返回最相关的知识片段。",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "query": {"type": "string", "description": "检索查询"},
+ "top_k": {"type": "integer", "description": "返回片段数量", "default": 3},
+ },
+ "required": ["query"],
+ },
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "python_executor",
+ "description": "在受限 Python 环境中使用 pandas/numpy 读取给定 CSV 并完成数据分析。",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "code": {
+ "type": "string",
+ "description": "Python 代码。把最终结果赋值给 result 或 answer。",
+ }
+ },
+ "required": ["code"],
+ },
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "calculator",
+ "description": "执行确定性数学表达式计算，适合简单四则运算、比例和幂运算。",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "expression": {"type": "string", "description": "数学表达式，如 1280 * 0.15 * 3.2"}
+ },
+ "required": ["expression"],
+ },
+ },
+ },
+]
+
+
+def get_agentic_tools(names=None):
+ if names is None:
+ names = ["retriever", "python_executor", "calculator"]
+ name_set = set(names)
+ return [tool for tool in AGENTIC_TOOL_SCHEMAS if tool["function"]["name"] in name_set]
+
+
+def json_dumps(data):
+ return json.dumps(data, ensure_ascii=False, default=str)
+
+
+def format_agentic_user_prompt(sample):
+ table_lines = []
+ for table in sample.get("tables", []):
+ table_lines.append(
+ f"- {table.get('name', '')}: path={table.get('path', '')}; {table.get('description', '')}"
+ )
+ tool_names = ", ".join(sample.get("tools", ["retriever", "python_executor", "calculator"]))
+ return (
+ f"用户问题：{sample.get('question', '')}\n\n"
+ f"可用数据表：\n" + "\n".join(table_lines) + "\n\n"
+ f"可用工具：{tool_names}\n"
+ "要求：先确认指标口径，再调用工具计算；最终答案只输出结论，不要编造数据。"
+ )
+
+
+def parse_tool_calls(text):
+ calls = []
+ for match in re.findall(r"(.*?)", text or "", re.DOTALL):
+ raw = match.strip()
+ try:
+ data = json.loads(raw)
+ args = data.get("arguments", {})
+ if isinstance(args, str):
+ try:
+ args = json.loads(args)
+ except Exception:
+ args = {}
+ calls.append(
+ {
+ "name": data.get("name", ""),
+ "arguments": args if isinstance(args, dict) else {},
+ "raw": raw,
+ "valid_json": True,
+ }
+ )
+ except Exception:
+ calls.append({"name": "", "arguments": {}, "raw": raw, "valid_json": False})
+ return calls
+
+
+def strip_thinking(text):
+ text = text or ""
+ return re.sub(r".*?", "", text, flags=re.DOTALL).strip()
+
+
+def extract_final_answer(turn_outputs):
+ if isinstance(turn_outputs, str):
+ turn_outputs = [turn_outputs]
+ if not turn_outputs:
+ return ""
+ text = strip_thinking(turn_outputs[-1])
+ if "" in text:
+ text = text.split("")[-1]
+ text = re.sub(r".*?", "", text, flags=re.DOTALL)
+ return text.strip()
+
+
+def _tokenize(text):
+ return re.findall(r"[\w\u4e00-\u9fff]+", str(text).lower())
+
+
+def _overlap_score(query, text):
+ q = Counter(_tokenize(query))
+ d = Counter(_tokenize(text))
+ return sum(min(q[token], d[token]) for token in q)
+
+
+def _safe_number(value, default=0.0):
+ try:
+ return float(value)
+ except Exception:
+ return default
+
+
+class _SafeArithmetic(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_funcs = {
+ "abs": abs,
+ "round": round,
+ "sqrt": math.sqrt,
+ "log": math.log,
+ "log10": math.log10,
+ "ceil": math.ceil,
+ "floor": math.floor,
+ "pow": pow,
+ }
+
+ 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("calculator only supports numeric constants")
+
+ def visit_BinOp(self, node):
+ op = type(node.op)
+ if op not in self.allowed_binops:
+ raise ValueError(f"operator not allowed: {op.__name__}")
+ return self.allowed_binops[op](self.visit(node.left), self.visit(node.right))
+
+ def visit_UnaryOp(self, node):
+ op = type(node.op)
+ if op not in self.allowed_unary:
+ raise ValueError(f"operator not allowed: {op.__name__}")
+ return self.allowed_unary[op](self.visit(node.operand))
+
+ def visit_Call(self, node):
+ if not isinstance(node.func, ast.Name) or node.func.id not in self.allowed_funcs:
+ raise ValueError("function not allowed")
+ args = [self.visit(arg) for arg in node.args]
+ return self.allowed_funcs[node.func.id](*args)
+
+ def generic_visit(self, node):
+ raise ValueError(f"node not allowed: {type(node).__name__}")
+
+
+def safe_calculate(expression):
+ normalized = (
+ str(expression)
+ .replace("^", "**")
+ .replace("×", "*")
+ .replace("÷", "/")
+ .replace("−", "-")
+ .replace("（", "(")
+ .replace("）", ")")
+ )
+ tree = ast.parse(normalized, mode="eval")
+ return _SafeArithmetic().visit(tree)
+
+
+def _jsonable(value):
+ if isinstance(value, pd.DataFrame):
+ return value.to_dict("records")
+ if isinstance(value, pd.Series):
+ return value.to_dict()
+ if isinstance(value, np.generic):
+ return value.item()
+ if isinstance(value, dict):
+ return {str(k): _jsonable(v) for k, v in value.items()}
+ if isinstance(value, (list, tuple, set)):
+ return [_jsonable(v) for v in value]
+ return value
+
+
+def _validate_python_ast(code):
+ banned_nodes = (ast.Import, ast.ImportFrom, ast.Global, ast.Nonlocal, ast.Lambda, ast.ClassDef)
+ banned_names = {
+ "__import__",
+ "eval",
+ "exec",
+ "compile",
+ "open",
+ "input",
+ "globals",
+ "locals",
+ "vars",
+ "dir",
+ "getattr",
+ "setattr",
+ "delattr",
+ "os",
+ "sys",
+ "subprocess",
+ "socket",
+ "requests",
+ "urllib",
+ "pathlib",
+ "shutil",
+ }
+ banned_attrs = {
+ "to_csv",
+ "to_excel",
+ "to_pickle",
+ "to_parquet",
+ "write",
+ "remove",
+ "unlink",
+ "rmdir",
+ "mkdir",
+ "rename",
+ "replace",
+ "system",
+ "popen",
+ }
+ tree = ast.parse(code, mode="exec")
+ for node in ast.walk(tree):
+ if isinstance(node, banned_nodes):
+ raise ValueError(f"python_executor does not allow {type(node).__name__}")
+ if isinstance(node, ast.Name) and node.id in banned_names:
+ raise ValueError(f"name not allowed: {node.id}")
+ if isinstance(node, ast.Attribute) and node.attr in banned_attrs:
+ raise ValueError(f"attribute not allowed: {node.attr}")
+ return tree
+
+
+class _PandasProxy:
+ def __init__(self, repo_root, allowed_paths):
+ self.repo_root = repo_root
+ self.allowed_paths = {os.path.abspath(path) for path in allowed_paths}
+
+ def read_csv(self, path, *args, **kwargs):
+ candidates = []
+ raw = str(path)
+ if os.path.isabs(raw):
+ candidates.append(os.path.abspath(raw))
+ candidates.append(os.path.abspath(raw))
+ candidates.append(os.path.abspath(os.path.join(self.repo_root, raw)))
+ allowed_by_basename = {os.path.basename(path): path for path in self.allowed_paths}
+ if raw in allowed_by_basename:
+ candidates.append(allowed_by_basename[raw])
+ for candidate in candidates:
+ if candidate in self.allowed_paths:
+ return pd.read_csv(candidate, *args, **kwargs)
+ raise PermissionError(f"read_csv path is not in sample tables: {raw}")
+
+ def __getattr__(self, name):
+ return getattr(pd, name)
+
+
+def _python_worker(code, repo_root, allowed_paths, queue):
+ try:
+ tree = _validate_python_ast(code)
+ output = io.StringIO()
+ safe_builtins = {
+ "abs": abs,
+ "all": all,
+ "any": any,
+ "bool": bool,
+ "dict": dict,
+ "enumerate": enumerate,
+ "float": float,
+ "int": int,
+ "len": len,
+ "list": list,
+ "max": max,
+ "min": min,
+ "pow": pow,
+ "print": print,
+ "range": range,
+ "round": round,
+ "set": set,
+ "sorted": sorted,
+ "str": str,
+ "sum": sum,
+ "tuple": tuple,
+ "zip": zip,
+ }
+ local_vars = {}
+ global_vars = {
+ "__builtins__": safe_builtins,
+ "pd": _PandasProxy(repo_root, allowed_paths),
+ "np": np,
+ "math": math,
+ "json": json,
+ "statistics": statistics,
+ }
+ with contextlib.redirect_stdout(output):
+ exec(compile(tree, "", "exec"), global_vars, local_vars)
+ value = local_vars.get("result", local_vars.get("answer", output.getvalue().strip()))
+ queue.put({"ok": True, "result": _jsonable(value), "error": "", "artifacts": []})
+ except Exception as exc:
+ queue.put({"ok": False, "result": "", "error": str(exc)[:500], "artifacts": []})
+
+
+class AgenticToolEnv:
+ def __init__(self, sample, repo_root=None, timeout=3):
+ self.sample = sample
+ self.repo_root = os.path.abspath(repo_root or os.getcwd())
+ self.timeout = timeout
+ self.allowed_paths = [self._resolve_table_path(table.get("path", "")) for table in sample.get("tables", [])]
+
+ def _resolve_table_path(self, path):
+ if os.path.isabs(path):
+ return os.path.abspath(path)
+ return os.path.abspath(os.path.join(self.repo_root, path))
+
+ def execute(self, name, arguments=None):
+ arguments = arguments or {}
+ if name == "retriever":
+ return self._retriever(arguments)
+ if name == "calculator":
+ return self._calculator(arguments)
+ if name == "python_executor":
+ return self._python_executor(arguments)
+ return {"ok": False, "result": "", "error": f"unknown tool: {name}", "artifacts": []}
+
+ def _retriever(self, arguments):
+ query = str(arguments.get("query", self.sample.get("question", "")))
+ top_k = max(1, min(int(arguments.get("top_k", 3) or 3), 8))
+ docs = self.sample.get("documents", [])
+ scored = []
+ for doc in docs:
+ text = f"{doc.get('title', '')} {doc.get('text', '')}"
+ scored.append((_overlap_score(query, text), doc))
+ scored.sort(key=lambda item: item[0], reverse=True)
+ result = []
+ for score, doc in scored[:top_k]:
+ result.append(
+ {
+ "id": doc.get("id", ""),
+ "title": doc.get("title", ""),
+ "text": doc.get("text", ""),
+ "score": score,
+ }
+ )
+ return {"ok": True, "result": result, "error": "", "artifacts": []}
+
+ def _calculator(self, arguments):
+ expression = arguments.get("expression", "")
+ try:
+ value = safe_calculate(expression)
+ except Exception as exc:
+ return {"ok": False, "result": "", "error": str(exc)[:500], "artifacts": []}
+ return {"ok": True, "result": value, "error": "", "artifacts": []}
+
+ def _python_executor(self, arguments):
+ code = str(arguments.get("code", ""))
+ if not code.strip():
+ return {"ok": False, "result": "", "error": "empty code", "artifacts": []}
+ queue = mp.Queue(maxsize=1)
+ proc = mp.Process(target=_python_worker, args=(code, self.repo_root, self.allowed_paths, queue))
+ proc.start()
+ proc.join(self.timeout)
+ if proc.is_alive():
+ proc.terminate()
+ proc.join()
+ return {"ok": False, "result": "", "error": f"python_executor timeout>{self.timeout}s", "artifacts": []}
+ if queue.empty():
+ return {"ok": False, "result": "", "error": "python_executor returned no result", "artifacts": []}
+ return queue.get()
+
+
+def validate_tool_call(call, available_names):
+ if not call.get("valid_json", False):
+ return False
+ name = call.get("name", "")
+ args = call.get("arguments", {})
+ if name not in available_names:
+ return False
+ if name == "retriever":
+ return bool(args.get("query"))
+ if name == "calculator":
+ return bool(args.get("expression"))
+ if name == "python_executor":
+ return bool(args.get("code"))
+ return False
+
+
+def _f1(pred_items, gold_items):
+ pred = set(pred_items)
+ gold = set(gold_items)
+ if not pred and not gold:
+ return 1.0
+ if not pred or not gold:
+ return 0.0
+ inter = len(pred & gold)
+ if inter == 0:
+ return 0.0
+ precision = inter / len(pred)
+ recall = inter / len(gold)
+ return 2 * precision * recall / (precision + recall)
+
+
+def _numbers(text):
+ return [float(x.replace(",", "")) for x in re.findall(r"[-+]?\d+(?:,\d{3})*(?:\.\d+)?", str(text))]
+
+
+def score_answer(final_answer, sample):
+ checks = sample.get("checks", {}) or {}
+ score_parts = []
+ contains = checks.get("contains") or []
+ if contains:
+ score_parts.append(float(all(str(item).lower() in final_answer.lower() for item in contains)))
+ if "number" in checks:
+ target = _safe_number(checks.get("number"))
+ tolerance = _safe_number(checks.get("tolerance"), 1e-3)
+ nums = _numbers(final_answer)
+ score_parts.append(float(any(abs(num - target) <= tolerance for num in nums)))
+ if not score_parts and sample.get("answer"):
+ score_parts.append(float(str(sample["answer"]).lower() in final_answer.lower()))
+ return sum(score_parts) / len(score_parts) if score_parts else 0.0
+
+
+def repetition_penalty(text, n=3, cap=0.5):
+ toks = _tokenize(text)
+ grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+ return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
+
+
+def score_agentic_trajectory(turn_outputs, sample, repo_root=None, unfinished=False, execute_tools=True):
+ if isinstance(turn_outputs, str):
+ turn_outputs = [turn_outputs]
+ turn_outputs = turn_outputs or []
+ available_names = set(sample.get("tools", ["retriever", "python_executor", "calculator"]))
+ expected_tools = sample.get("expected_tools", [])
+ env = AgenticToolEnv(sample, repo_root=repo_root)
+
+ calls = []
+ tag_errors = 0
+ for turn in turn_outputs:
+ tag_errors += abs(str(turn).count("") - str(turn).count(""))
+ calls.extend(parse_tool_calls(turn))
+
+ valid_flags = [validate_tool_call(call, available_names) for call in calls]
+ valid_count = sum(valid_flags)
+ invalid_count = len(calls) - valid_count
+ selected_tools = [call.get("name", "") for call, ok in zip(calls, valid_flags) if ok]
+ tool_f1 = _f1(selected_tools, expected_tools)
+
+ exec_success = 1.0
+ if execute_tools and calls:
+ successes = []
+ for call, ok in zip(calls, valid_flags):
+ if not ok:
+ successes.append(False)
+ continue
+ result = env.execute(call["name"], call["arguments"])
+ successes.append(bool(result.get("ok")))
+ exec_success = sum(successes) / len(successes) if successes else 0.0
+
+ final_answer = extract_final_answer(turn_outputs)
+ answer_score = 0.0 if unfinished else score_answer(final_answer, sample)
+ task_success = float(answer_score >= 0.999)
+ schema_valid_rate = valid_count / len(calls) if calls else 0.0
+ invalid_action_rate = invalid_count / len(calls) if calls else 0.0
+ format_score = 1.0 if tag_errors == 0 and all(call.get("valid_json") for call in calls) else 0.0
+ turn_count = len(turn_outputs)
+ expected_turns = max(1, int(sample.get("max_turns", 4)))
+ turn_efficiency = max(0.0, 1.0 - max(0, turn_count - expected_turns) / expected_turns)
+ repeat_penalty = repetition_penalty(final_answer)
+
+ reward = (
+ 2.5 * task_success
+ + 1.2 * answer_score
+ + 0.8 * tool_f1
+ + 0.6 * schema_valid_rate
+ + 0.6 * exec_success
+ + 0.4 * format_score
+ + 0.2 * turn_efficiency
+ - 0.8 * invalid_action_rate
+ - 0.3 * tag_errors
+ - repeat_penalty
+ - (0.5 if unfinished else 0.0)
+ )
+ reward = max(min(reward, 5.0), -3.0)
+ return reward, {
+ "reward_total": reward,
+ "task_success": task_success,
+ "answer_score": answer_score,
+ "tool_selection_f1": tool_f1,
+ "schema_valid_rate": schema_valid_rate,
+ "exec_success_rate": exec_success,
+ "invalid_action_rate": invalid_action_rate,
+ "format_score": format_score,
+ "turn_efficiency": turn_efficiency,
+ "repetition_penalty": repeat_penalty,
+ "num_tool_calls": len(calls),
+ "avg_turns": turn_count,
+ }
+
+
+def average_agentic_metrics(metrics):
+ if not metrics:
+ return {}
+ keys = sorted({key for row in metrics for key in row})
+ return {
+ key: sum(float(row.get(key, 0.0)) for row in metrics) / len(metrics)
+ for key in keys
+ }
diff --git a/dataset/agentic_dataset.py b/dataset/agentic_dataset.py
new file mode 100644
index 0000000..86ab134
--- /dev/null
+++ b/dataset/agentic_dataset.py
@@ -0,0 +1,40 @@
+import json
+
+from torch.utils.data import Dataset
+
+from agentic.data_analysis_env import AGENTIC_SYSTEM_PROMPT, format_agentic_user_prompt, get_agentic_tools
+
+
+def load_agentic_jsonl(path):
+ rows = []
+ with open(path, "r", encoding="utf-8") as f:
+ for line in f:
+ line = line.strip()
+ if line:
+ rows.append(json.loads(line))
+ return rows
+
+
+class AgenticRLDataset(Dataset):
+ def __init__(self, jsonl_path, tokenizer=None, max_length=2048):
+ super().__init__()
+ self.samples = load_agentic_jsonl(jsonl_path)
+ self.tokenizer = tokenizer
+ self.max_length = max_length
+
+ def __len__(self):
+ return len(self.samples)
+
+ def __getitem__(self, index):
+ sample = self.samples[index]
+ tools = get_agentic_tools(sample.get("tools"))
+ messages = [
+ {"role": "system", "content": AGENTIC_SYSTEM_PROMPT},
+ {"role": "user", "content": format_agentic_user_prompt(sample)},
+ ]
+ return {
+ "id": sample.get("id", str(index)),
+ "messages": messages,
+ "tools": tools,
+ "sample": sample,
+ }
diff --git a/docs/agentic_data_analysis_project.md b/docs/agentic_data_analysis_project.md
new file mode 100644
index 0000000..c03abd5
--- /dev/null
+++ b/docs/agentic_data_analysis_project.md
@@ -0,0 +1,183 @@
+# Agentic DataAnalysis：面向运营数据分析的 Agentic RL 系统
+
+## 项目定位
+
+面向运营数据分析场景，构建一个 300M 级 Agentic LLM：模型需要先理解指标口径，再调用 `retriever`、`python_executor`、`calculator` 等工具完成多步分析，最终输出可验证的业务结论。
+
+这条线不重新预训练，直接使用已有 `SearchLM-300M` 预训练权重：
+
+```text
+out/pretrain_searchlm_1024_mla.pth
+```
+
+后续主流程：
+
+```text
+Agent 轨迹 SFT -> Agentic GRPO/CISPO -> 离线评测 -> Demo 展示
+```
+
+## 数据构造
+
+生成运营数据表、指标文档、任务、标准答案和专家工具轨迹：
+
+```bash
+python scripts/build_agentic_data_tasks.py \
+ --output_dir dataset/agentic_data \
+ --train_size 30000 \
+ --dev_size 1000 \
+ --test_size 1000
+```
+
+输出：
+
+- `dataset/agentic_data/tables/orders.csv`
+- `dataset/agentic_data/tables/marketing.csv`
+- `dataset/agentic_data/tables/users.csv`
+- `dataset/agentic_data/train.jsonl`
+- `dataset/agentic_data/dev.jsonl`
+- `dataset/agentic_data/test.jsonl`
+- `dataset/agentic_data/manifest.json`
+
+训练前校验：
+
+```bash
+python scripts/validate_agentic_data.py \
+ --data dataset/agentic_data/train.jsonl
+```
+
+## SFT 冷启动
+
+把专家轨迹转换成现有 `SFTDataset` 支持的 conversations 格式：
+
+```bash
+python scripts/build_agentic_sft.py \
+ --input dataset/agentic_data/train.jsonl \
+ --output dataset/agentic_data/sft.jsonl
+```
+
+使用已有 300M 预训练模型进行 Agent SFT：
+
+```bash
+cd trainer
+
+deepspeed --num_gpus=7 train_full_sft_deepspeed.py \
+ --model_profile SearchLM-300M \
+ --from_weight pretrain_searchlm \
+ --save_weight agent_sft \
+ --data_path ../dataset/agentic_data/sft.jsonl \
+ --dtype float16 \
+ --save_interval 200
+```
+
+完整 14GB 通用 SFT 数据建议先训练一轮，再进入 Agent SFT。服务器可直接使用一键脚本：
+
+```bash
+bash scripts/run_general_sft_300m.sh
+```
+
+脚本默认使用 7 张 GPU、每卡 batch 1、梯度累积 8，有效 batch 为 56；如果检测到
+`checkpoints/general_sft_14g_1024_mla_latest`，会自动加 `--from_resume 1`。日志写入 `logs/`。
+
+常用覆盖参数：
+
+```bash
+USE_WANDB=1 SAVE_INTERVAL=200 bash scripts/run_general_sft_300m.sh
+```
+
+通用 SFT 完成后，将下面 Agent SFT 命令中的 `--from_weight` 改为 `general_sft_14g`。
+
+## Agentic RL
+
+使用 trajectory-level reward 优化多步工具调用行为：
+
+```bash
+cd trainer
+
+deepspeed --num_gpus=7 train_agent_deepspeed.py \
+ --model_profile SearchLM-300M \
+ --from_weight agent_sft \
+ --save_weight agent_grpo \
+ --data_path ../dataset/agentic_data/train.jsonl \
+ --num_generations 6 \
+ --batch_size 1 \
+ --accumulation_steps 4 \
+ --dtype float16 \
+ --save_interval 25
+```
+
+断点续训：
+
+```bash
+deepspeed --num_gpus=7 train_agent_deepspeed.py \
+ --model_profile SearchLM-300M \
+ --from_weight agent_sft \
+ --save_weight agent_grpo \
+ --data_path ../dataset/agentic_data/train.jsonl \
+ --num_generations 6 \
+ --batch_size 1 \
+ --accumulation_steps 4 \
+ --dtype float16 \
+ --save_interval 25 \
+ --from_resume 1
+```
+
+`../checkpoints/` 保存完整 DeepSpeed 训练状态；`../out/agent_grpo_1024_mla.pth` 是轻量推理权重。
+
+## Reward 指标
+
+奖励函数位于 `agentic/data_analysis_env.py`，按整条轨迹评分：
+
+- `task_success`：最终答案是否命中标准答案中的文本和数值校验。
+- `answer_score`：文本 contains 与数值 tolerance 的平均得分。
+- `tool_selection_f1`：调用工具集合与期望工具集合的 F1。
+- `schema_valid_rate`：工具名、JSON 和必填参数是否合法。
+- `exec_success_rate`：工具能否成功执行。
+- `invalid_action_rate`：非法 action 占比。
+- `avg_turns`：平均交互轮数。
+- `format_score`：`` 标签和 JSON 是否闭合合法。
+
+## 离线评测
+
+直接跑模型生成：
+
+```bash
+python scripts/eval_agentic_data.py \
+ --data dataset/agentic_data/test.jsonl \
+ --weight agent_grpo \
+ --save_dir out \
+ --output reports/agentic_grpo_pred.jsonl \
+ --json_output reports/agentic_grpo_metrics.json \
+ --bad_cases_output reports/agentic_grpo_bad_cases.jsonl
+```
+
+也可以评测已有预测文件：
+
+```bash
+python scripts/eval_agentic_data.py \
+ --data dataset/agentic_data/test.jsonl \
+ --pred reports/agentic_grpo_pred.jsonl \
+ --json_output reports/agentic_grpo_metrics.json
+```
+
+建议固定对比：
+
+```text
+pretrain_searchlm -> agent_sft -> agent_grpo
+```
+
+核心简历指标使用固定 test split，不使用训练集结果。
+
+## Demo
+
+```bash
+streamlit run scripts/agentic_data_demo.py -- \
+ --data dataset/agentic_data/dev.jsonl
+```
+
+Demo 展示问题、数据表、指标文档、工具轨迹、最终答案和 reward 分解，适合面试时说明 Agent 行为如何被训练和评估。
+
+## 简历表达
+
+可以写成：
+
+> 构建面向运营数据分析的 Agentic RL 系统，基于 300M Decoder-only 预训练模型完成工具轨迹 SFT 与 GRPO/CISPO 后训练；设计 retriever、python executor、calculator 交互环境和 trajectory-level reward，在固定测试集上评估任务成功率、工具选择准确率、无效 action 率、Python 执行成功率和平均交互轮数。
diff --git a/model/chat_template.jinja b/model/chat_template.jinja
index 7771b03..c7603da 100644
--- a/model/chat_template.jinja
+++ b/model/chat_template.jinja
@@ -1,129 +1,69 @@
{%- if tools %}
- {{- '<|im_start|>system
-' }}
+ {{- '<|im_start|>system\n' }}
{%- if messages[0].role == 'system' %}
- {{- messages[0].content + '
-
-' }}
+ {{- messages[0].content + '\n\n' }}
{%- endif %}
- {{- "# Tools
-
-You may call one or more functions to assist with the user query.
-
-You are provided with function signatures within XML tags:
-" }}
+ {{- '# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n' }}
{%- for tool in tools %}
- {{- "
-" }}
+ {{- '\n' }}
{{- tool | tojson }}
{%- endfor %}
- {{- "
-
-
-For each function call, return a json object with function name and arguments within XML tags:
-
-{"name": , "arguments": }
-<|im_end|>
-" }}
+ {{- '\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{"name": , "arguments": }\n<|im_end|>\n' }}
{%- else %}
{%- if messages[0].role == 'system' %}
- {{- '<|im_start|>system
-' + messages[0].content + '<|im_end|>
-' }}
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
-{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
-{%- for message in messages[::-1] %}
- {%- set index = (messages|length - 1) - loop.index0 %}
- {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %}
- {%- set ns.multi_step_tool = false %}
- {%- set ns.last_query_index = index %}
- {%- endif %}
-{%- endfor %}
{%- for message in messages %}
{%- if message.content is string %}
{%- set content = message.content %}
{%- else %}
{%- set content = '' %}
{%- endif %}
- {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
- {{- '<|im_start|>' + message.role + '
-' + content + '<|im_end|>' + '
-' }}
- {%- elif message.role == "assistant" %}
+ {%- if message.role == 'user' or (message.role == 'system' and not loop.first) %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>\n' }}
+ {%- elif message.role == 'assistant' %}
{%- set reasoning_content = '' %}
{%- if message.reasoning_content is string %}
{%- set reasoning_content = message.reasoning_content %}
- {%- else %}
- {%- if '' in content %}
- {%- set reasoning_content = content.split('')[0].rstrip('
-').split('')[-1].lstrip('
-') %}
- {%- set content = content.split('')[-1].lstrip('
-') %}
- {%- endif %}
- {%- endif %}
- {%- if true %}
- {{- '<|im_start|>' + message.role + '
-
-' + reasoning_content.strip('
-') + '
-
-
-' + content.lstrip('
-') }}
+ {%- elif '' in content %}
+ {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %}
+ {%- set content = content.split('')[-1].lstrip('\n') %}
{%- endif %}
+ {{- '<|im_start|>assistant\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }}
{%- if message.tool_calls %}
{%- for tool_call in message.tool_calls %}
- {%- if (loop.first and content) or (not loop.first) %}
- {{- '
-' }}
+ {%- if (loop.first and content) or not loop.first %}
+ {{- '\n' }}
{%- endif %}
{%- if tool_call.function %}
{%- set tool_call = tool_call.function %}
{%- endif %}
- {{- '
-{"name": "' }}
- {{- tool_call.name }}
- {{- '", "arguments": ' }}
+ {{- '\n{"name": "' + tool_call.name + '", "arguments": ' }}
{%- if tool_call.arguments is string %}
{{- tool_call.arguments }}
{%- else %}
{{- tool_call.arguments | tojson }}
{%- endif %}
- {{- '}
-' }}
+ {{- '}\n' }}
{%- endfor %}
{%- endif %}
- {{- '<|im_end|>
-' }}
- {%- elif message.role == "tool" %}
- {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == 'tool' %}
+ {%- if loop.first or messages[loop.index0 - 1].role != 'tool' %}
{{- '<|im_start|>user' }}
{%- endif %}
- {{- '
-
-' }}
- {{- content }}
- {{- '
-' }}
- {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
- {{- '<|im_end|>
-' }}
+ {{- '\n\n' + content + '\n' }}
+ {%- if loop.last or messages[loop.index0 + 1].role != 'tool' %}
+ {{- '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
- {{- '<|im_start|>assistant
-' }}
+ {{- '<|im_start|>assistant\n' }}
{%- if open_thinking is defined and open_thinking is true %}
- {{- '
-' }}
+ {{- '\n' }}
{%- else %}
- {{- '
-
-
-
-' }}
+ {{- '\n\n\n\n' }}
{%- endif %}
-{%- endif %}
\ No newline at end of file
+{%- endif %}
diff --git a/requirements.txt b/requirements.txt
index 0ac5398..7c3be3c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,6 +10,7 @@ ngrok==1.4.0
nltk==3.8
numpy==1.26.4
openai==1.59.6
+pandas==2.2.3
# peft==0.7.1
psutil==5.9.8
pydantic==2.11.5
diff --git a/scripts/agentic_data_demo.py b/scripts/agentic_data_demo.py
new file mode 100644
index 0000000..95f7a54
--- /dev/null
+++ b/scripts/agentic_data_demo.py
@@ -0,0 +1,97 @@
+import argparse
+import json
+import os
+import sys
+
+import pandas as pd
+import streamlit as st
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from agentic.data_analysis_env import extract_final_answer, score_agentic_trajectory
+from dataset.agentic_dataset import load_agentic_jsonl
+
+
+DEFAULT_DATA = "dataset/agentic_data/dev.jsonl"
+
+
+def trajectory_to_text(sample):
+ chunks = []
+ for item in sample.get("expert_trajectory", []):
+ if item.get("role") == "assistant" and item.get("tool_call"):
+ chunks.append("\n" + json.dumps(item["tool_call"], ensure_ascii=False, indent=2) + "\n")
+ elif item.get("role") == "assistant":
+ chunks.append(str(item.get("content", "")))
+ return "\n---\n".join(chunks)
+
+
+def parse_turns(text):
+ if "\n---\n" in text:
+ return [part.strip() for part in text.split("\n---\n") if part.strip()]
+ return [text.strip()] if text.strip() else []
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Agentic DataAnalysis Streamlit demo")
+ parser.add_argument("--data", default=DEFAULT_DATA)
+ parser.add_argument("--repo_root", default=".")
+ return parser.parse_args()
+
+
+def main():
+ args = parse_args()
+ st.set_page_config(page_title="Agentic DataAnalysis", layout="wide")
+ st.title("Agentic DataAnalysis")
+
+ rows = load_agentic_jsonl(args.data) if os.path.exists(args.data) else []
+ if not rows:
+ st.error(f"数据不存在或为空：{args.data}")
+ return
+
+ sample_ids = [row.get("id", str(i)) for i, row in enumerate(rows)]
+ selected = st.sidebar.selectbox("样本", sample_ids)
+ sample = rows[sample_ids.index(selected)]
+
+ st.subheader("问题")
+ st.write(sample.get("question", ""))
+
+ left, right = st.columns([1.1, 1.0])
+ with left:
+ st.subheader("数据表")
+ st.dataframe(pd.DataFrame(sample.get("tables", [])), use_container_width=True)
+ st.subheader("指标文档")
+ st.dataframe(pd.DataFrame(sample.get("documents", [])), use_container_width=True)
+
+ with right:
+ st.subheader("标准答案")
+ st.write(sample.get("answer", ""))
+ st.json(sample.get("checks", {}))
+ st.subheader("期望工具")
+ st.write(", ".join(sample.get("expected_tools", [])))
+
+ st.subheader("轨迹与奖励")
+ default_text = trajectory_to_text(sample)
+ trajectory_text = st.text_area("轨迹文本（用 --- 分隔多个 assistant turn 可手动评测模型输出）", value=default_text, height=360)
+ turn_outputs = [
+ part
+ for part in parse_turns(trajectory_text)
+ if "" not in part
+ ]
+ reward, parts = score_agentic_trajectory(turn_outputs, sample, repo_root=args.repo_root)
+
+ c1, c2, c3, c4, c5 = st.columns(5)
+ c1.metric("Reward", f"{reward:.3f}")
+ c2.metric("Success", f"{parts['task_success']:.3f}")
+ c3.metric("Tool F1", f"{parts['tool_selection_f1']:.3f}")
+ c4.metric("Invalid", f"{parts['invalid_action_rate']:.3f}")
+ c5.metric("Turns", f"{parts['avg_turns']:.1f}")
+
+ st.subheader("Reward 分解")
+ st.dataframe(pd.DataFrame([parts]).T.rename(columns={0: "score"}), use_container_width=True)
+
+ st.subheader("Final Answer")
+ st.write(extract_final_answer(turn_outputs))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/build_agentic_data_tasks.py b/scripts/build_agentic_data_tasks.py
new file mode 100644
index 0000000..8e7391a
--- /dev/null
+++ b/scripts/build_agentic_data_tasks.py
@@ -0,0 +1,488 @@
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+from collections import defaultdict
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from agentic.data_analysis_env import AgenticToolEnv, get_agentic_tools
+
+
+CHANNELS = ["抖音", "微信", "百度", "小红书", "自然流量"]
+CATEGORIES = ["美妆", "数码", "服饰", "食品", "家居"]
+PROVINCES = ["广东", "浙江", "北京", "上海", "四川", "湖北"]
+
+METRIC_DOCS = [
+ {
+ "id": "metric_refund_amount",
+ "title": "退款金额",
+ "text": "退款金额使用 orders 表的 refund_amount 字段求和；按渠道、品类、省份或时间过滤后再聚合。",
+ },
+ {
+ "id": "metric_net_revenue",
+ "title": "净收入",
+ "text": "净收入 = GMV - 退款金额。GMV 使用 orders.gmv，退款金额使用 orders.refund_amount。",
+ },
+ {
+ "id": "metric_refund_rate",
+ "title": "退款率",
+ "text": "退款率 = 退款金额 / GMV。按品类或渠道对 refund_amount 与 gmv 分别求和后再相除。",
+ },
+ {
+ "id": "metric_roas",
+ "title": "ROAS",
+ "text": "ROAS = GMV / 广告花费。GMV 来自 orders 表，广告花费来自 marketing 表的 spend 字段。",
+ },
+ {
+ "id": "metric_new_users",
+ "title": "新增用户",
+ "text": "新增用户使用 users 表的 new_users 字段求和；活跃用户使用 active_users 字段求和。",
+ },
+ {
+ "id": "schema_orders",
+ "title": "orders 表字段",
+ "text": "orders 包含 month、quarter、channel、category、province、gmv、refund_amount、orders_count。",
+ },
+ {
+ "id": "schema_marketing",
+ "title": "marketing 表字段",
+ "text": "marketing 包含 month、quarter、channel、spend、impressions、clicks。",
+ },
+ {
+ "id": "schema_users",
+ "title": "users 表字段",
+ "text": "users 包含 month、quarter、channel、province、active_users、new_users。",
+ },
+]
+
+
+def mkdir(path):
+ os.makedirs(path, exist_ok=True)
+
+
+def write_csv(path, rows):
+ if not rows:
+ raise ValueError(f"empty csv rows: {path}")
+ with open(path, "w", encoding="utf-8", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+ writer.writeheader()
+ writer.writerows(rows)
+
+
+def relpath(path):
+ cwd = os.getcwd()
+ try:
+ return os.path.relpath(path, cwd)
+ except ValueError:
+ return path
+
+
+def quarter_of(month):
+ return f"Q{(month - 1) // 3 + 1}"
+
+
+def generate_tables(output_dir, seed):
+ rng = random.Random(seed)
+ table_dir = os.path.join(output_dir, "tables")
+ mkdir(table_dir)
+
+ orders = []
+ for month in range(1, 13):
+ season = 1.0 + (0.18 if month in {6, 11, 12} else 0.0) + (0.08 if month in {1, 2} else 0.0)
+ for channel_idx, channel in enumerate(CHANNELS):
+ for category_idx, category in enumerate(CATEGORIES):
+ for province_idx, province in enumerate(PROVINCES):
+ base = 42000 + channel_idx * 4600 + category_idx * 3100 + province_idx * 1800
+ gmv = round((base + rng.uniform(-5200, 6200)) * season, 2)
+ refund_rate = 0.025 + 0.006 * category_idx + 0.003 * channel_idx + rng.uniform(0.0, 0.018)
+ refund_amount = round(gmv * refund_rate, 2)
+ orders_count = max(1, int(gmv / (90 + category_idx * 12 + rng.uniform(-8, 8))))
+ orders.append(
+ {
+ "month": month,
+ "quarter": quarter_of(month),
+ "channel": channel,
+ "category": category,
+ "province": province,
+ "gmv": gmv,
+ "refund_amount": refund_amount,
+ "orders_count": orders_count,
+ }
+ )
+
+ marketing = []
+ for month in range(1, 13):
+ for channel_idx, channel in enumerate(CHANNELS):
+ spend = round(18000 + channel_idx * 3500 + month * 420 + rng.uniform(-1800, 2200), 2)
+ clicks = max(1, int(spend / (2.2 + channel_idx * 0.25 + rng.uniform(-0.2, 0.25))))
+ impressions = clicks * rng.randint(18, 36)
+ marketing.append(
+ {
+ "month": month,
+ "quarter": quarter_of(month),
+ "channel": channel,
+ "spend": spend,
+ "impressions": impressions,
+ "clicks": clicks,
+ }
+ )
+
+ users = []
+ for month in range(1, 13):
+ for channel_idx, channel in enumerate(CHANNELS):
+ for province_idx, province in enumerate(PROVINCES):
+ active_users = max(100, int(5200 + month * 110 + channel_idx * 520 + province_idx * 260 + rng.uniform(-420, 520)))
+ new_users = max(20, int(active_users * (0.08 + channel_idx * 0.008 + rng.uniform(0.0, 0.04))))
+ users.append(
+ {
+ "month": month,
+ "quarter": quarter_of(month),
+ "channel": channel,
+ "province": province,
+ "active_users": active_users,
+ "new_users": new_users,
+ }
+ )
+
+ paths = {
+ "orders": os.path.join(table_dir, "orders.csv"),
+ "marketing": os.path.join(table_dir, "marketing.csv"),
+ "users": os.path.join(table_dir, "users.csv"),
+ }
+ write_csv(paths["orders"], orders)
+ write_csv(paths["marketing"], marketing)
+ write_csv(paths["users"], users)
+ return {"orders": orders, "marketing": marketing, "users": users, "paths": paths}
+
+
+def table_specs(paths):
+ return [
+ {
+ "name": "orders",
+ "path": relpath(paths["orders"]),
+ "description": "订单聚合表，包含 GMV、退款金额和订单数。",
+ },
+ {
+ "name": "marketing",
+ "path": relpath(paths["marketing"]),
+ "description": "投放聚合表，包含广告花费、曝光和点击。",
+ },
+ {
+ "name": "users",
+ "path": relpath(paths["users"]),
+ "description": "用户活跃聚合表，包含活跃用户和新增用户。",
+ },
+ ]
+
+
+def sum_field(rows, field, **filters):
+ total = 0.0
+ for row in rows:
+ if all(row.get(key) == value for key, value in filters.items()):
+ total += float(row[field])
+ return round(total, 2)
+
+
+def top_group(rows, group_key, value_field, filters):
+ agg = defaultdict(float)
+ for row in rows:
+ if all(row.get(key) == value for key, value in filters.items()):
+ agg[row[group_key]] += float(row[value_field])
+ key, value = max(agg.items(), key=lambda item: item[1])
+ return key, round(value, 2)
+
+
+def refund_rate_by_category(rows, month):
+ agg = defaultdict(lambda: {"gmv": 0.0, "refund": 0.0})
+ for row in rows:
+ if row["month"] == month:
+ agg[row["category"]]["gmv"] += float(row["gmv"])
+ agg[row["category"]]["refund"] += float(row["refund_amount"])
+ rates = {
+ category: values["refund"] / values["gmv"]
+ for category, values in agg.items()
+ if values["gmv"] > 0
+ }
+ category, rate = max(rates.items(), key=lambda item: item[1])
+ return category, round(rate, 4)
+
+
+def monthly_gmv_by_channel(orders, month, channel):
+ return sum(float(row["gmv"]) for row in orders if row["month"] == month and row["channel"] == channel)
+
+
+def marketing_spend(marketing, month, channel):
+ return sum(float(row["spend"]) for row in marketing if row["month"] == month and row["channel"] == channel)
+
+
+def build_tool_call(name, arguments):
+ return {"name": name, "arguments": arguments}
+
+
+def make_sample(sample_id, question, answer, checks, tables, expected_tools, actions, task_type, max_turns=4):
+ sample = {
+ "id": sample_id,
+ "question": question,
+ "tables": tables,
+ "documents": METRIC_DOCS,
+ "tools": ["retriever", "python_executor", "calculator"],
+ "expected_tools": expected_tools,
+ "answer": answer,
+ "checks": checks,
+ "max_turns": max_turns,
+ "metadata": {"task_type": task_type},
+ }
+ env = AgenticToolEnv(sample)
+ trajectory = []
+ for action in actions:
+ trajectory.append({"role": "assistant", "tool_call": action})
+ result = env.execute(action["name"], action["arguments"])
+ trajectory.append({"role": "tool", "name": action["name"], "content": result})
+ trajectory.append({"role": "assistant", "content": answer})
+ sample["expert_trajectory"] = trajectory
+ sample["tool_schemas"] = get_agentic_tools(sample["tools"])
+ return sample
+
+
+def task_top_refund_channel(sample_id, tables, data, rng):
+ quarter = rng.choice(["Q1", "Q2", "Q3", "Q4"])
+ channel, value = top_group(data["orders"], "channel", "refund_amount", {"quarter": quarter})
+ question = f"2024年{quarter}退款金额最高的渠道是什么？请给出退款金额。"
+ code = (
+ f'orders = pd.read_csv("{tables[0]["path"]}")\n'
+ f'df = orders[orders["quarter"] == "{quarter}"]\n'
+ 'agg = df.groupby("channel")["refund_amount"].sum().sort_values(ascending=False)\n'
+ 'result = {"channel": agg.index[0], "refund_amount": round(float(agg.iloc[0]), 2)}'
+ )
+ answer = f"2024年{quarter}退款金额最高的渠道是{channel}，退款金额为{value:.2f}元。"
+ return make_sample(
+ sample_id,
+ question,
+ answer,
+ {"contains": [channel], "number": value, "tolerance": 0.02},
+ tables,
+ ["retriever", "python_executor"],
+ [
+ build_tool_call("retriever", {"query": "退款金额 指标口径 orders refund_amount", "top_k": 2}),
+ build_tool_call("python_executor", {"code": code}),
+ ],
+ "top_refund_channel",
+ )
+
+
+def task_net_revenue(sample_id, tables, data, rng):
+ month = rng.randint(1, 12)
+ channel = rng.choice(CHANNELS)
+ gmv = sum_field(data["orders"], "gmv", month=month, channel=channel)
+ refund = sum_field(data["orders"], "refund_amount", month=month, channel=channel)
+ value = round(gmv - refund, 2)
+ question = f"2024年{month}月{channel}渠道的净收入是多少？"
+ code = (
+ f'orders = pd.read_csv("{tables[0]["path"]}")\n'
+ f'df = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]\n'
+ 'result = {"net_revenue": round(float(df["gmv"].sum() - df["refund_amount"].sum()), 2)}'
+ )
+ answer = f"2024年{month}月{channel}渠道的净收入为{value:.2f}元。"
+ return make_sample(
+ sample_id,
+ question,
+ answer,
+ {"contains": [channel], "number": value, "tolerance": 0.02},
+ tables,
+ ["retriever", "python_executor"],
+ [
+ build_tool_call("retriever", {"query": "净收入 GMV 退款金额 计算口径", "top_k": 2}),
+ build_tool_call("python_executor", {"code": code}),
+ ],
+ "net_revenue",
+ )
+
+
+def task_refund_rate_category(sample_id, tables, data, rng):
+ month = rng.randint(1, 12)
+ category, rate = refund_rate_by_category(data["orders"], month)
+ percent = round(rate * 100, 2)
+ question = f"2024年{month}月退款率最高的品类是什么？退款率是多少？"
+ code = (
+ f'orders = pd.read_csv("{tables[0]["path"]}")\n'
+ f'df = orders[orders["month"] == {month}]\n'
+ 'agg = df.groupby("category")[["refund_amount", "gmv"]].sum()\n'
+ 'agg["refund_rate"] = agg["refund_amount"] / agg["gmv"]\n'
+ 'top = agg.sort_values("refund_rate", ascending=False).iloc[0]\n'
+ 'result = {"category": agg.sort_values("refund_rate", ascending=False).index[0], "refund_rate": round(float(top["refund_rate"] * 100), 2)}'
+ )
+ answer = f"2024年{month}月退款率最高的品类是{category}，退款率为{percent:.2f}%。"
+ return make_sample(
+ sample_id,
+ question,
+ answer,
+ {"contains": [category], "number": percent, "tolerance": 0.02},
+ tables,
+ ["retriever", "python_executor"],
+ [
+ build_tool_call("retriever", {"query": "退款率 指标口径 refund_amount gmv", "top_k": 2}),
+ build_tool_call("python_executor", {"code": code}),
+ ],
+ "refund_rate_category",
+ )
+
+
+def task_roas(sample_id, tables, data, rng):
+ month = rng.randint(1, 12)
+ channel = rng.choice([c for c in CHANNELS if c != "自然流量"])
+ gmv = monthly_gmv_by_channel(data["orders"], month, channel)
+ spend = marketing_spend(data["marketing"], month, channel)
+ value = round(gmv / spend, 2)
+ question = f"2024年{month}月{channel}渠道的ROAS是多少？"
+ code = (
+ f'orders = pd.read_csv("{tables[0]["path"]}")\n'
+ f'marketing = pd.read_csv("{tables[1]["path"]}")\n'
+ f'gmv = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]["gmv"].sum()\n'
+ f'spend = marketing[(marketing["month"] == {month}) & (marketing["channel"] == "{channel}")]["spend"].sum()\n'
+ 'result = {"roas": round(float(gmv / spend), 2)}'
+ )
+ answer = f"2024年{month}月{channel}渠道的ROAS为{value:.2f}。"
+ return make_sample(
+ sample_id,
+ question,
+ answer,
+ {"contains": [channel], "number": value, "tolerance": 0.02},
+ tables,
+ ["retriever", "python_executor"],
+ [
+ build_tool_call("retriever", {"query": "ROAS GMV 广告花费 指标口径", "top_k": 2}),
+ build_tool_call("python_executor", {"code": code}),
+ ],
+ "roas",
+ )
+
+
+def task_new_users(sample_id, tables, data, rng):
+ month = rng.randint(1, 12)
+ channel = rng.choice(CHANNELS)
+ province, value = top_group(data["users"], "province", "new_users", {"month": month, "channel": channel})
+ question = f"2024年{month}月{channel}渠道新增用户最多的省份是哪个？新增用户数是多少？"
+ code = (
+ f'users = pd.read_csv("{tables[2]["path"]}")\n'
+ f'df = users[(users["month"] == {month}) & (users["channel"] == "{channel}")]\n'
+ 'agg = df.groupby("province")["new_users"].sum().sort_values(ascending=False)\n'
+ 'result = {"province": agg.index[0], "new_users": int(agg.iloc[0])}'
+ )
+ answer = f"2024年{month}月{channel}渠道新增用户最多的省份是{province}，新增用户数为{int(value)}。"
+ return make_sample(
+ sample_id,
+ question,
+ answer,
+ {"contains": [province], "number": int(value), "tolerance": 0.01},
+ tables,
+ ["retriever", "python_executor"],
+ [
+ build_tool_call("retriever", {"query": "新增用户 users new_users 字段口径", "top_k": 2}),
+ build_tool_call("python_executor", {"code": code}),
+ ],
+ "new_users",
+ )
+
+
+def task_incremental_gmv(sample_id, tables, data, rng):
+ month = rng.randint(1, 12)
+ channel = rng.choice([c for c in CHANNELS if c != "自然流量"])
+ pct = rng.choice([10, 12, 15, 20])
+ gmv = monthly_gmv_by_channel(data["orders"], month, channel)
+ spend = marketing_spend(data["marketing"], month, channel)
+ roas = gmv / spend
+ incremental = round(spend * pct / 100 * roas, 2)
+ question = f"如果2024年{month}月{channel}渠道广告花费增加{pct}%，且ROAS保持不变，预计GMV增加多少？"
+ code = (
+ f'orders = pd.read_csv("{tables[0]["path"]}")\n'
+ f'marketing = pd.read_csv("{tables[1]["path"]}")\n'
+ f'gmv = orders[(orders["month"] == {month}) & (orders["channel"] == "{channel}")]["gmv"].sum()\n'
+ f'spend = marketing[(marketing["month"] == {month}) & (marketing["channel"] == "{channel}")]["spend"].sum()\n'
+ 'result = {"spend": round(float(spend), 2), "roas": round(float(gmv / spend), 4)}'
+ )
+ expression = f"{spend:.2f} * {pct} / 100 * {roas:.6f}"
+ answer = f"预计GMV增加{incremental:.2f}元。"
+ return make_sample(
+ sample_id,
+ question,
+ answer,
+ {"contains": ["GMV"], "number": incremental, "tolerance": 0.05},
+ tables,
+ ["retriever", "python_executor", "calculator"],
+ [
+ build_tool_call("retriever", {"query": "ROAS GMV 广告花费 预算增加 计算口径", "top_k": 2}),
+ build_tool_call("python_executor", {"code": code}),
+ build_tool_call("calculator", {"expression": expression}),
+ ],
+ "incremental_gmv",
+ max_turns=5,
+ )
+
+
+TASK_BUILDERS = [
+ task_top_refund_channel,
+ task_net_revenue,
+ task_refund_rate_category,
+ task_roas,
+ task_new_users,
+ task_incremental_gmv,
+]
+
+
+def write_jsonl(path, rows):
+ with open(path, "w", encoding="utf-8") as f:
+ for row in rows:
+ f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+
+def build_split(name, size, tables, data, seed):
+ rng = random.Random(seed)
+ rows = []
+ for i in range(size):
+ task_fn = TASK_BUILDERS[i % len(TASK_BUILDERS)]
+ sample_id = f"agentic_{name}_{i:06d}"
+ rows.append(task_fn(sample_id, tables, data, rng))
+ rng.shuffle(rows)
+ return rows
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Build Agentic DataAnalysis tasks")
+ parser.add_argument("--output_dir", default="dataset/agentic_data")
+ parser.add_argument("--train_size", type=int, default=30000)
+ parser.add_argument("--dev_size", type=int, default=1000)
+ parser.add_argument("--test_size", type=int, default=1000)
+ parser.add_argument("--seed", type=int, default=42)
+ return parser.parse_args()
+
+
+def main():
+ args = parse_args()
+ mkdir(args.output_dir)
+ data = generate_tables(args.output_dir, args.seed)
+ tables = table_specs(data["paths"])
+ splits = {
+ "train": build_split("train", args.train_size, tables, data, args.seed + 1),
+ "dev": build_split("dev", args.dev_size, tables, data, args.seed + 2),
+ "test": build_split("test", args.test_size, tables, data, args.seed + 3),
+ }
+ for split, rows in splits.items():
+ write_jsonl(os.path.join(args.output_dir, f"{split}.jsonl"), rows)
+ manifest = {
+ "task": "AgenticDataAnalysis",
+ "seed": args.seed,
+ "splits": {split: len(rows) for split, rows in splits.items()},
+ "tables": tables,
+ "tools": [tool["function"]["name"] for tool in get_agentic_tools()],
+ "task_types": [fn.__name__.replace("task_", "") for fn in TASK_BUILDERS],
+ }
+ with open(os.path.join(args.output_dir, "manifest.json"), "w", encoding="utf-8") as f:
+ json.dump(manifest, f, ensure_ascii=False, indent=2)
+ print(json.dumps(manifest, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/build_agentic_sft.py b/scripts/build_agentic_sft.py
new file mode 100644
index 0000000..75e345d
--- /dev/null
+++ b/scripts/build_agentic_sft.py
@@ -0,0 +1,86 @@
+import argparse
+import json
+import os
+import sys
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from agentic.data_analysis_env import AGENTIC_SYSTEM_PROMPT, format_agentic_user_prompt, get_agentic_tools
+from dataset.agentic_dataset import load_agentic_jsonl
+
+
+def assistant_tool_message(tool_call):
+ return {
+ "role": "assistant",
+ "content": "",
+ "reasoning_content": "",
+ "tools": "",
+ "tool_calls": json.dumps(
+ [
+ {
+ "type": "function",
+ "function": {
+ "name": tool_call.get("name", ""),
+ "arguments": tool_call.get("arguments", {}),
+ },
+ }
+ ],
+ ensure_ascii=False,
+ ),
+ }
+
+
+def normal_message(role, content):
+ return {
+ "role": role,
+ "content": content,
+ "reasoning_content": "",
+ "tools": "",
+ "tool_calls": "",
+ }
+
+
+def convert_sample(sample):
+ tools = get_agentic_tools(sample.get("tools"))
+ conversations = [
+ {
+ "role": "system",
+ "content": AGENTIC_SYSTEM_PROMPT,
+ "reasoning_content": "",
+ "tools": json.dumps(tools, ensure_ascii=False),
+ "tool_calls": "",
+ },
+ normal_message("user", format_agentic_user_prompt(sample)),
+ ]
+ for item in sample.get("expert_trajectory", []):
+ if item.get("role") == "assistant" and item.get("tool_call"):
+ conversations.append(assistant_tool_message(item["tool_call"]))
+ elif item.get("role") == "tool":
+ conversations.append(normal_message("tool", json.dumps(item.get("content", {}), ensure_ascii=False)))
+ elif item.get("role") == "assistant":
+ conversations.append(normal_message("assistant", str(item.get("content", ""))))
+ return {"id": sample.get("id", ""), "conversations": conversations}
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Build Agentic DataAnalysis SFT conversations")
+ parser.add_argument("--input", required=True)
+ parser.add_argument("--output", required=True)
+ parser.add_argument("--limit", type=int, default=0)
+ return parser.parse_args()
+
+
+def main():
+ args = parse_args()
+ rows = load_agentic_jsonl(args.input)
+ if args.limit > 0:
+ rows = rows[:args.limit]
+ os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+ with open(args.output, "w", encoding="utf-8") as f:
+ for sample in rows:
+ f.write(json.dumps(convert_sample(sample), ensure_ascii=False) + "\n")
+ print(json.dumps({"input": args.input, "output": args.output, "samples": len(rows)}, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/eval_agentic_data.py b/scripts/eval_agentic_data.py
new file mode 100644
index 0000000..d089b93
--- /dev/null
+++ b/scripts/eval_agentic_data.py
@@ -0,0 +1,237 @@
+import argparse
+import json
+import os
+import sys
+import time
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from agentic.data_analysis_env import (
+ AgenticToolEnv,
+ average_agentic_metrics,
+ extract_final_answer,
+ format_agentic_user_prompt,
+ get_agentic_tools,
+ parse_tool_calls,
+ score_agentic_trajectory,
+)
+from dataset.agentic_dataset import load_agentic_jsonl
+
+
+def apply_eval_model_profile(args):
+ if not args.model_profile:
+ return
+ path = args.model_profiles_path
+ with open(path, "r", encoding="utf-8") as f:
+ profiles = json.load(f)
+ if args.model_profile not in profiles:
+ raise ValueError(f"Unknown model_profile={args.model_profile}. Available: {', '.join(profiles)}")
+ profile = profiles[args.model_profile]
+ for key, value in profile.items():
+ if hasattr(args, key):
+ setattr(args, key, value)
+
+
+def load_model(args):
+ import torch
+ from transformers import AutoTokenizer
+
+ from model.model_minimind import MiniMindForCausalLM
+ from model.model_minimind_mla import MiniMindMLAForCausalLM
+ from trainer.trainer_utils import build_lm_config, get_model_suffix, resolve_attention_type
+
+ tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
+ config = build_lm_config(
+ hidden_size=args.hidden_size,
+ num_hidden_layers=args.num_hidden_layers,
+ max_seq_len=args.max_total_len,
+ use_moe=bool(args.use_moe),
+ attention_type=resolve_attention_type(args),
+ kv_lora_rank=args.kv_lora_rank,
+ q_lora_rank=args.q_lora_rank,
+ rope_dim=args.rope_dim,
+ num_attention_heads=args.num_attention_heads,
+ num_key_value_heads=args.num_key_value_heads,
+ intermediate_size=args.intermediate_size,
+ )
+ model = MiniMindMLAForCausalLM(config) if args.attention_type == "mla" or args.use_mla else MiniMindForCausalLM(config)
+ suffix = get_model_suffix(config)
+ weight_path = os.path.join(args.save_dir, f"{args.weight}_{config.hidden_size}{suffix}.pth")
+ weights = torch.load(weight_path, map_location=args.device)
+ model.load_state_dict(weights, strict=False)
+ dtype = torch.float16 if args.dtype == "float16" else torch.float32
+ if args.dtype == "bfloat16":
+ dtype = torch.bfloat16
+ model = model.to(args.device)
+ if args.device != "cpu" and args.dtype != "float32":
+ model = model.to(dtype=dtype)
+ return model.eval(), tokenizer
+
+
+def generate_once(model, tokenizer, messages, tools, args):
+ import torch
+
+ prompt = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ tools=tools,
+ open_thinking=bool(args.open_thinking),
+ )
+ inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(args.device)
+ if inputs["input_ids"].size(1) > args.max_total_len:
+ inputs["input_ids"] = inputs["input_ids"][:, -args.max_total_len:]
+ inputs["attention_mask"] = inputs["attention_mask"][:, -args.max_total_len:]
+ st = time.time()
+ with torch.no_grad():
+ generated = model.generate(
+ inputs["input_ids"],
+ attention_mask=inputs["attention_mask"],
+ max_new_tokens=args.max_new_tokens,
+ do_sample=bool(args.do_sample),
+ temperature=args.temperature,
+ top_p=args.top_p,
+ pad_token_id=tokenizer.pad_token_id,
+ eos_token_id=tokenizer.eos_token_id,
+ )
+ text = tokenizer.decode(generated[0][inputs["input_ids"].size(1):], skip_special_tokens=True)
+ return text, time.time() - st
+
+
+def run_episode(model, tokenizer, sample, args):
+ tools = get_agentic_tools(sample.get("tools"))
+ messages = [
+ {"role": "system", "content": "你是一个面向运营数据分析的 Agent。"},
+ {"role": "user", "content": format_agentic_user_prompt(sample)},
+ ]
+ env = AgenticToolEnv(sample, repo_root=args.repo_root, timeout=args.tool_timeout)
+ turn_outputs = []
+ elapsed = 0.0
+ unfinished = False
+ for turn in range(args.max_turns):
+ text, dt = generate_once(model, tokenizer, messages, tools, args)
+ elapsed += dt
+ turn_outputs.append(text)
+ calls = parse_tool_calls(text)
+ if not calls:
+ break
+ unfinished = turn == args.max_turns - 1
+ messages.append({"role": "assistant", "content": text})
+ for call in calls:
+ result = env.execute(call.get("name", ""), call.get("arguments", {}))
+ messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)})
+ if unfinished:
+ break
+ reward, parts = score_agentic_trajectory(
+ turn_outputs,
+ sample,
+ repo_root=args.repo_root,
+ unfinished=unfinished,
+ execute_tools=True,
+ )
+ return {
+ "id": sample.get("id", ""),
+ "question": sample.get("question", ""),
+ "prediction": extract_final_answer(turn_outputs),
+ "turn_outputs": turn_outputs,
+ "reward": reward,
+ "metrics": parts,
+ "latency_sec": elapsed,
+ }
+
+
+def eval_predictions(args, rows):
+ pred_map = {}
+ with open(args.pred, "r", encoding="utf-8") as f:
+ for line in f:
+ if line.strip():
+ item = json.loads(line)
+ pred_map[item.get("id")] = item
+ outputs = []
+ for sample in rows:
+ item = pred_map.get(sample.get("id"), {})
+ turns = item.get("turn_outputs") or [item.get("prediction") or item.get("response") or item.get("output") or ""]
+ reward, parts = score_agentic_trajectory(turns, sample, repo_root=args.repo_root)
+ outputs.append(
+ {
+ "id": sample.get("id", ""),
+ "question": sample.get("question", ""),
+ "prediction": extract_final_answer(turns),
+ "turn_outputs": turns,
+ "reward": reward,
+ "metrics": parts,
+ }
+ )
+ return outputs
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Evaluate Agentic DataAnalysis model or predictions")
+ parser.add_argument("--data", required=True)
+ parser.add_argument("--pred", default="")
+ parser.add_argument("--output", default="reports/agentic_eval_predictions.jsonl")
+ parser.add_argument("--json_output", default="reports/agentic_eval_metrics.json")
+ parser.add_argument("--bad_cases_output", default="reports/agentic_eval_bad_cases.jsonl")
+ parser.add_argument("--max_samples", type=int, default=0)
+ parser.add_argument("--repo_root", default=".")
+ parser.add_argument("--device", default="cuda")
+ parser.add_argument("--save_dir", default="out")
+ parser.add_argument("--weight", default="agent_grpo")
+ parser.add_argument("--tokenizer_path", default="model")
+ parser.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
+ parser.add_argument("--hidden_size", default=1024, type=int)
+ parser.add_argument("--num_hidden_layers", default=24, type=int)
+ parser.add_argument("--num_attention_heads", default=8, type=int)
+ parser.add_argument("--num_key_value_heads", default=4, type=int)
+ parser.add_argument("--intermediate_size", default=3264, type=int)
+ parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--attention_type", default="mla", choices=["gqa", "mha", "mqa", "mla"])
+ parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--kv_lora_rank", default=192, type=int)
+ parser.add_argument("--q_lora_rank", default=384, type=int)
+ parser.add_argument("--rope_dim", default=None, type=int)
+ parser.add_argument("--max_total_len", default=2500, type=int)
+ parser.add_argument("--max_new_tokens", default=512, type=int)
+ parser.add_argument("--max_turns", default=4, type=int)
+ parser.add_argument("--tool_timeout", default=3, type=int)
+ parser.add_argument("--temperature", default=0.8, type=float)
+ parser.add_argument("--top_p", default=0.9, type=float)
+ parser.add_argument("--do_sample", default=1, type=int)
+ parser.add_argument("--open_thinking", default=0, type=int)
+ parser.add_argument("--model_profile", type=str, default="")
+ parser.add_argument("--model_profiles_path", type=str, default="configs/searchlm_profiles.json")
+ return parser.parse_args()
+
+
+def main():
+ args = parse_args()
+ apply_eval_model_profile(args)
+ rows = load_agentic_jsonl(args.data)
+ if args.max_samples > 0:
+ rows = rows[: args.max_samples]
+
+ if args.pred:
+ outputs = eval_predictions(args, rows)
+ else:
+ model, tokenizer = load_model(args)
+ outputs = [run_episode(model, tokenizer, sample, args) for sample in rows]
+
+ os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+ os.makedirs(os.path.dirname(args.json_output) or ".", exist_ok=True)
+ os.makedirs(os.path.dirname(args.bad_cases_output) or ".", exist_ok=True)
+ with open(args.output, "w", encoding="utf-8") as f:
+ for item in outputs:
+ f.write(json.dumps(item, ensure_ascii=False) + "\n")
+ metrics = average_agentic_metrics([item["metrics"] for item in outputs])
+ metrics["samples"] = len(outputs)
+ with open(args.json_output, "w", encoding="utf-8") as f:
+ json.dump(metrics, f, ensure_ascii=False, indent=2)
+ with open(args.bad_cases_output, "w", encoding="utf-8") as f:
+ for item in outputs:
+ if item["metrics"].get("task_success", 0) < 1:
+ f.write(json.dumps(item, ensure_ascii=False) + "\n")
+ print(json.dumps(metrics, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/run_general_sft_300m.sh b/scripts/run_general_sft_300m.sh
new file mode 100755
index 0000000..4f4692b
--- /dev/null
+++ b/scripts/run_general_sft_300m.sh
@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+DATA_PATH="${DATA_PATH:-${REPO_ROOT}/dataset/sft_t2t.jsonl}"
+WEIGHT_PATH="${WEIGHT_PATH:-${REPO_ROOT}/out/pretrain_searchlm_1024_mla.pth}"
+SAVE_WEIGHT="${SAVE_WEIGHT:-general_sft_14g}"
+NUM_GPUS="${NUM_GPUS:-7}"
+EPOCHS="${EPOCHS:-1}"
+MICRO_BATCH="${MICRO_BATCH:-1}"
+ACCUMULATION_STEPS="${ACCUMULATION_STEPS:-8}"
+LEARNING_RATE="${LEARNING_RATE:-1e-5}"
+SAVE_INTERVAL="${SAVE_INTERVAL:-500}"
+LOG_INTERVAL="${LOG_INTERVAL:-20}"
+NUM_WORKERS="${NUM_WORKERS:-8}"
+USE_WANDB="${USE_WANDB:-0}"
+FORCE_RESUME="${FORCE_RESUME:-auto}"
+
+if [[ ! -s "${DATA_PATH}" ]]; then
+ echo "SFT dataset not found or empty: ${DATA_PATH}" >&2
+ exit 1
+fi
+
+if [[ ! -s "${WEIGHT_PATH}" ]]; then
+ echo "Pretrained checkpoint not found: ${WEIGHT_PATH}" >&2
+ echo "Expected from_weight=pretrain_searchlm -> out/pretrain_searchlm_1024_mla.pth" >&2
+ exit 1
+fi
+
+if ! command -v deepspeed >/dev/null 2>&1; then
+ echo "deepspeed command not found. Install it in the server training environment first." >&2
+ exit 1
+fi
+
+mkdir -p "${REPO_ROOT}/logs" "${REPO_ROOT}/out" "${REPO_ROOT}/checkpoints"
+
+LATEST_FILE="${REPO_ROOT}/checkpoints/${SAVE_WEIGHT}_1024_mla_latest"
+RESUME=0
+if [[ "${FORCE_RESUME}" == "1" ]]; then
+ RESUME=1
+elif [[ "${FORCE_RESUME}" == "auto" && -s "${LATEST_FILE}" ]]; then
+ RESUME=1
+fi
+
+STAMP="$(date +%Y%m%d_%H%M%S)"
+LOG_FILE="${REPO_ROOT}/logs/${SAVE_WEIGHT}_${STAMP}.log"
+
+echo "repo_root=${REPO_ROOT}"
+echo "data_path=${DATA_PATH}"
+echo "weight_path=${WEIGHT_PATH}"
+echo "save_weight=${SAVE_WEIGHT}"
+echo "num_gpus=${NUM_GPUS}"
+echo "effective_batch=$((NUM_GPUS * MICRO_BATCH * ACCUMULATION_STEPS))"
+echo "resume=${RESUME}"
+echo "log_file=${LOG_FILE}"
+
+WANDB_ARGS=()
+if [[ "${USE_WANDB}" == "1" ]]; then
+ WANDB_ARGS+=(--use_wandb)
+fi
+
+cd "${REPO_ROOT}/trainer"
+
+deepspeed --num_gpus="${NUM_GPUS}" train_full_sft_deepspeed.py \
+ --model_profile SearchLM-300M \
+ --from_weight pretrain_searchlm \
+ --save_weight "${SAVE_WEIGHT}" \
+ --data_path "${DATA_PATH}" \
+ --epochs "${EPOCHS}" \
+ --batch_size "${MICRO_BATCH}" \
+ --accumulation_steps "${ACCUMULATION_STEPS}" \
+ --learning_rate "${LEARNING_RATE}" \
+ --dtype float16 \
+ --num_workers "${NUM_WORKERS}" \
+ --log_interval "${LOG_INTERVAL}" \
+ --save_interval "${SAVE_INTERVAL}" \
+ --from_resume "${RESUME}" \
+ "${WANDB_ARGS[@]}" 2>&1 | tee "${LOG_FILE}"
diff --git a/scripts/validate_agentic_data.py b/scripts/validate_agentic_data.py
new file mode 100644
index 0000000..6850eee
--- /dev/null
+++ b/scripts/validate_agentic_data.py
@@ -0,0 +1,69 @@
+import argparse
+import json
+import os
+import sys
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from agentic.data_analysis_env import score_agentic_trajectory
+from dataset.agentic_dataset import load_agentic_jsonl
+
+
+REQUIRED_FIELDS = ["id", "question", "tables", "documents", "tools", "expected_tools", "answer", "checks"]
+
+
+def trajectory_to_turn_outputs(sample):
+ turns = []
+ for item in sample.get("expert_trajectory", []):
+ if item.get("role") == "assistant" and item.get("tool_call"):
+ turns.append("\n" + json.dumps(item["tool_call"], ensure_ascii=False) + "\n")
+ elif item.get("role") == "assistant":
+ turns.append(str(item.get("content", "")))
+ return turns
+
+
+def validate_sample(sample):
+ errors = []
+ for field in REQUIRED_FIELDS:
+ if field not in sample:
+ errors.append(f"missing field: {field}")
+ for table in sample.get("tables", []):
+ path = table.get("path", "")
+ if not path:
+ errors.append("table path is empty")
+ elif not os.path.exists(path):
+ errors.append(f"table path not found: {path}")
+ if not sample.get("expert_trajectory"):
+ errors.append("missing expert_trajectory")
+ if not errors:
+ reward, parts = score_agentic_trajectory(trajectory_to_turn_outputs(sample), sample)
+ if parts.get("task_success", 0) < 1:
+ errors.append(f"expert trajectory not successful: reward={reward:.3f}, parts={parts}")
+ return errors
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Validate Agentic DataAnalysis JSONL")
+ parser.add_argument("--data", required=True)
+ parser.add_argument("--max_errors", type=int, default=20)
+ return parser.parse_args()
+
+
+def main():
+ args = parse_args()
+ rows = load_agentic_jsonl(args.data)
+ total_errors = 0
+ for row in rows:
+ errors = validate_sample(row)
+ if errors:
+ total_errors += len(errors)
+ if total_errors <= args.max_errors:
+ print(json.dumps({"id": row.get("id"), "errors": errors}, ensure_ascii=False))
+ summary = {"data": args.data, "samples": len(rows), "errors": total_errors}
+ print(json.dumps(summary, ensure_ascii=False))
+ if total_errors:
+ raise SystemExit(1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/trainer/train_agent_deepspeed.py b/trainer/train_agent_deepspeed.py
new file mode 100644
index 0000000..65a7645
--- /dev/null
+++ b/trainer/train_agent_deepspeed.py
@@ -0,0 +1,576 @@
+"""
+DeepSpeed Agentic DataAnalysis GRPO/CISPO 入口。
+
+推荐在 trainer/ 目录启动：
+ deepspeed --num_gpus=7 train_agent_deepspeed.py --model_profile SearchLM-300M
+"""
+import os
+import sys
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+import argparse
+import json
+import math
+import random
+import warnings
+from contextlib import nullcontext
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, DistributedSampler
+
+from agentic.data_analysis_env import average_agentic_metrics, parse_tool_calls, score_agentic_trajectory
+from dataset.agentic_dataset import AgenticRLDataset
+from trainer.rollout_engine import compute_per_token_logps, create_rollout_engine
+from trainer.trainer_utils import (
+ Logger,
+ SkipBatchSampler,
+ add_deepspeed_args,
+ add_model_profile_args,
+ apply_model_profile,
+ assert_no_training_output_collision,
+ build_lm_config,
+ get_cuda_peak_memory_gb,
+ init_model,
+ is_main_process,
+ load_deepspeed_checkpoint,
+ load_deepspeed_config,
+ log_training_setup,
+ reduce_metrics,
+ resolve_attention_type,
+ save_deepspeed_checkpoint,
+ save_model_weights,
+ setup_seed,
+)
+
+warnings.filterwarnings("ignore")
+
+
+def init_deepspeed_runtime(args, deepspeed):
+ deepspeed.init_distributed(dist_backend="nccl")
+ local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank if args.local_rank >= 0 else 0))
+ if torch.cuda.is_available():
+ torch.cuda.set_device(local_rank)
+ args.device = f"cuda:{local_rank}"
+ else:
+ args.device = "cpu"
+ args.training_backend = "DeepSpeed-ZeRO"
+ setup_seed(42 + (dist.get_rank() if dist.is_initialized() else local_rank))
+ return local_rank
+
+
+def collate_fn(batch):
+ return {
+ "id": [x["id"] for x in batch],
+ "messages": [x["messages"] for x in batch],
+ "tools": [x["tools"] for x in batch],
+ "sample": [x["sample"] for x in batch],
+ }
+
+
+def rollout_single(args, rollout_engine, tokenizer, messages, tools, sample):
+ all_outputs = []
+ prompt_ids = None
+ response_ids = []
+ response_mask = []
+ response_old_logps = []
+ final_context = ""
+ unfinished = False
+ open_thinking = random.random() < args.thinking_ratio
+
+ for turn in range(args.max_turns):
+ context = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ tools=tools,
+ open_thinking=open_thinking,
+ )
+ inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(args.device)
+ context_ids = inputs["input_ids"][0].tolist()
+ if prompt_ids is None:
+ prompt_ids = context_ids
+
+ rollout_result = rollout_engine.rollout(
+ prompt_ids=inputs["input_ids"],
+ attention_mask=inputs["attention_mask"],
+ num_generations=1,
+ max_new_tokens=args.max_gen_len,
+ temperature=args.temperature,
+ )
+ new_ids = rollout_result.completion_ids[0].tolist()
+ new_logps = rollout_result.per_token_logps[0].tolist()
+ pairs = [
+ (token_id, logp)
+ for token_id, logp in zip(new_ids, new_logps)
+ if token_id != tokenizer.pad_token_id and token_id != tokenizer.eos_token_id
+ ]
+ new_ids = [token_id for token_id, _ in pairs]
+ new_logps = [logp for _, logp in pairs]
+ new_text = rollout_result.completions[0]
+ all_outputs.append(new_text)
+ response_ids.extend(new_ids)
+ response_mask.extend([1] * len(new_ids))
+ response_old_logps.extend(new_logps)
+ final_context = context + new_text
+
+ calls = parse_tool_calls(new_text)
+ if not calls:
+ break
+
+ unfinished = turn == args.max_turns - 1
+ messages.append({"role": "assistant", "content": new_text})
+ from agentic.data_analysis_env import AgenticToolEnv
+
+ env = AgenticToolEnv(sample, repo_root=args.repo_root, timeout=args.tool_timeout)
+ for call in calls:
+ result = env.execute(call.get("name", ""), call.get("arguments", {}))
+ result_str = json.dumps(result, ensure_ascii=False)[: args.max_observation_chars]
+ messages.append({"role": "tool", "content": result_str})
+
+ observe_context = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=not unfinished,
+ tools=tools,
+ open_thinking=open_thinking,
+ )
+ observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
+ current_len = len(prompt_ids) + len(response_ids)
+ obs_delta = observe_ids[current_len:]
+ response_ids.extend(obs_delta)
+ response_mask.extend([0] * len(obs_delta))
+ response_old_logps.extend([0.0] * len(obs_delta))
+ final_context = observe_context
+ if unfinished:
+ break
+
+ return (
+ final_context,
+ prompt_ids or [],
+ response_ids,
+ response_mask,
+ response_old_logps,
+ all_outputs,
+ unfinished,
+ )
+
+
+def rollout_batch(args, rollout_engine, tokenizer, messages_batch, tools_batch, samples):
+ contexts = []
+ prompt_ids_batch = []
+ response_ids_batch = []
+ response_masks_batch = []
+ response_old_logps_batch = []
+ turn_outputs_batch = []
+ unfinished_batch = []
+ for messages, tools, sample in zip(messages_batch, tools_batch, samples):
+ for _ in range(args.num_generations):
+ messages_copy = [dict(message) for message in messages]
+ context, prompt_ids, response_ids, response_mask, old_logps, turn_outputs, unfinished = rollout_single(
+ args,
+ rollout_engine,
+ tokenizer,
+ messages_copy,
+ tools,
+ sample,
+ )
+ contexts.append(context)
+ prompt_ids_batch.append(prompt_ids)
+ response_ids_batch.append(response_ids)
+ response_masks_batch.append(response_mask)
+ response_old_logps_batch.append(old_logps)
+ turn_outputs_batch.append(turn_outputs)
+ unfinished_batch.append(unfinished)
+ return contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch
+
+
+def calculate_agentic_rewards(args, turn_outputs_batch, samples, unfinished_batch):
+ rewards = []
+ metrics = []
+ for idx, turn_outputs in enumerate(turn_outputs_batch):
+ sample_idx = idx // args.num_generations
+ reward, parts = score_agentic_trajectory(
+ turn_outputs,
+ samples[sample_idx],
+ repo_root=args.repo_root,
+ unfinished=unfinished_batch[idx],
+ execute_tools=True,
+ )
+ rewards.append(reward)
+ metrics.append(parts)
+ return torch.tensor(rewards, dtype=torch.float32, device=args.device), average_agentic_metrics(metrics)
+
+
+def pack_rollout_samples(args, tokenizer, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch):
+ packed = []
+ for prompt_ids, response_ids, response_mask, old_logps in zip(
+ prompt_ids_batch,
+ response_ids_batch,
+ response_masks_batch,
+ response_old_logps_batch,
+ ):
+ ids = prompt_ids + response_ids
+ mask = [0] * len(prompt_ids) + response_mask
+ old_per_token = [0.0] * max(len(prompt_ids) - 1, 0) + old_logps
+ if len(ids) > args.max_total_len:
+ ids = ids[-args.max_total_len:]
+ mask = mask[-args.max_total_len:]
+ old_per_token = old_per_token[-(len(ids) - 1):]
+ target_len = max(len(ids) - 1, 0)
+ if len(old_per_token) < target_len:
+ old_per_token += [0.0] * (target_len - len(old_per_token))
+ elif len(old_per_token) > target_len:
+ old_per_token = old_per_token[-target_len:]
+ packed.append((ids, mask, old_per_token))
+
+ max_len = max(len(ids) for ids, _, _ in packed)
+ input_ids = torch.tensor(
+ [ids + [tokenizer.pad_token_id] * (max_len - len(ids)) for ids, _, _ in packed],
+ device=args.device,
+ dtype=torch.long,
+ )
+ full_response_masks = torch.tensor(
+ [mask + [0] * (max_len - len(mask)) for _, mask, _ in packed],
+ device=args.device,
+ dtype=torch.float32,
+ )
+ old_per_token_logps = torch.tensor(
+ [old + [0.0] * ((max_len - 1) - len(old)) for _, _, old in packed],
+ device=args.device,
+ dtype=torch.float32,
+ )
+ return input_ids, full_response_masks, old_per_token_logps
+
+
+def train_epoch(args, lm_config, model_engine, tokenizer, loader, iters, rollout_engine, ref_model, start_step=0, wandb=None):
+ for step, batch in enumerate(loader, start=start_step + 1):
+ contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch = rollout_batch(
+ args,
+ rollout_engine,
+ tokenizer,
+ batch["messages"],
+ batch["tools"],
+ batch["sample"],
+ )
+ input_ids, full_response_masks, old_per_token_logps = pack_rollout_samples(
+ args,
+ tokenizer,
+ prompt_ids_batch,
+ response_ids_batch,
+ response_masks_batch,
+ response_old_logps_batch,
+ )
+ full_mask = (input_ids != tokenizer.pad_token_id).long()
+ rewards, reward_stats = calculate_agentic_rewards(args, turn_outputs_batch, batch["sample"], unfinished_batch)
+
+ res = model_engine(input_ids, attention_mask=full_mask)
+ aux_loss = res.aux_loss if getattr(lm_config, "use_moe", False) and res.aux_loss is not None else torch.tensor(0.0, device=args.device)
+ per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
+
+ with torch.no_grad():
+ ref_per_token_logps = compute_per_token_logps(
+ ref_model,
+ input_ids,
+ input_ids.size(1) - 1,
+ attention_mask=full_mask,
+ )
+
+ completion_mask = full_response_masks[:, 1:]
+ token_counts = completion_mask.sum(dim=1)
+ valid_rows = token_counts > 0
+
+ grouped_rewards = rewards.view(-1, args.num_generations)
+ mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)
+ std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)
+ advantages = (rewards - mean_r) / (std_r + 1e-4)
+
+ kl_div = ref_per_token_logps - per_token_logps
+ per_token_kl = torch.exp(kl_div) - kl_div - 1
+ ratio = torch.exp(per_token_logps - old_per_token_logps)
+ if args.loss_type == "cispo":
+ clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
+ per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
+ else:
+ clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
+ per_token_loss1 = ratio * advantages.unsqueeze(1)
+ per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
+ per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
+ policy_loss = (
+ ((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
+ if valid_rows.any()
+ else per_token_loss.sum() * 0.0
+ )
+ loss = policy_loss + aux_loss
+
+ model_engine.backward(loss)
+ model_engine.step()
+
+ if step % args.log_interval == 0 or step == iters:
+ reward_val = rewards.mean().item()
+ avg_len_val = token_counts.float().mean().item()
+ kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(token_counts.sum().item(), 1)
+ lr_list = model_engine.get_lr()
+ current_lr = lr_list[0] if lr_list else args.learning_rate
+ metrics = reduce_metrics(
+ {
+ "reward": reward_val,
+ "avg_len": avg_len_val,
+ "kl_ref": kl_ref_val,
+ "policy_loss": policy_loss.item(),
+ "aux_loss": aux_loss.item(),
+ **reward_stats,
+ }
+ )
+ peak_memory_gb = get_cuda_peak_memory_gb()
+ memory_text = f", peak_memory: {peak_memory_gb:.2f} GB" if peak_memory_gb is not None else ""
+ Logger(
+ f"Epoch:[{args.current_epoch + 1}/{args.epochs}]({step}/{iters}), "
+ f"Reward:{metrics['reward']:.4f}, Success:{metrics.get('task_success', 0):.4f}, "
+ f"ToolF1:{metrics.get('tool_selection_f1', 0):.4f}, Invalid:{metrics.get('invalid_action_rate', 0):.4f}, "
+ f"Exec:{metrics.get('exec_success_rate', 0):.4f}, Turns:{metrics.get('avg_turns', 0):.2f}, "
+ f"KL:{metrics['kl_ref']:.4f}, Loss:{metrics['policy_loss']:.4f}, LR:{current_lr:.8f}"
+ f"{memory_text}"
+ )
+ if wandb and is_main_process():
+ log_data = {
+ "reward": metrics["reward"],
+ "kl_ref": metrics["kl_ref"],
+ "policy_loss": metrics["policy_loss"],
+ "aux_loss": metrics["aux_loss"],
+ "avg_response_len": metrics["avg_len"],
+ "learning_rate": current_lr,
+ }
+ log_data.update({key: value for key, value in metrics.items() if key not in log_data})
+ wandb.log(log_data)
+
+ if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+ for i, sample_id in enumerate(batch["id"][:2]):
+ Logger(f"[DEBUG] sample={sample_id}")
+ for j in range(args.num_generations):
+ idx = i * args.num_generations + j
+ Logger(f"[DEBUG] gen={j}, reward={rewards[idx].item():.4f}, unfinished={unfinished_batch[idx]}")
+ Logger("\n---\n".join(turn_outputs_batch[idx]))
+ Logger("=" * 80)
+
+ if step % args.save_interval == 0 or step == iters:
+ model_engine.eval()
+ save_deepspeed_checkpoint(
+ model_engine,
+ lm_config,
+ weight=args.save_weight,
+ epoch=args.current_epoch,
+ step=step,
+ wandb=wandb,
+ save_dir="../checkpoints",
+ )
+ if is_main_process():
+ ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
+ Logger(f"Lightweight model saved: {ckp}")
+ model_engine.train()
+ rollout_engine.update_policy(model_engine)
+
+ del contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch
+ del turn_outputs_batch, unfinished_batch, input_ids, full_response_masks, old_per_token_logps
+ del per_token_logps, ref_per_token_logps, rewards, grouped_rewards, mean_r, std_r, advantages
+ del completion_mask, token_counts, res, loss, policy_loss, aux_loss
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Agentic DataAnalysis GRPO/CISPO with DeepSpeed ZeRO")
+ parser.add_argument("--save_dir", type=str, default="../out")
+ parser.add_argument("--save_weight", default="agent_grpo", type=str)
+ parser.add_argument("--epochs", type=int, default=1)
+ parser.add_argument("--batch_size", type=int, default=1)
+ parser.add_argument("--learning_rate", type=float, default=3e-7)
+ parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
+ parser.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16", "float32"])
+ parser.add_argument("--num_workers", type=int, default=4)
+ parser.add_argument("--accumulation_steps", type=int, default=4)
+ parser.add_argument("--grad_clip", type=float, default=1.0)
+ parser.add_argument("--log_interval", type=int, default=1)
+ parser.add_argument("--save_interval", type=int, default=25)
+ parser.add_argument("--hidden_size", default=768, type=int)
+ parser.add_argument("--num_hidden_layers", default=8, type=int)
+ parser.add_argument("--num_attention_heads", default=8, type=int)
+ parser.add_argument("--num_key_value_heads", default=4, type=int)
+ parser.add_argument("--intermediate_size", default=None, type=int)
+ parser.add_argument("--use_moe", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--attention_type", default="gqa", choices=["gqa", "mha", "mqa", "mla"])
+ parser.add_argument("--use_mla", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--kv_lora_rank", default=128, type=int)
+ parser.add_argument("--q_lora_rank", default=256, type=int)
+ parser.add_argument("--rope_dim", default=None, type=int)
+ parser.add_argument("--max_seq_len", default=1536, type=int)
+ parser.add_argument("--max_gen_len", type=int, default=512)
+ parser.add_argument("--max_total_len", type=int, default=2500)
+ parser.add_argument("--max_turns", type=int, default=4)
+ parser.add_argument("--max_observation_chars", type=int, default=2048)
+ parser.add_argument("--tool_timeout", type=int, default=3)
+ parser.add_argument("--repo_root", type=str, default="..")
+ parser.add_argument("--data_path", type=str, default="../dataset/agentic_data/train.jsonl")
+ parser.add_argument("--num_generations", type=int, default=6)
+ parser.add_argument("--beta", type=float, default=0.05)
+ parser.add_argument("--loss_type", type=str, default="cispo", choices=["grpo", "cispo"])
+ parser.add_argument("--epsilon", type=float, default=0.2)
+ parser.add_argument("--epsilon_high", type=float, default=5.0)
+ parser.add_argument("--temperature", type=float, default=0.8)
+ parser.add_argument("--thinking_ratio", type=float, default=0.1)
+ parser.add_argument("--from_weight", default="agent_sft", type=str)
+ parser.add_argument("--from_resume", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--overwrite_output", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--use_wandb", action="store_true")
+ parser.add_argument("--wandb_project", type=str, default="Agentic-DataAnalysis-RL-DS")
+ parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1])
+ parser.add_argument("--debug_mode", action="store_true")
+ parser.add_argument("--debug_interval", type=int, default=20)
+ parser.add_argument("--rollout_engine", type=str, default="torch", choices=["torch", "sglang"])
+ parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998")
+ parser.add_argument("--sglang_model_path", type=str, default="../model")
+ parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agentic")
+ add_model_profile_args(parser)
+ add_deepspeed_args(parser)
+ return parser.parse_args()
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ model_profile = apply_model_profile(args)
+
+ try:
+ import deepspeed
+ except ImportError as exc:
+ raise RuntimeError("DeepSpeed未安装。服务器环境执行: pip install deepspeed") from exc
+
+ local_rank = init_deepspeed_runtime(args, deepspeed)
+ os.makedirs(args.save_dir, exist_ok=True)
+
+ lm_config = build_lm_config(
+ 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),
+ use_moe=bool(args.use_moe),
+ attention_type=resolve_attention_type(args),
+ kv_lora_rank=args.kv_lora_rank,
+ q_lora_rank=args.q_lora_rank,
+ rope_dim=args.rope_dim,
+ num_attention_heads=args.num_attention_heads,
+ num_key_value_heads=args.num_key_value_heads,
+ intermediate_size=args.intermediate_size,
+ )
+ assert_no_training_output_collision(
+ lm_config,
+ weight=args.save_weight,
+ save_dir=args.save_dir,
+ checkpoint_dir="../checkpoints",
+ from_resume=args.from_resume,
+ overwrite=args.overwrite_output,
+ )
+
+ model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+ if args.use_compile == 1:
+ model = torch.compile(model)
+ Logger("torch.compile enabled before DeepSpeed initialize")
+
+ train_ds = AgenticRLDataset(args.data_path, tokenizer, max_length=args.max_total_len)
+ train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+ loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler, collate_fn=collate_fn)
+ iters = len(loader_for_count)
+ total_optimizer_steps = max(1, math.ceil(iters / args.accumulation_steps) * args.epochs)
+ ds_config = load_deepspeed_config(args, total_optimizer_steps)
+
+ model_engine, _, _, _ = deepspeed.initialize(
+ model=model,
+ model_parameters=model.parameters(),
+ config=ds_config,
+ )
+
+ ckp_data = None
+ if args.from_resume == 1:
+ ckp_data = load_deepspeed_checkpoint(
+ model_engine,
+ lm_config,
+ weight=args.save_weight,
+ save_dir="../checkpoints",
+ )
+
+ ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
+ ref_model = ref_model.eval().requires_grad_(False)
+ if torch.cuda.is_available() and args.dtype in {"float16", "bfloat16"}:
+ ref_dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
+ ref_model = ref_model.to(dtype=ref_dtype)
+
+ device_type = "cuda" if "cuda" in args.device else "cpu"
+ amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+ autocast_ctx = nullcontext() if device_type == "cpu" or args.dtype == "float32" else torch.cuda.amp.autocast(dtype=amp_dtype)
+
+ rollout_engine = create_rollout_engine(
+ engine_type=args.rollout_engine,
+ policy_model=model_engine,
+ tokenizer=tokenizer,
+ device=args.device,
+ autocast_ctx=autocast_ctx,
+ sglang_base_url=args.sglang_base_url,
+ sglang_model_path=args.sglang_model_path,
+ sglang_shared_path=args.sglang_shared_path,
+ )
+ rollout_engine.update_policy(model_engine)
+
+ wandb = None
+ if args.use_wandb and is_main_process():
+ import swanlab as wandb
+
+ wandb_id = ckp_data.get("wandb_id") if ckp_data else None
+ wandb.init(
+ project=args.wandb_project,
+ name=f"Agentic-RL-DS-{args.attention_type}-{args.hidden_size}",
+ id=wandb_id,
+ resume="must" if wandb_id else None,
+ )
+
+ log_training_setup(
+ args,
+ lm_config,
+ stage="agentic_rl_deepspeed",
+ dataset_len=len(train_ds),
+ iters=iters,
+ tokens_per_sample=args.max_total_len,
+ extra={
+ "model_profile": args.model_profile if model_profile else "",
+ "deepspeed_config": args.deepspeed_config,
+ "zero_stage": ds_config.get("zero_optimization", {}).get("stage", "none"),
+ "num_generations": args.num_generations,
+ "rollout_engine": args.rollout_engine,
+ "loss_type": args.loss_type,
+ "beta": args.beta,
+ "max_turns": args.max_turns,
+ "total_optimizer_steps": total_optimizer_steps,
+ },
+ )
+
+ start_epoch = ckp_data.get("epoch", 0) if ckp_data else 0
+ start_step = ckp_data.get("step", 0) if ckp_data else 0
+
+ for epoch in range(start_epoch, args.epochs):
+ args.current_epoch = epoch
+ train_sampler and train_sampler.set_epoch(epoch)
+ setup_seed(42 + epoch + local_rank)
+ indices = torch.randperm(len(train_ds)).tolist()
+ skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+ batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+ loader = DataLoader(
+ train_ds,
+ batch_sampler=batch_sampler,
+ num_workers=args.num_workers,
+ pin_memory=True,
+ collate_fn=collate_fn,
+ )
+ if skip > 0:
+ Logger(f"Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始")
+ train_epoch(args, lm_config, model_engine, tokenizer, loader, len(loader) + skip, rollout_engine, ref_model, skip, wandb)
+
+ if dist.is_initialized():
+ dist.destroy_process_group()
diff --git a/trainer/train_full_sft_deepspeed.py b/trainer/train_full_sft_deepspeed.py
index 567a8d6..cef0b88 100644
--- a/trainer/train_full_sft_deepspeed.py
+++ b/trainer/train_full_sft_deepspeed.py
@@ -71,6 +71,29 @@ def validate_sft_supervision(train_ds, max_samples=32):
)
+def validate_chat_template(tokenizer):
+ try:
+ rendered = tokenizer.apply_chat_template(
+ [
+ {"role": "user", "content": "template check"},
+ {"role": "assistant", "content": "ok"},
+ ],
+ tokenize=False,
+ add_generation_prompt=False,
+ )
+ except Exception as exc:
+ raise RuntimeError(
+ "Tokenizer chat template is invalid. Check model/chat_template.jinja and "
+ "model/tokenizer_config.json before loading the SFT dataset."
+ ) from exc
+ required_markers = ("<|im_start|>user", "<|im_start|>assistant", "<|im_end|>")
+ missing = [marker for marker in required_markers if marker not in rendered]
+ if missing:
+ raise RuntimeError(f"Tokenizer chat template is missing required markers: {missing}")
+ if is_main_process():
+ Logger("Tokenizer chat template validation passed")
+
+
def init_deepspeed_runtime(args, deepspeed):
deepspeed.init_distributed(dist_backend="nccl")
local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank if args.local_rank >= 0 else 0))
@@ -129,7 +152,12 @@ def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wand
}
)
- if step % args.save_interval == 0 or step == iters:
+ should_save = step % args.save_interval == 0 or step == iters
+ del input_ids, labels, res, loss, aux_loss
+
+ if should_save:
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
model_engine.eval()
save_deepspeed_checkpoint(
model_engine,
@@ -144,10 +172,13 @@ def train_epoch(args, lm_config, model_engine, loader, iters, start_step=0, wand
if is_main_process():
ckp = save_model_weights(lm_config, model_engine, save_dir=args.save_dir, weight=args.save_weight)
Logger(f"Lightweight model saved: {ckp}")
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ # rank0 导出轻量权重时，其他 rank 必须等待，避免提前销毁 NCCL process group。
+ if dist.is_initialized():
+ dist.barrier()
model_engine.train()
- del input_ids, labels, res, loss, aux_loss
-
def parse_args():
parser = argparse.ArgumentParser(description="Full SFT with DeepSpeed ZeRO")
@@ -222,6 +253,7 @@ if __name__ == "__main__":
)
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+ validate_chat_template(tokenizer)
if args.use_compile == 1:
model = torch.compile(model)
Logger("torch.compile enabled before DeepSpeed initialize")