NanoJev 模型解析:架构、训练与部署
NanoJev 模型解析:架构、训练与部署
数据集:C-Tianyu/NanoJev-Data · Datasets at Hugging Face
模型:C-Tianyu/NanoJev · Hugging Face
目录
- 项目功能与涉及模型
- “Nano 复刻版”的含义
- 模型结构:主干 + 决策头
- 参数量精确验证
- 训练方法与数据集
- 模型组合、保存与部署
- 概率分布的使用:模型与控制器解耦
- 生产环境热更新分析
- 关键事实与文件索引
1. 项目功能与涉及模型
1.1 项目定位
NanoJev 是 Jev(typesafe.ai 的 System One 模型)的 0.6B 参数 nano 复刻版:输入状态与问题,直接输出完整概率分布,不做自回归生成(decode 步数为 0)。
核心功能模块:
- 并行决策引擎:独立状态、问题、候选路径放进同一次 backbone 前向批量处理。三类问题:
Choice:返回 2–255 个候选的完整分布(集合注意力 + softmax 共享评分头)Boolean:sigmoid 输出命题成立概率Score:2–10 个有序等级的概率分布及加权期望
- 四款游戏统一支持(同一个 step-400 checkpoint):
- ViZDoom Basic:瞄准射击,测试 128/128 成功
- ViZDoom Predict Position:移动目标火箭射击,27/128
- 50×50 Maze:225 次行动到达出口
- Snake:完整存活 256 步内吃到 30 个食物,8/8
- 推理服务:
scripts/serve_decisions.py,模型只加载一次,POST /api/evaluate提交状态与问题批次 - 三模型并排网页回放:Jev / NanoJev / 未微调 Qwen 按同一游戏时钟同步展示
- 训练与数据管线:混合任务 SFT、五分区数据划分、独立重放评测
- 发布:Hugging Face 模型(
C-Tianyu/NanoJev)+ 数据集(C-Tianyu/NanoJev-Data),版本unified-games-v1
1.2 测试结果(274 个测试案例,相同观察接口与固定种子 epsilon-greedy 控制器)
| 模型 | Maze | Snake | Basic | Predict Position |
|---|---|---|---|---|
| NanoJev | 4/10 | 8/8 | 128/128 | 27/128 |
| Jev | 7/10 | 8/8 | 56/128 | 11/128 |
| 未微调 Qwen3-0.6B | 2/10 | 0/8 | 56/128 | 11/128 |
1.3 项目中出现过的模型角色
| 模型 | 角色 | 是否参与 NanoJev 运行时组合 |
|---|---|---|
| Qwen3-0.6B 主干 + 决策头 | NanoJev 本体 | 是(合为一个 DecisionModel) |
| 未微调 Qwen3-0.6B | 评测基线(原始词表 option logits) | 否,仅对比 |
| Jev | 评测基线(API 调用) | 否,仅对比 |
| Sonic Doom 视觉专家(CNN+GRU) | 生成 Predict Position 专家监督数据(896 局轨迹) | 否,只产数据 |
| APPO 检查点(约 340 万参数 CNN+循环网络) | 另一条 RL 实验线,生成 Basic 动作演示 | 否,独立实验 |
对外部署的模型只有”Qwen3-0.6B + 决策头”这一个(两组件合成的一个 checkpoint)。
2. “Nano 复刻版”的含义
README 标题为 “A nano replica of Jev”,拆开理解:
- Jev:typesafe.ai 发布的商业 “System One” 模型。核心思想是”系统一”式决策——不做链式推理、不生成 token,一次前向直接输出候选的完整概率分布(对应 Choice / Noul / Score 等接口)。
- 复刻(replica):复现的是这套决策范式和接口,并做到可直接同条件对比——Jev 通过官方 API 参与评测,两边使用相同观察接口、候选动作与固定随机种子。项目文档
docs/TYPESAFE_CONTRACT.md记录了本地接口与 TypeSafe API 的逐项映射(如本地boolean对应其noul)。 - nano:指规模极小。Jev 权重闭源,无法逐参数复制;NanoJev 改用开源小模型 Qwen3-0.6B 底座 + 自训练决策头,把同一能力压缩到 5.96 亿参数、单卡可跑的量级,并统一四款游戏。
一句话:NanoJev = 用 0.6B 开源小模型 + 自训练决策头,复现”一次前向直接输出决策概率分布”范式的 nano 级实现。
3. 模型结构:主干 + 决策头
3.1 整体架构对比
1 | |
主干(词嵌入 + 28 层)两图完全相同:NanoJev 没有改动主干任何一层,差异全部在输出端。
3.2 切掉与保留的内容
| 部件 | 处理 |
|---|---|
| 28 层 transformer(attention/MLP/RMSNorm/RoPE) | 保留,完整微调 |
| 词嵌入 + tokenizer | 保留(输入编码仍走 Qwen 词表和文本模板) |
lm_head(词表投影,输出”下一个 token 的概率”) |
切掉 |
| 自回归解码循环、KV cache | 不用(use_cache=False,无生成步骤) |
| 输出路径 | 换成决策头:候选路径末 token 的 hidden → 标量分数 → 分布 |
代码证据:
scripts/train_toy_decisions.py第 307 行:model = DecisionModel(lm.model, args.set_head).cuda()——lm是AutoModelForCausalLM,.model取出纯主干(Qwen3Model,不含 lm_head),下一行del lm丢弃整个生成式包装;scripts/predict_toy_decisions.py第 221–224 行:用AutoModel.from_config(body_config)只构造主干结构,再叠加决策头,load_state_dict(..., strict=True)从best.safetensors严格加载——若 checkpoint 带 lm_head 权重会直接报错。
一个概念上的校正:这不叫”用决策头代替 lm_head 的功能”——两者输出语义不同。lm_head 服务”生成答案 token”,决策头服务”给候选打分排序”。准确描述是:保留主干的全部”理解状态”能力,把生成式的输出路径整体换成判别式的打分路径。
旁证:未微调 Qwen 基线恰恰靠被切掉的这部分工作——frozen_native_baseline(train_toy_decisions.py 第 231 行)里 F.linear(h, lm.get_output_embeddings().weight[label_ids]) 直接读词表输出头做选项 logits。
3.3 决策头内部结构与数据流
1 | |
DecisionModel 定义在 scripts/train_toy_decisions.py 第 86–141 行:
backbone:Qwen3 transformer 主干(保留全部 28 层 + 词嵌入)norm:LayerNorm(hidden)(主干自带final RMSNorm,决策头入口再加独立 LayerNorm,属实际实现)scalar:Linear(hidden → 1),std=0.02随机初始化(避免首步死区)set_head='attention'(仅 Choice 启用):set_project:Linear(hidden+1 → 128),拼接 hidden 与log Kset_attention:MultiheadAttention(128, 4 heads, dropout=0.0)set_output:Linear(128 → 1),零初始化残差修正
细节说明:
- 集合注意力是 Choice 专属(第 127–133 行
z = z.index_add(0, choice, δ));Boolean/Score 只走scalar基础分; log K特征让模型感知”总共有多少候选”,是动态候选 2–255 能力的来源;- 输入侧
State:/Question type:/Question:/Candidate:/Decision:模板见load_examples(第 45–57 行),每条候选路径以 EOS 结尾。
4. 参数量精确验证
按 Qwen3-0.6B 官方 config(hidden=1024、28 层、16Q/8KV、head_dim=128、intermediate=3072、vocab=151936、tied、attention_bias=false)逐项计算:
| 组件 | 计算 | 参数量 |
|---|---|---|
| embed_tokens | 151,936 × 1,024 | 155,582,464 |
| 每层 Qwen3 Decoder | q 2,097,152 + k 1,048,576 + v 1,048,576 + o 2,097,152 + q_norm/k_norm 各 128 + LLN 1,024 + MLP 3×3,145,728 + PLN 1,024 | 15,730,944 |
| 28 层小计 | 15,730,944 × 28 | 440,466,432 |
| final RMSNorm | 1,024 | 1,024 |
| 主干合计 | 596,049,920 | |
| 决策头:norm + scalar | LayerNorm 2,048 + Linear 1,025 | 3,073 |
| 决策头:集合注意力三件套 | set_project 131,328 + attention 66,048 + set_output 129 | 197,505 |
| 总计 | 596,049,920 + 200,578 | 596,250,498 |
最后一行与训练记录 results/training_run_receipts.json 的 parameter_count 完全一致到个位,反证两件事:
- checkpoint 中确实没有独立的 lm_head——
tie_word_embeddings: true使词表投影与 embed_tokens 共享,切掉零损失;若独立保存会多 155,582,464(约 7.52 亿); - 主干 596,049,920 + 决策头 200,578(占 0.034%)。
官方 config 字段与项目实践对照
| config 字段 | 原版含义 | 项目中的实际使用 |
|---|---|---|
tie_word_embeddings: true |
词表投影共享 | 支撑”切 lm_head 零损失”;基线 Qwen 的选项 logits 直接从 embedding 权重取行 |
max_position_embeddings: 40960 |
支持 4 万 token 上下文 | 项目用不满:训练 max_length 8192(v1 记录)、toy 512,推理默认由 checkpoint config 决定 |
use_cache: true |
官方默认开 KV cache | 项目强制 model.backbone.config.use_cache = False |
torch_dtype: bfloat16 |
官方分发精度 | 项目 parameter_storage: float32 加载与训练,仅前向 autocast 到 bf16 |
sliding_window: null / use_sliding_window: false |
全注意力 | 与决策头布局一致(本版本未做前缀共享优化) |
rope_theta: 1000000、attention_bias: false |
RoPE base 1M;注意力无 bias | 参数量计算无 bias 项;每层带 q_norm/k_norm(results/unified_td_v1/validation.json 探针参数 layers.0.self_attn.q_norm.weight 可证) |
5. 训练方法与数据集
5.1 决策头的训练流程
以主训练脚本 scripts/train_unified_games.py 为例(unified-games-v1 由它训练),核心是两阶段全参数训练:
1 | |
具体机制(第 765–805 行):
- 优化器
AdamW分两个参数组:body(backbone)与head(非 backbone 参数),weight decay 0.01; - 每步:按四任务 1/3、1/3、1/6、1/6 权重采样一批问题 → 完整问题交叉熵 →
loss.backward()→ 梯度裁剪 1.0 →optimizer.step(); - 热身阶段 backbone 完全不更新,避免随机初始化的头破坏预训练主干;
- 每 50 步在 dev 上评估,按加权 dev CE 选最优 checkpoint 存为
best.safetensors。
5.2 微调方法:全参数微调,不是 LoRA
- 全库搜索
lora|peft|qlora无真实匹配; - backbone 参数直接
requires_grad_(True)后由 AdamW 端到端更新; - 准确说法:全参微调主干 + 从头训练决策头——主干从官方预训练权重出发微调(lr 1e-5),决策头随机初始化从零学(lr 1e-4),两者联合优化;
- 后续
--stage critic(RLCD 后训练)、TD 目标等也是继续训练同一模型,同样不是 LoRA。
5.3 数据集格式
单条样本 = 一个 state + 多个带类型的 question + 监督目标:
1 | |
- 张量化后每条候选路径 =
State:…\nQuestion:…\nCandidate:\n{候选文本}\nDecision:+ EOS;同一问题的候选共享前缀、一次前向打分; - 目标来源三类(
POLICY_TARGET_KINDS):api_policy_distribution(Jev API 概率)、expert_action(专家动作 one-hot)、expert_distribution(专家分布);另有outcome行记录真实游戏结局; - 规模:每版本 18,760 条问题(ViZDoom Predict Position 11,173 + Basic 5,160 + Maze 1,469 + Snake 958),train 分区 10,898;896 局 Predict Position 专家轨迹(17,498 决策步);五分区 train / dev / calibration / test / OOD;
- 来源:游戏局部决策问题由项目生成管线产出;ViZDoom 动作监督来自 Sonic Doom 视觉专家;部分软目标来自 Jev API 真实输出概率。
与常规 NLP 训练集的对比:
| 维度 | 常规 LLM SFT 数据 | NanoJev 数据 |
|---|---|---|
| 样本形态 | prompt → 答案文本序列 | 一个 state + N 个独立问题(问题间不消费彼此答案) |
| 答案形态 | token 序列 | 候选集合上的概率分布(choice 2–255 / boolean / score 2–10) |
| 监督来源 | 人写的标准答案 | ① Jev API 真实概率 ② 专家动作 ③ 真实游戏结局 |
| 损失 | next-token 交叉熵 | 完整问题的分布交叉熵(对全部候选一次算) |
| 数据审计 | 一般无 | 严格 provenance:teacher 概率精度校验、无效目标 quarantined、episode/策略 ID 交叉校验 |
5.4 toy_inference_example.json 的定位
research/toy_inference_example.json 是推理请求样例(文档中称”无teacher推理输入”),验收目标是”不读取 teacher/gold,只加载本地 checkpoint,一次 forward 返回 2 states / 6 questions / 15 paths 的概率”——故意不含答案字段。推理本身不需要答案:模型输出概率分布,排序、贪心、采样都是调用方的事。
该样例的 state 是客服工单域(退款、软件故障),三种题型齐全,用于展示通用决策输入格式。
6. 模型组合、保存与部署
6.1 组合方式
NanoJev 是一个整体:Qwen3-0.6B 主干 + 决策头合为一个 nn.Module(DecisionModel),一起训练、存在同一个 best.safetensors、推理时一起加载。不是两个模型的级联或路由,而是”主干 + 输出头”关系(类似 BERT + 分类头):
1 | |
Checkpoint 目录 = best.safetensors(全部权重)+ config.json(记录 set_head 等)+ backbone_config/ + tokenizer/。
6.2 推理服务
scripts/serve_decisions.py:启动时构造 DecisionPredictor(”构造时加载一次权重”,predict_toy_decisions.py 第 177 行),Handler 通过闭包捕获,/api/health 返回 {'ready':True,'model_loaded_once':True,'provider_calls':0}——model_loaded_once 为写死的状态,当前实现不支持热更新,属本地演示服务设计。
7. 概率分布的使用:模型与控制器解耦
一句话核心:NanoJev 是”感知/判断”组件,控制器才是”决策者”。
7.1 模型侧:只产出概率
一次前向返回每个候选的概率分布(answer_from_probabilities 会先校验”概率有限、总和为 1”),模型自身不做 argmax、不采样、不带探索逻辑。
7.2 调用方侧:拿到分布后做四件事
控制器逻辑在 scripts/unified_game_pipeline.py 第 183–202 行:
1 | |
各题型取答案规则:
| 题型 | 调用方怎么用分布 |
|---|---|
| Choice | max(ids, key=probabilities.__getitem__) 贪心,或整体按分布采样 |
| Boolean | 取 p_true(或与阈值比较),也可用于 TD/Brier 目标 |
| Score | 概率加权期望 score = Σ i·p_i,或取 argmax 等级 |
7.3 为什么解耦
- 同一 checkpoint,换控制器即换策略:三模型对比同时跑
greedy和T=1 sample(verify_nanojev_comparison.py的POLICIES); - 可复现、可审计:控制器声明(controller/epsilon/seed/tie_break)是冻结策略的一部分,行为分布可离线重算复核——
unified_td.py要求记录动作与冻结控制器 RNG 抽样完全一致,否则校验失败; - 概率本身是分析对象:迷宫评测统计”最优动作”上的概率质量
p_optimal,只有输出完整分布才可能。
8. 生产环境热更新分析
结论先行:热更新能力不取决于输出端接的是决策头还是语言头,两者在权重层面是同一个问题;差异只在工程配套。
8.1 为什么本质相同
两种架构的 checkpoint 都是”主干参数 + 输出层参数”(决策头的 norm/scalar/set_* 或词表投影)。热更新要解决的都是三个问题:
- 权重替换的原子性:新权重写到新目录、完整校验后原子切换(项目本身有 SHA256 校验文化);
- 在途请求的一致性:
load_state_dict非原子,原地重载时并发请求可能读到混合权重,需排空或双缓冲; - 结构兼容性:同结构可原地/双缓冲;结构变了(
set_head从none换attention,或换词表大小)参数键/形状不匹配,只能重启或双实例换流。项目代码有明确拦截:train_pipeline_decisions.py的"Warm-start architecture cannot change set_head"。
8.2 两种架构的工程差异
| 维度 | 主干 + 语言头(LLM 服务) | 主干 + 决策头(NanoJev 形态) |
|---|---|---|
| 生态工具 | 成熟:vLLM 支持运行时 update_weights,多副本滚动重启是常态 |
自研服务,热更需自己写 |
| 显存约束 | 大模型几百 GB,双驻留成本高 | 5.96 亿参数,bf16 约 1.2GB,双驻留轻松 |
| 缓存失效 | 有 KV cache / prefix cache,换权重后要处理失效与流式请求收尾 | 无 KV cache、无自回归,无缓存失效问题 |
| 契约风险 | 换模型要重新评估 prompt、输出格式、对齐 | 接口不变(状态+问题→概率),调用方无感 |
8.3 当前现状与改造方案
当前实现是”加载一次”(见第 6.2 节)。若要生产热更,推荐双缓冲 + 原子指针切换:
1 | |
需要改的点:Handler 不能闭包捕获固定的 engine,要持有可变引用(如带锁的容器)。
共同红线:结构不兼容的更新不能原地热更,只能重启或双实例换流。
9. 关键事实与文件索引
核心数字
| 项目 | 数值 |
|---|---|
| 主干参数量 | 596,049,920 |
| 决策头参数量 | 200,578(占 0.034%) |
| 总参数量 | 596,250,498(= 训练记录 parameter_count) |
| 数据规模(每版本) | 18,760 条问题;train 分区 10,898 |
| PP 专家轨迹 | 896 局 / 17,498 决策步 |
| 问题候选范围 | Choice 2–255;Score 2–10 等级;Boolean 1 命题 |
文件索引
| 主题 | 文件 |
|---|---|
DecisionModel 与决策头定义 |
scripts/train_toy_decisions.py(第 86–141 行) |
| 主干选取与生成式包装丢弃 | scripts/train_toy_decisions.py(第 307 行) |
| 未微调 Qwen 基线(词表头取行) | scripts/train_toy_decisions.py(第 231 行) |
| 主训练脚本(两阶段、全参) | scripts/train_unified_games.py(第 775–805 行) |
| 推理对象与 checkpoint 结构 | scripts/predict_toy_decisions.py(第 149–156、177、221–224 行) |
| 推理服务(加载一次) | scripts/serve_decisions.py(第 24、60–69 行) |
| 控制器(贪心/采样/epsilon) | scripts/unified_game_pipeline.py(第 183–202 行) |
| 冻结控制器 RNG 校验 | scripts/unified_td.py(第 107–111 行) |
| 训练记录(参数量) | results/training_run_receipts.json(第 51 行) |
| QK-Norm 探针参数 | results/unified_td_v1/validation.json(第 15 行) |
| 接口映射(boolean↔noul) | docs/TYPESAFE_CONTRACT.md |
| 推理请求样例(无答案) | research/toy_inference_example.json |
| 超参变更拦截 | scripts/train_pipeline_decisions.py(第 452–453 行) |
总结
NanoJev = Qwen3-0.6B 主干(保留 28 层与词嵌入、切掉 lm_head)+ 自训练决策头(约 20 万参数、两阶段全参训练),一次前向直接输出 2–255 个候选的完整概率分布;模型只负责”算概率”,排序、贪心、采样与探索全部交给外部控制器。