- 人工智能
- 大模型
- 基础模型
- 推理模型
- NLP
【免费下载链接】Nemotron-3-Labs-Ultra-Math-SFT
Nemotron-3-Labs-Ultra-Math-SFT 是 NVIDIA 发布的 550B 参数数学推理大模型(技术报告中称 Nemotron-3-Ultra-SFT),它采用Mamba2-Transformer 混合结构 + LatentMoE 潜在专家路由 + MTP 多词元预测的前沿架构,支持最长 1M tokens 上下文,并以约 55B 的激活参数量实现了国际数学奥林匹克(IMO 2026)金牌级别的解题能力。本文将带你逐层拆解这套 Mamba2 与 LatentMoE 混合架构的设计思路,并附上核心配置速查表与 vLLM 部署要点。
📐 模型速览:550B 总参数,仅 55B 激活
| 项目 | 规格 |
|---|---|
| 总参数量 / 激活参数量 | 550B / 55B |
| 架构类型 | Mamba2-Transformer Hybrid LatentMoE + MTP |
隐藏层维度hidden_size | 8192 |
总层数num_hidden_layers | 108 |
| 最大上下文 | 262,144(可放宽至 1,048,576) |
词表大小vocab_size | 131,072 |
| 训练数据 | Nemotron-Math-Proofs-v3-SFT(414,890 条长程数学推理样本,覆盖 15,818 道难题) |
| 基座模型 | NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16(SFT 微调) |
关键信息可在 README.md 的 Model Architecture 一节查看,训练数据细节见 README.md。模型权重按 529 个分片存放,路由清单见 model.safetensors.index.json。
它为什么能解奥数题?答案就藏在架构里:用 Mamba2 层以"恒定内存"处理超长推理链,用稀疏 MoE 层以少量激活参数容纳巨大知识容量,再用 MTP 层加速推理。下面逐块拆解。
🔍 架构总览:108 层的 Mamba2-Transformer 混合堆叠
传统 Transformer 每一层都是"注意力 + FFN",而 Nemotron Hybrid 架构把每一层替换为四种积木之一:
| 积木 | 含义 | 代码实现 |
|---|---|---|
M | Mamba2 状态空间层 | NemotronHMamba2Mixer |
* | Attention 注意力层 | NemotronHAttention |
- | 普通 MLP 层 | NemotronHMLP |
E | MoE(专家混合)层 | NemotronHMOE |
每层由统一的 NemotronHBlock 按类型装配对应的 mixer,层类型映射逻辑见 configuration_nemotron_h.py 的 layers_block_type 属性。
如何读懂 hybrid_override_pattern
整网结构由 config.json 中的hybrid_override_pattern字符串定义,每个字符对应一层,共 108 个字符:
MEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*... 48个M + 48个E + 12个*统计结果一目了然:
- 48 层 Mamba2(
M) - 48 层 LatentMoE(
E) - 12 层 Attention(
*) - 0 层普通 MLP(
-)——所有前馈层全部 MoE 化
值得注意的设计规律:Mamba2 层与 MoE 层严格交替出现,注意力层大约每 6~8 层插入一次。也就是说,模型 90% 的计算量都发生在"线性复杂度的 Mamba2 + 稀疏激活的 MoE"上,注意力只保留在关键位置承担"精确检索"职责——这正是长上下文数学推理的理想分工。
Mamba2 层:恒定内存的选择性状态空间
Mamba2 层是这套架构处理长序列的主力,核心参数见 config.json:
- 256 个 SSM 头,每头维度 64(
mamba_head_dim),中间维度达 256 × 64 = 16,384 - 状态维度
ssm_state_size为 128,8 个状态分组 - 因果卷积核宽度 4,激活函数
silu,时间步 Δ 限制在 [0.001, 0.1]
与注意力机制需要把历史 KV 全部留存(内存随序列长度线性增长)不同,Mamba2 通过选择性状态空间模型把全部历史信息压缩进一个固定大小的递归状态中。无论输入 10 万还是 100 万 tokens,每层 Mamba2 的缓存开销都是常数——这是它能支撑 1M 上下文的根本原因。
代码实现见 NemotronHMamba2Mixer:输入投影后,经因果卷积、选择性离散化(计算 ∆、B、C)与门控 RMSNorm 输出。SSM 状态缓存默认使用 float32 精度(config.json),部署时可通过参数放宽为 float16 以省显存。
注意力层:64:2 极限 GQA 的精确检索
12 层 Attention 承担"全局精确检索"职责(数学推理中定位远距离引用、对齐公式符号等场景)。其参数设计相当激进:
- 64 个查询头(
num_attention_heads: 64),头维度 128 - 仅 2 个 KV 头(num_key_value_heads),分组查询注意力(GQA)比例为 32:1
KV 头数从常见的 8 压缩到 2,意味着注意力层的 KV 缓存只有常规设计的 1/4,12 层注意力层的缓存总量被压到极致,为超长上下文腾出显存空间。注意力权重与 KV 缓存更新逻辑见 NemotronHAttention。
LatentMoE 层:512 个专家如何通过 2048 维瓶颈工作
E层是模型容量的主要来源,每个 MoE 层包含512 个路由专家 + 1 个共享专家(config.json),每 token 只激活22 个专家(num_experts_per_tok: 22)。
什么是 Latent 瓶颈
与普通 MoE 让专家直接处理 8192 维向量不同,NemotronHMOE 在专家前后加了两道低维投影:
fc1_latent_proj:把 8192 维隐藏态压缩到2048 维潜空间(moe_latent_size)- 每个专家只在 2048 维上工作:2048 → 5120(
moe_intermediate_size)→ 2048,即 NemotronHMLP 以is_expert=True实例化 fc2_latent_proj:再把 2048 维升回 8192 维
专家在"2048 维潜空间"内交流,token 在专家间流动的数据量直接缩小为原来的 1/4,多卡专家并行(expert parallel)时的通信开销大幅降低——这是 LatentMoE 命名的由来。
路由与共享专家
- 路由器 NemotronHTopkRouter 采用 sigmoid 打分 + 组内 Top-k 筛选,并用
e_score_correction_bias做免辅助损失的负载平衡 - 选中专家的权重会归一化后乘以
routed_scaling_factor = 5.0,放大路由专家的贡献 - 1 个共享专家(中间维度 10,240,是路由专家的 2 倍宽)始终参与计算,并直接作用于压缩前的残差(
moe_shared_expert_overlap: false,即并行相加),保证每个 token 都有稳定的"保底"容量
MTP 多词元预测:内置的投机解码加速
模型末尾挂载了 1 个 MTP(Multi-Token Prediction)模块(num_nextn_predict_layers),它本身是一个 2 层的迷你混合块,结构模式*E= 1 层注意力 + 1 层 MoE(config.json)。首层会将输入 embedding 与主干隐藏态拼接(eh_proj,见 NemotronHMultiTokenPredictor),从而"预训练式"地学习预测下一个 token。
推理时可直接启用 vLLM 的nemotron_h_mtp投机解码(建议num_speculative_tokens: 5),由 MTP 模块一次性草拟多个 token,主干模型一步验证,显著提升解码吞吐——对需要长链推理的数学任务尤为有效。
核心配置速查表
| 参数 | 值 | 说明 |
|---|---|---|
hidden_size | 8192 | 主干隐藏维度 |
num_hidden_layers | 108 | 总层数(48 M + 48 E + 12 *) |
mamba_num_heads/mamba_head_dim | 256 / 64 | Mamba2 头配置 |
ssm_state_size | 128 | SSM 状态维度 |
conv_kernel | 4 | Mamba2 因果卷积核宽 |
num_attention_heads/num_key_value_heads | 64 / 2 | 极限 GQA 比例 |
head_dim | 128 | 注意力头维度 |
n_routed_experts | 512 | 每层路由专家数 |
n_shared_experts | 1 | 每层共享专家数 |
num_experts_per_tok | 22 | 每 token 激活专家数 |
moe_latent_size | 2048 | LatentMoE 潜空间维度 |
moe_intermediate_size | 5120 | 专家中间维度 |
moe_shared_expert_intermediate_size | 10240 | 共享专家中间维度 |
routed_scaling_factor | 5.0 | 路由专家输出缩放 |
max_position_embeddings | 262144 | 默认上下文(可至 1M) |
⚡ 部署要点:vLLM 单机 8×B200 起步
该模型属于前沿规模,官方推荐最低8× B200 单节点(约 1.5 TB 聚合 HBM)或跨节点 ≥8 张 H100/H200/GB200/GB300,运行时为 vLLM(推荐容器vllm/vllm-openai:v0.22.0)。关键启动参数:
--tensor-parallel-size 8:8 路张量并行--enable-expert-parallel:专家并行(512 专家分片到各卡)--dtype bfloat16--mamba-backend flashinfer+--mamba-ssm-cache-dtype float16:Mamba2 快速内核与低精度状态缓存--speculative-config '{"method": "nemotron_h_mtp", "num_speculative_tokens": 5}':启用 MTP 投机解码--reasoning-parser nemotron_v3:解析思考内容,实现对应 ultra_v3_reasoning_parser.py- 使用 1M 上下文时需设置
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1并将--max-model-len调到 1048576
对话行为由 chat_template.jinja 与 special_tokens_map.json 定义,生成默认 token 见 generation_config.json。
文件地图:架构定义在哪里
| 文件 | 作用 |
|---|---|
| config.json | 全部架构超参数(层模式、MoE、Mamba2、MTP) |
| configuration_nemotron_h.py | NemotronHConfig配置类与层类型解析 |
| modeling_nemotron_h.py | Mamba2Mixer、Attention、MOE、Router、MTP 完整实现 |
| ultra_v3_reasoning_parser.py | 推理内容(reasoning)解析器 |
| model.safetensors.index.json | 529 个权重的分片路由表 |
| README.md | 模型卡:架构说明、部署、数据集、许可 |
总结
Nemotron-3-Labs-Ultra-Math-SFT 的架构是一次清晰的"长上下文 × 大参数 × 低激活"权衡实践:48 层 Mamba2用恒定状态承载超长推理链,48 层 LatentMoE以 2048 维瓶颈让 512 个专家低开销协作(每 token 仅激活 22 个,激活参数约 55B),12 层极限 GQA 注意力负责精确检索,末尾MTP 模块提供内置投机解码。如果你想深入研究某一层的实现细节,从 modeling_nemotron_h.py 中的 NemotronHBlock 入手,按hybrid_override_pattern逐层对照,是最快的上手路径。
- 人工智能
- 大模型
- 基础模型
- 推理模型
- NLP
【免费下载链接】Nemotron-3-Labs-Ultra-Math-SFT
相关推荐
Rx.NET 构建体系现代化:基于 ADR-0001 的 .NET 7.0 时代工具链升级全解析(TFM 迁移、WinRT 依赖清理与 UWP 测试框架迁移)
Rx.NET 构建体系现代化:基于 ADR 0001 的 .NET 7.0 时代工具链升级全解析(TFM 迁移、WinRT 依赖清理与 UWP 测试框架迁移)
Nemotron-3-Labs-Ultra-Math-RL config.json 参数详解:混合 LatentMoE 模型关键超参数完整清单
Nemotron 3 Labs Ultra Math RL config.json 参数详解:混合 LatentMoE 模型关键超参数完整清单 Nemotron
人工智能大模型基础模型强化学习推理模型拆解NVIDIA Nemotron推理教师模型架构:Mamba-2+MoE+Attention混合与LatentMoE是这样工作的
拆解NVIDIA Nemotron推理教师模型架构:Mamba 2+MoE+Attention混合与LatentMoE是这样工作的 NVIDIA Nemotro
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考