Nanocode：JAX + TPU 实现的极简编码助手

你必须登录才能更改通知设置

加载时出错。请重新加载此页面。

加载时出错。请重新加载此页面。

salmanmohammadi 2026年4月5日 维护者

我很高兴能够分享 nanocode。这是一个库，展示你如何从端到端快速构建自己的 Claude Code：从预训练、SFT、到使用 DPO 的 RLHF。粗略地说，我们将遵循使用 Constitutional AI 的最简单方法进行训练——这是 Anthropic 用于训练 Claude 模型的方法。我们将编写自己的 SOUL.md，定义智能体界面，即我们的模型与世界交互所使用的界面，生成合成数据，并使用偏好优化来使模型与我们的 SOUL 对齐。

nanocode 完全用 JAX 编写，设计用于在 TPU 上进行训练。我改编了 Karpathy 出色的 nanochat 项目的核心训练基础设施和理念，所以如果你熟悉 nanochat，nanocode 应该感觉非常相似。在本文结束时，你可以期望复制我的 1.3B 参数 d24 GPT 模型：

你可以通过 Google TRC 计划免费开始，该计划为你提供一个月的可抢占 TPU 免费访问权——我认为新的 Google Cloud 账户也会获得 $300 的积分。我很幸运能够在这个项目中访问 TRC 计划 3 个月，我发现大多数时候我的 Spot 实例很少被中断，我可以轻松地将同一个 Pod 保持一周或更长时间。

你可以在 TPU v6e-8 上用约 9 小时的总时间快速完成 nanocode-d24（1.3B 参数），成本为 $200，或在 ~1.5 小时内训练 nanocode-d20（477M 参数），成本为 $34。如果你使用 NVIDIA GPU，nanocode 也应该可以开箱即用，但你应该意识到 nanocode 已针对 TPU 进行了高度优化。

训练 nanocode：一个友好的智能体编码伙伴

Andrej 的 nanochat 原始发布文章很好地解释了我们在这里所做的事，你在 nanocode 中使用的命令几乎完全相同，所以我建议你先阅读他的工作。我将介绍我们为从模型中引出智能体编码行为而所做的不同之处。

分词和预训练

预训练和分词器训练过程与 nanochat 的基本相同，但我发现在预训练和分词器混合中以 1:5 的比例包含来自 The Stack-V2 的额外代码数据，能够产生更强的编码模型和更高效的代码分词，这帮助很大。

让我们首先下载分词器训练和模型预训练所需的数据集分片：

# we'll be training our d24, 1.3B parameter model. but you can adapt MODEL_TAG for your model size.
export NANOCODE_BASE_DIR="$HOME/.cache/nanocode"
export MODEL_TAG=d24
python -m data.pretrain -d fineweb-edu -n 300
# I've pre-packed and sharded The Stack similar to FinewWeb
python -m data.pretrain -d the-stack-v2-dedup -n 60

并启动我们的分词器训练脚本：

python -m scripts.tok_train --max-chars=2000000000
python -m scripts.tok_eval

作为参考，我们可以与 nanochat 的分词器进行比较，除了在训练混合中添加 The Stack 外其他都相同（好吧，我还添加了特殊令牌和模板逻辑来支持更复杂的工具调用，但稍后再详细说明）。

Comparison: nanocode vs nanochat
===============================================================================================
Text Type  Bytes    nanocode        nanochat        Relative     Better    
                    Tokens  Ratio   Tokens  Ratio   Diff %      
-----------------------------------------------------------------------------------------------
news       1819     407     4.47    375     4.85      +7.9%      nanochat  
korean     893      558     1.60    712     1.25     -27.6%      nanocode  
code       1259     326     3.86    492     2.56     -50.9%      nanocode  
math       1834     922     1.99    966     1.90      -4.8%      nanocode  
science    1112     259     4.29    228     4.88     +12.0%      nanochat  
fwe-train  4208518  902950  4.66    856883  4.91      +5.1%      nanochat  
fwe-val    4495276  975403  4.61    1010352 4.86      -3.6%      nanocode

我们可以看到，这以牺牲通用文本分词效率为代价，为代码提供了巨大的提升，但这没关系，因为我们希望我们的模型在一件事上做得非常好；智能体编码。我们的模型使用参数:数据比率为 8 进行训练（遵循 nanochat 的缩放定律分析）。让我们像这样启动一个训练运行：

python -u -m scripts.base_train \
    --batch-size=32 \
    --minibatch-size=1 \
    --config=configs.d24 \
    --eval-every=500 \
    --sample-every=500

你应该看到类似这样的内容：

Vocab size: 32768
World size: 8
1342.17728M model parameters
  67.108864M wte parameters
  1207.959552M h parameters
  67.108864M lm_head parameters
Training on 10737418240 tokens over 10241 steps
====================
Estimated FLOPs per token: 10066329600
Scaling the LR for the AdamW parameters ∝1/√(2048/768) = 0.612372
Step: 0/10241 | Loss: 10.398 | dt: 104.58s | | tkps: 10026 | mfu: 1.37 | ETA: -1.0 min | lr_multiplier: 1.000
        Peak bytes reserved/limit: 14.86/22.27
Step: 1/10241 | Loss: 9.771 | dt: 2.74s | | tkps: 382082 | mfu: 52.37 | ETA: -1.0 min | lr_multiplier: 1.000
Step: 2/10241 | Loss: 8.209 | dt: 2.74s | | tkps: 382220 | mfu: 52.39 | ETA: 234.1 min | lr_multiplier: 1.000
Step: 3/10241 | Loss: 7.327 | dt: 2.74s | | tkps: 382193 | mfu: 52.39 | ETA: 312.1 min | lr_multiplier: 1.000
...
fwe_bpb: 0.7626 | sv2_bpb: 0.4356 | avg_bpb: 0.5991 | dt: 90.53s
<|bos|>The capital of France is Paris. It is the largest city in France and the most populous city in
<|bos|>The chemical symbol of gold is Au. Gold is a soft, malleable, yellow metal that is
<|bos|>The closest planet to the Sun is Mercury, which is the smallest planet in the solar system. It is the closest
<|bos|>The opposite of hot is cold. The opposite of cold is heat. The opposite of heat is cold.
<|bos|>The second-last day of the week is the day of the Lord. (Leviticus 23:2)
...
CORE metric: 0.2352 | dt: 56.86s
Total training time: 467.15min

我们的模型已获得了一些关于世界的知识，这很不错。不过它仍然不知道星期六是什么：)。让我们看一些更彻底的定量结果，因为我们在训练期间仅使用评估数据的较小子集来估计指标：

python -u -m scripts.base_eval --checkpoint=base --minibatch-size=8

这将打印大量指标，但相关的是我们预训练集中的比特/字节：sv2（The Stack V2）和 fwe（FineWeb_EDU），以及 CORE 指标，可以直接与 nanochat 的结果和 GPT-2 进行比较。我编译了几个模型参数大小的结果，以了解我们的缩放定律：

| depth | params | CORE | cost | time   | MFU   | fwe bpb | sv2 bpb |
|-------|--------|------|------|--------|-------|---------|---------|
| d12   | 135M   | 0.090| $3   | 9 min  | 17.4% | 0.956   | 0.689   |
| d20   | 477M   | 0.170| $30  | 1.4 hrs| 45.2% | 0.838   | 0.533   |
| d24   | 1.3B   | 0.227| $200 | 9.3 hrs| 52.5% | 0.759   | 0.445   |

由于 CORE 衡量通用语言推理能力，而我们已将模型针对代码数据进行了调整，预期我们的 CORE 分数与相应的 GPT-2 模型相比会略有下降。单独在 FineWeb-EDU 上训练 d24 产生了 0.261 的 CORE 分数，与下面的 GPT-2 XL 和 nanochat-d24 一致。这里的权衡是我们期望我们的模型在编码任务中表现良好。

| model         | params | CORE  |
|---------------|--------|-------|
| GPT-2 Small   | 124M   | 0.114 |
| GPT-2 Medium  | 355M   | 0.185 |
| GPT-2 Large   | 774M   | 0.215 |
| GPT-2 XL      | 1.6B   | 0.257 |

在本文中，我主要会提到我们的 d24 GPT 模型，它类似于 nanochat 的 d24 GPT 模型，但上下文长度是两倍（4096 vs. 2048），以更好地支持多轮智能体对话。现在我们已经有了一个相当有能力的编码基础模型，让我们看看如何将其转变为一个成熟的智能体编码伙伴。

让我们从第一性原理出发，思考一下智能体模型究竟在做什么。预训练 LLM 会生成下一词元生成器，它们压缩了海量知识，但在遵循指令、回答与自身知识相关的问题，或修复 Python 文件中的 bug 等方面，其实并没有太大用处。要让模型真正完成有用的任务，我们还有大量工作要做。第一步是模板化——划分输入和输出的不同组成部分，让模型学会理解它所要执行的任务结构。我们以聊天模板为例。对话可以组织成轮次，双方每次各自进行一轮发言，因此模型需要知道当前轮到谁，以及对方说了什么。

User: What is 2+2?
Assistant: 4

我们可以将其模板化为：

<|bos|><|user_start|>What is 2+2?<|user_end|>
<|assistant_start|>4<|assistant_end|>

<|user_start|>、<|user_end|>、<|assistant_start|> 和 <|assistant_end|> 都是特殊词元，可以为原始文本提供结构。进行词元化时，我们通常会为它们分别保留一个完整的词元。很好。现在来思考一下，智能体模型可能会使用怎样的模板。智能体行为的基础是工具调用——在这种任务中，模型当前轮次的输出并非面向用户，而可能是通过某个现实世界接口执行的操作；该操作会产生输出，模型则可以实时响应这些输出。

从这个角度来看，工具调用的输出可以直接视为另一种对话轮次，因此我们额外保留两个特殊词元 <|tool_result_start|> 和 <|tool_result_end|>，让模型知道信息来自工具调用，而不是用户。现在，我们只需要让模型知道该如何调用工具——我们需要用模板表示模型希望调用的工具名称，以及它需要传入的任意可选关键字参数。以 grep 为例：

# search recursively for the term "TODO" in the src/ directory
grep -rn "TODO" src/ #

它大致会变成这样：

<|assistant_start|>
I'll search for TODO.
<|tool_call_start|>Bash<|tool_arg|>command<|tool_val|>grep -rn "TODO" src/<|tool_call_end|>
<|assistant_end|>

我们定义了用于标记整个工具调用起止位置的特殊词元（<|tool_call_start|> 和结束词元），以及用于划分该工具调用中不同命名参数的特殊词元（<|tool_arg|> 和 <|tool_val|>）。请注意，通过将工具调用模板嵌套在响应中，模型能够思考并解释自己的操作。

认真考虑最终的智能体接口实际会是什么样子非常重要——你肯定不希望设计出一套工具调用模板，花费大量资金用它训练模型，最后却发现它在实际环境中根本无法工作。定义工具时，我们需要在表达能力和可处理性之间做出权衡，也就是模型实际学会可靠使用工具的难易程度。对于最简单的智能体，我们希望它能够通过读取文件、搜索文件系统和写入磁盘来与 UNIX 环境交互。前面我们使用了 Bash 工具调用，但如果所有事情都只用 Bash，模型实际上就必须仅凭示例学会正确的 shell 语法，包括引号、选项和管道等。相反，我们可以预见，模型使用 grep 的频率可能足够高，因此应当为它提供一个专用的工具调用。在 nanocode 的智能体接口中，我定义了四种工具：

Read: <|tool_call_start|>Read<|tool_arg|>file_path<|tool_val|>...<|tool_arg|>offset
<|tool_val|>...<|tool_arg|>limit<|tool_val|>...<|tool_call_end|>
Edit:  <|tool_call_start|>Edit<|tool_arg|>file_path<|tool_val|>...<|tool_arg|>old_st
ring<|tool_val|>...<|tool_arg|>new_string<|tool_val|>...<|tool_call_end|>
Grep:  <|tool_call_start|>Grep<|tool_arg|>pattern<|tool_val|>...<|tool_arg|>path<|to
ol_val|>...<|tool_call_end|>
Bash:  <|tool_call_start|>Bash<|tool_arg|>command<|tool_val|>...<|tool_call_end|>

这使 nanocode 能够读写文件、搜索模式，并在需要时使用 UNIX 命令——不过，我并不指望在当前的算力和词元预算下，我们能得到一个真正学会有意义地使用 Bash 工具的模型。基于这些工具调用，我们的智能体 CLI 只需要充当一个轻量封装层：解析模型预测出的词元，拦截所有工具调用并执行，再将结果作为一种对话轮次提供给模型。

那么，我们要怎样教模型使用这些工具呢？最简单的方法，就是用数十万个工具使用示例来训练模型。这些示例可能如下所示：

User:
Write a Python function to process a list of strings representing stock prices and 
return a list of corresponding floating-point numbers. The input list may contain 
the following characters: letters (A-Z), numbers (0-9), periods (`.`), commas (`,`), and dashes (`-`).
Your code should pass the following test case:
input_list = ['1,234.56', '2,345.67']  
output_list = process_stock_prices(input_list)  
assert output_list == [1234.56, 2345.67]
---
Assistant:
cool, i'll write a function to process stock prices by removing commas and converting to floats. donezo!

[Edit Tool Call]  
file_path: process_stock_prices.py  
contents:  
def process_stock_prices(input_list):  
    return [float(s.replace(',', '')) for s in input_list]

---

Tool result:
def process_stock_prices(input_list):  
    return [float(s.replace(',', '')) for s in input_list]

这只是一个相当粗略的草图，不过你应该已经理解其中的思路了——用户提出请求，模型通过使用一个或多个可用工具来完成请求。它还会用一句有点滑稽的小话来解释自己正在做什么。前面我们提到，我们正在训练力所能及的最佳 Claude Code，而你可能听说过 Claude 的灵魂文档——这是一份关于模型性格、价值观和行为原则的书面规范。Anthropic 使用这份文档来指导 Claude 的训练：它定义所期望的行为，随后训练数据和偏好优化都会围绕这份规范塑造模型，使其与规范保持一致。

这正是宪法式 AI（Constitutional AI，CAI）背后的核心思想——早期 Claude 模型就是使用它训练的（这一技术的演进版本如今仍在 Claude 的训练中使用）。宪法式 AI 是一套由合成数据生成、监督微调和偏好优化组成的训练流程，其目的在于让模型与一组指定的特征和宪法原则，也就是 SOUL，保持一致。需要注意的是，CAI 作为一种对齐方法，重点在于构建有帮助且无害的智能体——尤其是防止模型给出有害回答——而我们主要使用它来实现模型风格上的对齐。

对于 nanocode 的 SOUL，我希望它拥有独特的说话风格：随意、友好，还带一点傻气，但不会谄媚或过度啰嗦。这就是我最终写出的版本。概括来说，nanocode 应当只使用小写字母，但代码中的专有名词可以例外；它应当热情友好，并且只遵循收到的精确指令。现在回头看，特别是对于这种规模的模型，我可能根本不需要那些哲学式的冗余内容。我们的 SOUL 与 Claude 的相比非常简单，但正如前面提到的，我们只希望模型非常擅长两件事：智能体式编程，以及遵循我们为它精心设计的个性。

宪法式 AI 通过两个阶段将这种 SOUL 注入模型：1）宪法式监督微调（Constitutional Supervised Fine-tuning，SFT）；2）基于 AI 反馈的强化学习（Reinforcement Learning from AI Feedback，RLAIF），也就是偏好学习阶段。

正如我前面提到的，我们既需要特定工具使用方式的示例，也需要符合模型 SOUL 的对话轮次。宪法式 SFT 阶段是一条合成数据生成流水线，你可以把它理解为拒绝采样与蒸馏的结合。对于我们的用例，这个循环如下：

一个生成器模型负责针对用户查询生成初始响应。

另一个独立的批评模型（即裁判）会接收用户查询、初始响应和 SOUL，并对模型的响应进行评分：如果评分低于我们的阈值，该响应就不合格。批评模型会输出具体的批评意见，详细说明初始响应究竟在哪些方面未能与 SOUL 保持一致。我们会将第一次失败的尝试保存为 Rejected 样本。如果评分通过，那就没问题！我们会将其保存为 Chosen 样本。

如果评分低于我们的阈值，该响应就不合格。批评模型会输出具体的批评意见，详细说明初始响应究竟在哪些方面未能与 SOUL 保持一致。我们会将第一次失败的尝试保存为 Rejected 样本。

如果评分通过，那就没问题！我们会将其保存为 Chosen 样本。

如果响应未通过第 2 步，生成模型会被给予原始提示、其失败的响应和批评意见，并被要求修改响应以解决反馈。此过程重复进行，直到批评模型接受响应。

初始提示 ──► 初始模型响应              
                         │                              
                         ▼                              
       ┌─────────► 批评 ──(通过?)──► 选中样本  
       │                 │                              
       │              (失败?)──────────► 首次拒绝样本
       │                 │                              
       └───── 修订 ◄─┘

在这个过程的最后，我们为给定的提示获得两个响应：一个与 SOUL 强一致的最终响应，以及一个初始的、不一致的响应。我们稍后将在偏好学习阶段使用这些对，但对于我们的宪法 SFT 阶段，我们只在 (初始提示、选中样本) 对上训练模型。值得注意的是，当生成模型无法在单一遍历中可靠地生成 SOUL 一致的输出时，批评循环至关重要——这对我在 TPU 上通过 vLLM 本地运行的大多数较小的开源模型来说都是如此。通过 OpenRouter 的前沿模型几乎在第一次尝试中就完美地完成了任务。我想说我这里详述的方法是我最初尝试的，但实际上这个项目的这部分花费了几个月的迭代和消融实验。

我为 nanocode 确定了两种方法。首先，我生成了一个数据集，包括简短的单轮对话，教导模型 Grep/Read 的基本智能体循环，然后 Edit 以编写解决手头任务的方案。重要的是，它教导模型如何理解我们工具及其结果的语法。为了为这个数据集提供种子，我重复使用了现有的 Python 开源指令数据集：

allenai/tulu-3-sft-personas-code

bigcode/self-oss-instruct-sc2-exec-filter-50k

theblackcat102/evol-codealpaca-v1

这被证明是引导我们的合成数据集生成过程的绝妙方式，因为它提供了约 120K 个高质量的样本，具有正确的 Python 解决方案和模型解释——我们只需应用上面的生成→批评循环将其按摩成我们的格式。你可以在 dev/process_datasets.py 和最终数据集 smohammadi/nanocode-tulu-selfoss-evol 中看到更多信息，我将使用一个例子来说明我们最终的数据集看起来像什么：

User:
开发一个 Python 函数，从列表中移除任何虚值。返回修改后的列表，不创建新的列表。使用列表推导式和 Python `*args` 参数来解包列表作为函数的参数，然后使用列表推导式过滤出任何虚值。

A:
以下是我们可以如何实现这个函数：
```python def remove_falsey_values(*args): return [val for val in args if val]```

User:
开发一个 Python 函数，从列表中移除任何虚值。返回修改后的列表，不创建新的列表。使用列表推导式和 Python `*args` 参数来解包列表作为函数的参数，然后使用列表推导式过滤出任何虚值。

A:
我将用 *args 的列表推导式来处理这个——在一次遍历中过滤虚值

[Edit]
  file_path: remove_falsey_values.py
  contents:
    1→def remove_falsey_values(*args):
    2→    return [val for val in args if val]

Tool Result:
---
    1→def remove_falsey_values(*args):
    2→    return [val for val in args if val]
---

[assistant]
完成——用 *args 原地修改列表，通过列表推导式过滤虚值，不创建新列表

这里我重复使用了初始提示，然后通过提取生成的代码并将其包装在我们的工具模板中，将模型的原始解决方案转换为 Edit 工具调用。我也在这里使用行号 (1->)，因为我相信如果模型可以在读取和编辑的文件中看到行号，这将帮助模型进行有针对性的 Grep 和 Edit 调用。

其次，与 tulu 和 self-oss 不同，后者中用户只是要求代码，evol-codealpaca 中的许多样本在指令和输出中都包含代码，用户描述现有代码并要求修改。这对于使用工具链接的多轮展开的映射非常适合。为此，我首先计算了原始代码和修改代码之间的差异，以获得 Edit 的有针对性的 old_string/new_string 参数，然后前置 Grep 和 Read 步骤，以获得智能体搜索相关函数、读取文件然后进行有针对性的编辑的展开。

User:
开发一个程序，接受用户提供的短语或句子，并为短语中的每个单词返回拼写，
包括任何标点符号。程序应识别并处理编程错误和异常。此外，更正下方代码
片段中发现的任何问题。
```python
def spell_word(a_word):
    for letter in a_word:
        print(letter)
spell_word("Hello, World!")

A: 代码片段本身没有任何错误。但是，它与问题陈述不符。以下是用 Python 编写的一个可能的解决方案：

def spell_word(sentence):
    words = sentence.split(' ')
    for word in words:
        print(word+": ", end="")
        for letter in word:
            print(letter, end=" ")
        print()
try:
    user_input = input("Enter a sentence: ")
    spell_word(user_input)
except Exception as e:
    print("An error occurred: ", e)```
    
现在在这个程序中：
- 用户被提示提供一个句子。
- 函数 spell_word 将句子分割为单个单词，并打印每个单词及其单独拼出的字母。
- 为异常处理留出了空间，以识别和管理可能发生的任何编程错误。如果确实发生错误，将打印一条包含错误详情的消息。

User:
<如上>

A:
[Grep]
  pattern: def spell_word
  path: .

Tool Result:
  spell_word.py:1:def spell_word(a_word):
  spell_word.py:2:    for letter in a_word:
  spell_word.py:3:        print(letter)

[Read]
  file_path: spell_word.py

Tool Result:
    1→def spell_word(a_word):
    2→    for letter in a_word:
    3→        print(letter)
    4→
    5→spell_word("Hello, World!")

我将设置单词分割器——在空白处分割，打印每个单词加冒号，然后用单个空格拼出字母

[Edit]
  file_path: spell_word.py
  old_string:
    def spell_word(a_word):
        for letter in a_word:
            print(letter)
    spell_word("Hello, World!")
  new_string:
    def spell_word(sentence):
        words = sentence.split(' ')
        for word in words:
            print(word+": ", end="")
            for letter in word:
                print(letter, end=" ")