单卡部署的 Lean 4 自动形式化模型架构实践

项目概览

该项目面向「自动形式化（autoformalization）」这一形式化数学的关键瓶颈：把自然语言数学陈述翻译成 Lean 4 等交互式定理证明器可编译、机器可校验的代码。当前主流自动形式化模型依赖稠密、多卡部署的大模型，作者选择在单卡 RTX 5090ti（32 GB VRAM）上，使用稀疏 MoE 模型 Qwen3-Coder-30B-A3B，通过一整套工程改造完成训练与推理。

核心目标有三：把小型 Transformer 的形式化数学能力推到极限，验证一组微调方法的有效性，并搭建可快速迭代的多智能体系统。

基础模型与单卡适配

模型选型

• 基座策略模型选用 Qwen3-Coder-30B-A3B，30B 总参数量、每 token 激活约 3B 参数、48 层 Transformer。
• 该模型在代码与软件合成任务上做过预训练，对类型论语义、结构化算法逻辑具有更好的基线理解。
• 稀疏 MoE 架构在保留大参数容量的同时，将前向激活限制在 3B 左右，有利于推理吞吐。

显存与量化

• bf16 精度下 30B 模型约需 60 GB 显存，超出单卡预算。
• 整体方案：
  • 整模型 NF4 双量化：模型直接载入主机内存，使用 NormalFloat4 + 双量化方案量化到位。
  • GPU 分配：参数全部驻留显存，加载后约 16.7 GB，训练峰值约 26–30 GB（micro-batch size = 2）。
  • 关闭 CPU 卸载，避免 PCIe 延迟瓶颈。

MoE 专家解耦工程

Qwen3Moe 的实现把专家权重存为 3D fused tensor，现有量化库无法处理这种结构。该项目为此自定义了一层解耦机制：

• 将每层的 fused expert block 拆解为独立的 _Expert 子模块，包含标准的 nn.Linear（gate_proj、up_proj、down_proj）。
• 把预训练权重从 fused 3D 张量拷贝到解耦后的 nn.Linear。
• 通过 bitsandbytes.nn.Linear4bit 就地位化量化。
• 前向传播在显式线性层上执行专家路由，保持基座模型输出语义不变。

为兼容 PEFT：在 PEFT 注入阶段把 model.config.model_type 动态改写为 qwen3_moe_unfused，使其能匹配自定义解耦子模块作为目标。

适配器与变体

• 注意力模块（q_proj、k_proj、v_proj、o_proj）：LoRA 秩 r = 64，α = 128。
• 专家模块（gate_proj、up_proj、down_proj）：秩 r = 8，α = 16。
• 可训练参数：约 4.68 亿，约占总参数的 2.92%。
• 变体对比：标准 QLoRA（线性低秩更新） vs. QDoRA（把更新分解为方向向量与幅度标量）。

训练与检索增强

GRPO + Lean REPL 反馈

训练阶段使用 Group Relative Policy Optimization（GRPO），通过持续运行的 Lean REPL 编译器反馈提供奖励信号。值得注意的是，正确的编译并不等于成功的翻译，因此评测引擎本身的准确性也是该领域的关键难题。

双层检索增强（Premise RAG）

系统 B 负责把自然语言陈述接地到形式化前提：

• 稠密语义向量化（句向量） + BM25 词法再排序。
• 检索结果作为提示词的上下文，喂给 NF4 基座 + QLoRA 适配器构成的策略模型。
• 模型采样生成 Lean 4 代码，再进入编译器反馈循环。

小结

整套方案围绕「单卡、稀疏 MoE、量化 + 解耦 + 微调 + 强化学习 + 检索增强」这条主线，把自动形式化从需要多卡集群的工作量压缩到了消费级硬件。其工程细节——尤其是 MoE 专家解耦与 PEFT 目标模块改写——对其他在单卡上微调 MoE 的项目也有借鉴意义。