agentic

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",
+]

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 数据进行计算；最终答案必须简洁，"
+    "包含关键结论和必要数值。工具调用必须使用 <tool_call> 包裹的 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"<tool_call>(.*?)</tool_call>", 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"<think>.*?</think>", "", 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 "</tool_call>" in text:
+        text = text.split("</tool_call>")[-1]
+    text = re.sub(r"<tool_call>.*?</tool_call>", "", 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, "<agentic_python_executor>", "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("<tool_call>") - str(turn).count("</tool_call>"))
+        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
+    }

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,
+        }

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`：`<tool_call>` 标签和 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 执行成功率和平均交互轮数。

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
+    {{- '# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>' }}
-
-You may call one or more functions to assist with the user query.
-
-You are provided with function signatures within <tools></tools> XML tags:
-<tools>" }}
     {%- for tool in tools %}
-        {{- "
+        {{- '\n' }}
-" }}
         {{- tool | tojson }}
     {%- endfor %}
-    {{- "
+    {{- '\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}\n</tool_call><|im_end|>\n' }}
-</tools>
-
-For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
-<tool_call>
-{"name": <function-name>, "arguments": <args-json-object>}
-</tool_call><|im_end|>
-" }}
 {%- else %}
     {%- if messages[0].role == 'system' %}
-        {{- '<|im_start|>system
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
-' + messages[0].content + '<|im_end|>
-' }}
     {%- 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('<tool_response>') and message.content.endswith('</tool_response>')) %}
-        {%- 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) %}
+    {%- if message.role == 'user' or (message.role == 'system' and not loop.first) %}
-        {{- '<|im_start|>' + message.role + '
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>\n' }}
-' + content + '<|im_end|>' + '
+    {%- elif message.role == 'assistant' %}
-' }}
-    {%- elif message.role == "assistant" %}
         {%- set reasoning_content = '' %}
         {%- if message.reasoning_content is string %}
             {%- set reasoning_content = message.reasoning_content %}
-        {%- else %}
+        {%- elif '</think>' in content %}
-            {%- if '</think>' in content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
-                {%- set reasoning_content = content.split('</think>')[0].rstrip('
+            {%- set content = content.split('</think>')[-1].lstrip('\n') %}
-').split('<think>')[-1].lstrip('
-') %}
-                {%- set content = content.split('</think>')[-1].lstrip('
-') %}
-            {%- endif %}
-        {%- endif %}
-        {%- if true %}
-            {{- '<|im_start|>' + message.role + '
-<think>
-' + reasoning_content.strip('
-') + '
-</think>
-
-' + content.lstrip('
-') }}
         {%- endif %}
+        {{- '<|im_start|>assistant\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\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 %}
-                {{- '<tool_call>
+                {{- '<tool_call>\n{"name": "' + tool_call.name + '", "arguments": ' }}
-{"name": "' }}
-                {{- tool_call.name }}
-                {{- '", "arguments": ' }}
                 {%- if tool_call.arguments is string %}
                     {{- tool_call.arguments }}
                 {%- else %}
                     {{- tool_call.arguments | tojson }}
                 {%- endif %}
-                {{- '}
+                {{- '}\n</tool_call>' }}
-</tool_call>' }}
             {%- endfor %}
         {%- endif %}
-        {{- '<|im_end|>
+        {{- '<|im_end|>\n' }}
-' }}
+    {%- elif message.role == 'tool' %}
-    {%- elif message.role == "tool" %}
+        {%- if loop.first or messages[loop.index0 - 1].role != 'tool' %}
-        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
             {{- '<|im_start|>user' }}
         {%- endif %}
-        {{- '
+        {{- '\n<tool_response>\n' + content + '\n</tool_response>' }}
-<tool_response>
+        {%- if loop.last or messages[loop.index0 + 1].role != 'tool' %}
-' }}
+            {{- '<|im_end|>\n' }}
-        {{- content }}
-        {{- '
-</tool_response>' }}
-        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
-            {{- '<|im_end|>
-' }}
         {%- endif %}
     {%- endif %}
 {%- endfor %}
 {%- if add_generation_prompt %}
-    {{- '<|im_start|>assistant
+    {{- '<|im_start|>assistant\n' }}
-' }}
     {%- if open_thinking is defined and open_thinking is true %}
-        {{- '<think>
+        {{- '<think>\n' }}
-' }}
     {%- else %}
-        {{- '<think>
+        {{- '<think>\n\n</think>\n\n' }}
-
-</think>
-
-' }}
     {%- endif %}
 {%- endif %}

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

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("<tool_call>\n" + json.dumps(item["tool_call"], ensure_ascii=False, indent=2) + "\n</tool_call>")
+        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 "<tool_response>" 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()

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()

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()

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()

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}"

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("<tool_call>\n" + json.dumps(item["tool_call"], ensure_ascii=False) + "\n</tool_call>")
+        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()

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()

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")