完全 from-scratch 构建一个能调工具的 LLM Agent

1. 引言

1.1 动机

现代 LLM 工程是一个抽象的塔——transformer 架构 → 训练框架 → tokenizer → 推理服务 → tool use protocol → web 前端。每一层都被高度优化的开源库覆盖（PyTorch、HuggingFace transformers、tiktoken、vllm、LangChain），单从用户视角"调一下 API 就行"。

但这种封装让系统成为黑盒。一个 ML 工程师在被问到"为什么 KV cache 让 decode 复杂度从 O(N²) 降到 O(N)？"或者"为什么 flash-attention 在长 context 下加速十倍？"或者"agent 调用 tool 时模型其实在 emit 什么 token？"时，往往说不清细节。这些问题的答案不在论文里——在代码里。

1.2 项目目标

构建一个可以独立运行的 LLM 全栈系统，满足以下约束：

1.3 与既有教程的关系

业界已有许多优秀的"从零写 GPT"教程，包括 karpathy/nanoGPT（预训练 + 模型）、karpathy/minbpe（BPE 算法）、lucidrains 的 transformer 实现。这些聚焦于模型本身。本项目的差异化贡献是把链路从模型扩展到端到端 stack：包含推理服务的 KV cache、agent 循环、web SSE 流式输出、CUDA/Triton kernel 对比——这些在教学项目中通常缺失。

2. 系统设计

2.1 整体架构

8 个独立但可协作的层：

模型结构
L1   05_gpu/           CUDA matmul + Triton flash-attention
L2   04_transformer/   GPT-2 架构 + 手写 BPE

模型训练
L3   00_train/         预训练循环（AdamW + cosine schedule，140 行）
L4   00b_sft/          指令 SFT（loss masking on prompts）
L5   00c_agent_sft/    Agent SFT（ReAct 格式 + 工具）

推理与应用
L6   03_model/         推理服务（自实现 KV cache）
L7   01_app/           App / Web UI（FastAPI + SSE）
L8   02_agent/         HTTP chat 客户端 + ReAct agent loop

每层都有独立的 README（讲为什么 + 怎么做 + 实测数字），可单独启动验证。

2.2 关键设计决策

(a) Tokenizer 用 OpenAI 公开 vocab，但 BPE 算法手写。vocab 文件（encoder.json + vocab.bpe）是 OpenAI 在 GPT-2 时代公开发布的数据；BPE 编解码算法是我们 230 行 Python 重写，并验证 bit-for-bit 等价 tiktoken。这样保持与 OpenAI 模型权重兼容，但算法完全自家实现。

(b) 训练数据 bundle 进 repo。1.1 MB 的 Tiny Shakespeare 直接 commit 到 git，避免运行时网络依赖。生成的 train.bin / val.bin 是 .gitignore 的，每次 prepare.py 1 秒重生成。

(c) HF mirror auto-fallback。GPT.from_pretrained("gpt2") 内部 probe huggingface.co，3 秒超时则 fallback 到 https://hf-mirror.com + 设 HF_HUB_DISABLE_XET=1。这让中国大陆地区无须任何手动配置就能跑通需要 OpenAI 权重的路径。

(d) Agent SFT 的 loss masking。详见第 4 节。这是项目中唯一一处需要"想清楚"的训练数据细节——把 OBSERVATION: 的 prefix 标为"学"、内容标为"mask"，使模型既能 emit 该 token 作为 stop signal，又不会幻觉工具输出。

2.3 贯穿示例

整个系统用一个具体 query "What is 1234 plus 5678?" 作为贯穿测试。这个数对 SFT 训练数据外（仅训练过 1-99 范围的加法）。期望流程：

用户：What is 1234 plus 5678?
模型：💭 I need to compute 1234 + 5678.        (THOUGHT)
模型：🔧 calc(1234 + 5678)                       (ACTION)
工具：6912                                       (OBSERVATION，由 Python 计算)
模型：6912.                                      (ANSWER)

这个 demo 同时验证：

• Tokenizer 能正确编码 "1234 plus 5678"
• 推理服务能流式吐 token
• Agent loop 能截取 ACTION 行 + 调真工具 + 注入 OBSERVATION
• Model 学到了 ReAct 格式契约
• 端到端通过浏览器 SSE 渲染

3. 训练方法

3.1 L3 预训练

初始化：N(0, 0.02) for all Linear / Embedding；residual c_proj 缩 1/sqrt(2N)（GPT-2 标准做法）。这使 step 0 loss = 10.815，与理论值 ln(50,257) = 10.825 一致。早期版本未做正确初始化时初始 loss = 80（softmax 严重偏向某些 token，cross-entropy 巨大）。

3.2 L4 SFT

两条独立路径：

• Path A：在 L3 自训的 7M base 上 SFT。
• Path B：在 OpenAI 公开的 gpt2-124M base 上 SFT（通过 GPT.from_pretrained 加载权重并 reshape 到我们的架构）。

Loss masking：每条 (Q, A) 样本中，仅 A 的 token 计 loss（prompt token 用 ignore_index=-1 mask 掉）。

3.3 L5 Agent SFT

在 L4 path-B（124M instruction-tuned）基础上继续 SFT，使用 ReAct 格式：

Q: <question>
THOUGHT: <reasoning>
ACTION: tool(args)
OBSERVATION: <tool output>
ANSWER: <final><|endoftext|>

数据生成：258 条 traces 由 build_data.py 程序合成，从 105 条 KB 事实 + 85 条算术问题生成。observation 字段在生成时调用真工具产生，保证训练数据与运行时输出完全一致。

关键 Loss masking 技术

parts = []
parts.append((f"Q: {q}\n",                   False))   # 用户输入
parts.append((f"THOUGHT: {thought}\n",       True))    # model emit
parts.append((f"ACTION: {action}\n",         True))    # model emit
parts.append(("OBSERVATION:",                True))    # model emit (stop signal)
parts.append((f" {observation}\n",           False))   # tool output, 不让 model 幻觉
parts.append((f"ANSWER: {answer}",           True))    # model emit
parts.append((EOT_ID,                        True))    # 学会停止

OBSERVATION: prefix 必须是"学"——这是模型告知 agent loop "我要工具结果" 的 stop signal。如果整段 OBSERVATION 都 mask，模型不会 emit 该 prefix，会一直循环 emit 更多 ACTION（我们在第一版迭代中观察到这个 bug）。

3.4 训练结果

实测 RTX 5090 (server #3, 2026-05-03)：

训练总用时（不含一次性的 124M 权重下载）：~100 秒。

4. 推理与服务

4.1 KV Cache 实现

L2 的 GPT 类支持两种 forward：

# 训练 / 简单推理：
logits = model(input_ids)
logits, loss = model(input_ids, targets=y)

# 流式推理（带 KV cache）：
logits, kv_caches = model.step(input_ids)              # prefill
logits, kv_caches = model.step(next_id, kv_caches)     # decode

step() 方法在每个 attention 层把当前 K/V 与缓存的 past K/V 拼接，使 decode step 复杂度从 O(N²) 降到 O(N)。

数值正确性验证

# Test: full forward (T tokens at once) ≡ step() with prefill + 1 decode
with torch.no_grad():
    logits_full = m(full_input)
    logits_pre, kvs = m.step(full_input[:, :T-1])
    logits_dec, _ = m.step(full_input[:, T-1:T], kvs)

diff = (logits_full[0, -1] - logits_dec[0, -1]).abs().max()
# Result: 2.4e-7  (浮点累加误差，可接受)

4.2 推理服务 (L6)

03_model/server.py 是一个 FastAPI 应用，~140 行，提供 POST /generate 流式 SSE。运行时不依赖 transformers 库（from_pretrained 仅在启动时用一次）。

实测延迟（RTX 5090, fp32, batch=1）：

注意 prefill < decode 主要是 launch overhead 在小 batch 下占主导——实际 forward 计算 < 1ms。

4.3 Agent Loop (L8)

ReAct 驱动器 02_agent/agent_loop.py，~120 行。核心循环：

def run_agent(query):
    prompt = f"Q: {query}\n"
    for step in range(MAX_STEPS):
        # 流式 generate，遇到 "OBSERVATION:" 或 EOT 停
        chunk, stop = generate_until_stop(
            prompt,
            stop_strings=["OBSERVATION:", "<|endoftext|>"],
            max_tokens=100,
        )
        # 把 chunk 中的 THOUGHT/ACTION/ANSWER 抛给 L7 渲染
        for event in parse_events(chunk):
            yield event

        if stop == "<|endoftext|>":
            yield {"type": "done"}
            return

        # stop == "OBSERVATION:" → 模型在等工具结果
        action = last_action(chunk)            # 例: "calc(1234 + 5678)"
        obs = tools.call(action)               # 真调 Python，得 "6912"
        yield {"type": "observation", "v": obs}

        # 注入真实 observation 后继续循环
        prompt = prompt + chunk + f"OBSERVATION: {obs}\n"

5. 实验结果

5.1 端到端 Agent 测试 (10 query)

测试设置：L6 加载 00c_agent_sft/out/agent.pt（124M, agent SFT），L8 启用 AGENT_MODE=1，greedy decoding。

5.2 Tokenizer 验证

我们的手写 BPE 与 tiktoken 在以下样本上 bit-for-bit 等价（每个样本 token 序列完全一致）：

7/7 通过。涵盖 ASCII、Unicode (CJK, emoji)、连续空白、特殊字符。

5.3 GPU Kernel 性能

实测 RTX 5090 (sm_120, fp32, 2048×2048×2048 matmul)：

Attention (B=4, H=16, T=1024, D=64, fp32)：

讨论

• tiled / naive = 1.29× 远低于 A100 上典型的 ~6×。原因是 5090 的 HBM 带宽非常高，naive 实现没有真正 memory-bound——L8 cache 已经吃住了大部分复读。这反而让 tile 优化收益降低。
• cuBLAS / tiled = 7.5× 在所有现代 GPU 上稳定。这个 gap 来源于：(a) Tensor Core 用 TF32 而非 fp32，理论峰值高 2-4 倍；(b) cuBLAS 还有 register tiling、async copy、software pipelining 等更深优化。
• flash-attention 8.4× 加速跨硬件代际成立。这是结构性收益（省去中间 [B, h, T, T] attention matrix 的 HBM round-trip），不依赖具体硬件——T=8192 时差距会更大。

5.4 跨硬件可重复性

我们在三台不同硬件、不同 region、不同时间 cold-start 验证了完整 pipeline：

每次 cold-start 都是从 git clone 开始，没有任何手动 environment 配置（除了 pip mirror 和 ssh 密码）。所有 fallback（HF mirror、Xet disable、bundled tinyshakespeare）都自动生效。

每次实测的关键数字（loss、TFLOPS、agent 答对率）误差 < 5%，证明 pipeline 在不同硬件上确定性可重现。

6. 讨论

6.1 工程取舍

本项目刻意省略了多项生产级特性：

每个删除都是为了让"这一层在做什么"显式化。教学完备性 vs 性能完备性是对立的——我们选了前者。

6.2 局限性

(a) 模型规模过小：7M base + 124M SFT'd model 远低于实际可用 LLM 规模。L5 agent 在 OOD（非训练分布的对话）上失败（"How are you?"），这是 124M 量级的固有限制。复现 GPT-2 small (124M × 10B WebText token) 需要 1×A100 上 ~4 天，超出本项目教学范围。

(b) Agent 是单步单工具：258 条 SFT trace 都是 1 步 ACTION，没有 multi-tool / multi-step 推理样本。

(c) BPE vocab 是 OpenAI 公开数据：算法手写但 50,257 个 token 的具体合并规则来自 OpenAI 2019 年公开发布。完全独立的 BPE 训练需要参考 karpathy/minbpe。

(d) GPU kernel 不追 SOTA：手写 tiled matmul 比 cuBLAS 慢 7.5×。追近需要 Tensor Core、async copy 等深层优化，超出"分块原理演示"目标。

6.3 教学贡献

我们认为本项目的主要教学贡献在于：

1. 抽象坍缩：读完 04_transformer/model.py 330 行，读者不会再说"transformer 里有 attention"——会说"c_attn 是 [D, 3D] 的 Linear，把 x 投影成 q/k/v 三块，reshape 成 [B, n_head, T, head_dim]，然后调 scaled_dot_product_attention 后投影回 D"。
2. 数字感：loss = 10.815 ≈ ln(50257)、prefill 1.8ms、Triton 8.4× 加速——这些不是论文里的概念，是读者自己机器上能跑出的具体值。
3. 形态感：8 层端到端串通，让读者建立"一个 query 从 HTTP body 到 GPU SM 的完整心智模型"。下次面对 "我们推理服务慢了" 这种工程问题，能在脑子里算 KV cache 占用、HBM 流量、Tensor Core 利用率。

7. 可重复性

完整代码、数据、训练日志都开源在 github.com/fxp/LLM-from-query-to-result。

7.1 完整 cold-start 命令

# 1. Clone (CN 区用 gh-proxy)
git clone https://gh-proxy.com/https://github.com/fxp/LLM-from-query-to-result.git
cd LLM-from-query-to-result

# 2. 安装（建议 PyPI 用 aliyun mirror）
mkdir -p ~/.pip
echo -e "[global]\nindex-url = https://mirrors.aliyun.com/pypi/simple/\ntrusted-host = mirrors.aliyun.com" > ~/.pip/pip.conf
pip install -r requirements.txt

# 3. L2 验证（首次会下 BPE vocab + 可选 OpenAI gpt2 weights）
cd 04_transformer && python bpe.py && python inference.py "Hello, I am" && cd ..

# 4. L3 训练
cd 00_train && python prepare.py && python train.py && python sample.py "ROMEO:" && cd ..

# 5. L4 SFT (path A 自训 base 或 path B OpenAI 124M)
cd 00b_sft && python train.py                    # path A, ~28 sec on 5090
# OR
cd 00b_sft && python train_from_gpt2.py          # path B, ~35 sec
cd ..

# 6. L5 Agent SFT (依赖 path B 输出)
cd 00c_agent_sft && python build_data.py && python train.py && cd ..

# 7. L1 GPU benchmark
cd 05_gpu
nvcc -O3 -arch=sm_$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | tr -d '.') matmul_naive.cu -o matmul_naive
./matmul_naive
python benchmark.py
cd ..

# 8. 起服务，浏览器打开 localhost:8000
MODEL_PATH=$(pwd)/00c_agent_sft/out/agent.pt python 03_model/server.py &
AGENT_MODE=1 uvicorn 01_app.backend.main:app --port 8000 &

7.2 验证关键 invariants

# Tokenizer bit-for-bit ≡ tiktoken
python 04_transformer/bpe.py

# KV cache numerically ≡ full forward
python -c "
import sys, torch
sys.path.insert(0, '04_transformer')
from model import GPT, GPTConfig
import tokenizer
torch.manual_seed(0)
m = GPT(GPTConfig(n_layer=2, n_head=2, n_embd=64, block_size=32)).eval()
x = torch.tensor([tokenizer.encode('Hello world')])
with torch.no_grad():
    full = m(x)
    pre, kvs = m.step(x[:, :-1])
    dec, _ = m.step(x[:, -1:], kvs)
print('max diff:', (full[0, -1] - dec[0, -1]).abs().max().item())
"

预期输出：max diff: 2.4e-07（< 1e-6 即通过）。

7.3 报告中的所有数字均可复现

每个 TFLOPS、loss、user-visible latency 都有对应的源代码 + 实测命令。完整原始 stdout 见 reports/cold-start-2026-05-03-server3.raw.log（426 行）。

8. 结论

我们论证了一个 reasonable size 的 LLM 全栈系统——~10K 行代码、零外部 LLM 依赖、可在 5090 上 70 秒训练 + 1 秒推理——是可教学且可独立运行的。每一行代码、每一个数字都在 GitHub 仓库中。

更重要的是，这个项目展示了 agent 的本质机制：不是更聪明的模型，而是模型 + 外部工具的协调。124M 模型自己绝对算不对 1234+5678；带一个 7 行 Python 写的 calc 工具就行。这个机制和 ChatGPT 接 web_search、Cursor 接 grep+edit 是同一个本质——只是规模不同。

希望这个项目能帮助更多工程师从"调 API"过渡到"理解栈"，从"凭直觉调参"过渡到"凭数字调参"。

引用

• Yao et al. 2022. ReAct: Synergizing Reasoning and Acting in Language Models. arXiv:2210.03629
• Dao et al. 2022. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135
• Radford et al. 2019. Language Models are Unsupervised Multitask Learners (GPT-2 paper).
• Karpathy. nanoGPT (2022). github.com/karpathy/nanoGPT
• Karpathy. minbpe (2024). github.com/karpathy/minbpe

附录 A：项目数据汇总

附录 B：术语表

base model

仅经过 next-token prediction 预训练，未 fine-tune 的模型

instruction-tuned model

经过 SFT 的 model，能 follow 指令格式

KV cache

在 attention 中缓存历史 K 和 V 张量，使 autoregressive decode 复杂度从 O(N²) 降到 O(N)

prefill

处理 prompt 全部 token 的一次 forward，建立 KV cache

decode

基于 KV cache，每次 forward 1 个新 token 的 autoregressive 步骤

ReAct

Yao et al. 提出的 agent 推理格式，THOUGHT → ACTION → OBSERVATION 循环

flash-attention

Dao et al. 提出的 attention 实现，将 Q@K.T、softmax、@V 三步融合在 SM 内 shared memory 完成，避免大型 attention matrix 落 HBM

TF32

NVIDIA Ampere+ 的张量精度，10-bit 尾数，作为 fp32 替代品在 Tensor Core 上跑

本报告所有实验均可在仓库内复现。如发现任何与报告数字不一致的实测，欢迎在 GitHub Issues 反馈。