Granite-4.1-3B-Base 长上下文扩展原理：396B Token 如何把窗口从 128K 扩到 512K

Granite-4.1-3B-Base 是 IBM Granite 团队推出的 3B 参数开源大语言模型，其最大的亮点之一是长上下文能力：模型通过额外消耗约 396B Token 的训练数据，把上下文窗口从 128K 扩展到 512K，一次可"读完"约 40 万字的长文档。本文将用通俗的方式讲清楚这套长上下文扩展原理。

一、先认识 Granite-4.1-3B-Base

Granite-4.1-3B-Base 是一款 decoder-only（仅解码器）的稠密 Transformer 模型，与 Granite 4.0 3B Micro 共享同一套底层架构和权重，采用 Apache 2.0 开源协议，由 IBM Granite Team 开发，面向文本生成、摘要、问答、代码补全（含 FIM 中间补全）等通用任务。

核心架构参数如下（完整信息见 README.md）：

💡 512K Token 是什么概念？按中英文混合约 1.5 字符/Token 估算，一次可容纳约70 万到 100 万字符，相当于同时喂入十几份技术报告或一整本技术书籍。

二、扩展的第一步：为什么 RoPE 让 512K 成为可能

想让模型"看"得更远，首先要让它能"数"到更远的位置。Granite 4.1 使用的是RoPE（旋转位置编码，Rotary Position Embedding），它是当前主流长上下文模型的事实标准。

1. RoPE 的直观理解

可以把每个 Token 的位置想象成钟表上的指针角度：模型不记忆"我在第几页"，而是通过两个 Token 之间的相对角度差来判断"你离我有多远"。这种相对性让模型天然更容易外推到训练时没见过的更长序列。

2. 关键细节：超大的 RoPE base 值

打开 config.json，会发现一个值得注意的参数：

• rope_theta: 10000000（即 1e7）

普通模型的 RoPE base 多为 10000，而 Granite 4.1 把 base 提高到了1000 倍。base 越大，不同频率分量的旋转周期越长、位置信号区分度越平滑，模型在超过训练长度时越不容易"迷失"。再配合 config.json 中的max_position_embeddings: 131072（128K）这一"基础窗口"设定，就为后续平滑扩展到 512K 打下了底子。

三、扩展的核心：396B Token 的"第 5 阶段训练"

真正让窗口从 128K 变成 512K 的，是模型五阶段训练策略中的最后一棒。根据 README.md 的训练数据表，整个训练流程如下：

可以看到，前四个阶段累计消耗约14.5T Token，把基础能力打磨到位；而阶段 V 专门用396B Token（约占前四阶段总量的 2.7%）做长上下文扩展。这一阶段的几个设计要点：

1. 渐进式变长：从 128K 起步，逐步拉长训练序列到 512K，而不是"一步到位"。就像跑步者不能直接冲刺全程，先建立长距离耐力。
2. 指数学习率调度：用缓慢衰减的学习率避免"灾难性遗忘"——模型在学会长距离依赖的同时，没有丢掉前 14.5T Token 学到的通用能力（MMLU、HumanEval 等基准成绩依然在线）。
3. 数据效率可观：396B Token 相比从零预训练的 15T，仅约 2.6% 的开销，就换来 4 倍的上下文窗口。

📊 这一训练在 NVIDIA GB200 NVL72 集群（CoreWeave）上完成，万卡级互联保障了超长序列的分布式训练稳定性，详见 README.md。

四、512K 为什么"用得起"：GQA 让显存压力可控

长上下文最大的敌人是显存。每生成一个新 Token，模型都要为之前所有 Token 保存 K/V 缓存，窗口越长缓存越大。Granite 4.1 用GQA（分组查询注意力）化解了这一难题：40 个查询头只共享8 个 KV 头（见 config.json），K/V 缓存体积直接降到 MHA 的 1/5。

算一笔账（bf16 精度下）：

也就是说，512K 满窗口推理时，KV 缓存约 40 GB——加上模型权重本身（约 6.8 GB，见 model.safetensors.index.json），一台带 80GB 显存的 GPU（如 A100/H100）即可承载，门槛并不离谱。

五、如何在仓库里验证这些配置

想在本地核对长上下文相关的关键设置，只需要看这几个文件：

• config.json：max_position_embeddings（128K 基础窗口）、rope_theta（1e7 大 base）、num_key_value_heads（GQA 的 8 个 KV 头）
• tokenizer_config.json：model_max_length被设为极大值，意味着分词器层面不限制输入长度，512K 上限由模型侧控制
• generation_config.json：生成时使用的特殊 Token 配置（<|end_of_text|>等）
• model-00001-of-00002.safetensors、model-00002-of-00002.safetensors：模型权重本体（约 6.8 GB，bf16）

六、实用建议与注意事项

1. 推理长度建议：虽然支持 512K，但注意力计算量随长度平方增长（512K 的理论 FLOPs 是 128K 的 16 倍），实际使用建议按需截取，"够用即可"。
2. Base 模型未做安全对齐：这是基础模型，可能产生不当输出，企业部署时建议配合 Granite Guardian 类安全模型过滤输入输出（参见 README.md 的伦理说明）。
3. 多语言友好：模型支持中文、英文、日语等 12 种语言，长上下文场景下中英文混合文档同样适用。
4. FIM 代码补全：仓库内置<FIM_BEGIN>等特殊 Token（见 tokenizer_config.json），代码场景可启用 Fill-in-the-Middle 补全模式。

总结

Granite-4.1-3B-Base 的长上下文扩展，本质上是一次**"架构预留 + 定向加练"**的组合拳：

• 架构侧：RoPE 大 base（1e7）+ GQA，让模型天生具备长序列外推能力和低显存占用；
• 训练侧：在 14.5T Token 基础预训练之上，用仅 396B Token 的第 5 阶段渐进式训练，把窗口从 128K 平滑扩展到 512K，且几乎不牺牲原有能力。

这套"小成本换大窗口"的思路，也为其他开源模型做长上下文扩展提供了可参考的工程范式。