研究对象先钉死
本章只讨论 https://github.com/xai-org/grok-1 的 main 分支 commit 7050ed204b8206bb8645c7b7bbef7252f79561b0，核验日期 2026-07-26。仓库包含 Python/JAX 推理示例、tokenizer 和 checkpoint 下载说明；权重另行下载。
一张规格表
| 项目 | Grok-1 公开值 | 在哪里 |
|---|---|---|
| 总参数 | 314B | README / 发布说明 |
| Transformer 层 | 64 | run.py: TransformerConfig |
| hidden / embedding size | 6,144 | emb_size=48*128 |
| Query heads | 48 | num_q_heads=48 |
| KV heads | 8 | num_kv_heads=8 |
| head size | 128 | key_size=128 |
| 专家 | 8 个，token 选 2 个 | num_experts / num_selected_experts |
| widening factor | 8 | widening_factor=8 |
| 词表 | 131,072 | vocab_size=128*1024 |
| 上下文 | 8,192 | sequence_len=8192 |
| 数值与部署 | bfloat16 前向、8-bit 权重结构、JAX mesh 分片 | model.py, runners.py |
数据怎样流过 64 层
LanguageModel.__call__() 把 token ID 查成 embedding，乘以输入缩放；Transformer.__call__() 建立 padding 与 causal mask，循环 64 次创建 DecoderLayer；最后 RMSNorm，使用 tied embedding 矩阵投影到词表 logits。
每个 DecoderLayer 有两个残差子层：MHABlock 完成归一化、Attention 与残差；随后归一化结果进入 MoELayer，两个专家结果按 router gate 合并，再加回主干。
Attention 子层
Grok-1 使用 48Q/8KV 的 GQA、128 维 head、RoPE 和 KV Cache。Attention logit 经过缩放与 tanh soft-cap，softmax 在 fp32 中完成。这样的工程细节比一句“用了 Transformer”更能解释数值稳定性和推理成本。
MoE 子层
Router 将每个 6,144 维 token 表示投影为 8 个 logits，softmax 后 Top-2。每个专家是带门控的前馈块：两个输入投影，其中一路 GELU，逐元素相乘，再投影回 hidden size。被选输出乘 gate 后求和。
314B 应该怎样理解
“314B”是所有参数的总量，不是激活参数的精确值。粗略地说，MoE 把大量前馈参数分散到 8 个专家，每 token 只访问两个；但共享 embedding、Attention、norm、router 与输出投影仍被使用。实际 FLOPs 还取决于序列长度、batch、量化、kernel、设备通信和缓存。
不要做这两个错误计算：
314B × 2/8 = 78.5B：把全模型都当成专家参数，忽略共享部分。- “只激活 25%，所以显存只需 25%”：即使单 token 只走两个专家，推理设备仍需能访问全部专家权重，除非按专家分布与调度；README 也警告需要极大 GPU 内存。
开放权重不等于完整开源训练
这次发布包括权重和推理代码，许可证宽松，但没有完整训练数据、训练代码、优化器配置、数据治理过程、后训练流水线或生产服务栈。准确称呼是 open weights + inference example，而不是“整个 Grok 系统开源”。
此外，checkpoint 是未经对话微调的基础模型。直接运行它的输出风格、指令跟随与安全行为，不应与当时或今天的 Grok 产品等同。
设计取舍
- 大 MoE 容量：增加参数容量，同时限制每 token 激活专家；换来路由、通信和负载均衡复杂度。
- GQA：降低 KV Cache；共享 K/V 可能限制部分表达。
- 8-bit 权重支持：降低存储和带宽；量化误差与 kernel 支持成为新变量。
- JAX + mesh sharding：清楚表达张量如何映射设备；部署需要理解设备拓扑。
本章检查点
Grok-1 是一座公开的“现代解码器实验室”：你能确认每个关键张量和调用。但它是一张 2023 年基础模型快照，不是当前 Grok 产品的剖面图。