从‘你好’到完整回复：一步步图解ChatGLM2-6B的推理循环（附代码片段）
深入解析ChatGLM2-6B的推理机制：从输入到生成的完整流程
当开发者第一次接触大型语言模型时，最令人着迷的莫过于观察一个简单输入如何逐步转化为连贯的输出。本文将以ChatGLM2-6B为例，详细拆解这个"思考"过程的每一个环节，特别聚焦于模型推理阶段的两层关键循环结构。通过实际代码片段和流程图解，我们将揭示这个看似"黑盒"的系统内部如何运作。
1. 模型推理的整体架构
ChatGLM2-6B采用Prefix Decoder-only架构，这是一种在传统自回归模型基础上改进的设计。与GPT系列使用的Causal Decoder-only不同，它在处理前缀部分时采用双向注意力机制，而在生成部分则切换为单向注意力。这种混合设计使其既能理解上下文，又能流畅生成文本。
推理过程的核心是两层嵌套循环：
- 外层循环：控制token的逐个生成，直到遇到结束符
- 内层循环：28层GLMBlock的连续处理，计算每个位置的表示
while True: # 外层循环：生成token序列 for i in range(28): # 内层循环：28层GLMBlock处理 # 每一层的计算逻辑 # 生成下一个token if token == eos_token: break2. 输入预处理与分词
当用户输入"你好"这样的简单文本时，模型会先进行一系列预处理步骤。ChatGLM2-6B会自动将原始输入包装成特定的对话格式：
原始输入："你好" → 处理后："[Round 1]\n\n问：你好\n\n答："
分词阶段采用WordPiece算法，将文本转换为token ID序列。ChatGLM2-6B的词表大小为65024，包含以下几种元素：
- 基础词汇
- 特殊控制符号（如[CLS]、[SEP]）
- 对话专用标记（如[Round]、\n\n问：）
- 多语言字符
分词后的结果是一个整数数组，例如：
[64790, 64792, 36474, 36147, 64286, 64286, 35180, 36474, 36147, 64286, 64286, 35181]3. 嵌入层与位置编码
分词后的ID序列通过嵌入层转换为高维向量表示。ChatGLM2-6B的嵌入维度为4096，这意味着每个token被映射为一个4096维的稠密向量。
嵌入层的关键特性：
| 特性 | 说明 |
|---|---|
| 可训练性 | 预训练阶段学习得到，推理时固定 |
| 共享权重 | 嵌入层与输出层共享部分参数 |
| 位置编码 | 采用旋转位置编码(RoPE)注入位置信息 |
嵌入层的输出形状为[序列长度, 批大小, 隐藏维度]，对于单个输入"你好"，典型形状为[17, 1, 4096]。
4. GLMBlock的28层处理
ChatGLM2-6B的核心由28个相同的GLMBlock堆叠而成，每个Block包含以下关键组件：
- RMSNorm：替代传统LayerNorm的归一化方法
- 注意力机制：混合使用双向和单向注意力
- MLP层：采用SwiGLU激活函数增强非线性能力
4.1 注意力机制详解
每个GLMBlock中的注意力模块执行以下操作：
# 伪代码表示注意力计算过程 def attention(q, k, v): scores = q @ k.transpose(-2, -1) / sqrt(dim) weights = softmax(scores) output = weights @ v return output实际实现中，ChatGLM2-6B采用了以下优化：
- KV Cache：缓存历史Key-Value对，避免重复计算
- 多查询注意力：Key和Value共享部分注意力头
- 旋转位置编码：更好地处理长序列
4.2 MLP层的独特设计
GLMBlock中的MLP层采用扩展-压缩的设计：
输入(4096) → 上投影(27392) → SwiGLU → 下投影(4096)这种设计大幅增加了中间表示能力，同时保持输入输出维度一致。SwiGLU激活函数的公式为：
SwiGLU(x) = x * sigmoid(βx) * Wx5. 生成下一个token
经过28层GLMBlock处理后，模型需要决定输出哪个token。这一过程分为三个步骤：
- 最终归一化：对最后一层的输出进行RMSNorm
- 线性投影：将隐藏状态映射到词表空间(65024维)
- 采样策略：根据logits选择下一个token
常见的采样方法包括：
- 贪心搜索：直接选择概率最高的token
- 温度采样：通过温度参数控制随机性
- Top-k/p采样：限制候选token范围
# 生成下一个token的示例代码 logits = last_hidden_state @ embedding_matrix.T probs = softmax(logits / temperature) next_token = sample(probs) # 根据策略采样6. 推理优化技术
在实际部署中，ChatGLM2-6B采用了多项推理优化技术：
KV Cache：
- 缓存历史轮次的Key-Value矩阵
- 避免重复计算，显著提升生成速度
- 内存占用随序列长度线性增长
量化推理：
- 支持FP16/INT8/INT4量化
- 减少显存占用，提升吞吐量
- 精度损失可控
批处理优化：
- 动态批处理技术
- 处理不同长度输入时自动填充
- 最大化GPU利用率
7. 完整推理流程示例
让我们通过一个具体例子，观察输入"你好"的完整处理过程：
预处理：
- 原始输入："你好"
- 格式化后："[Round 1]\n\n问：你好\n\n答："
分词：
- 转换为ID序列：[64790, 64792, ..., 35181]
嵌入查找：
- 每个ID映射为4096维向量
- 输出形状：[17, 1, 4096]
GLMBlock处理：
- 28层顺序处理
- 每层更新隐藏状态
生成第一个token：
- 计算logits
- 采样得到"你"(ID:36474)
循环生成：
- 将"你"加入输入
- 重复过程直到生成
在实际测试中，输入"你好"可能得到如下响应：
"[Round 1]\n\n问：你好\n\n答：你好！有什么我可以帮助你的吗？"理解ChatGLM2-6B的推理机制后，开发者可以更有效地进行模型调优和应用开发。例如，通过调整温度参数控制生成多样性，或使用特定的停止标记来精确控制输出长度。在部署时，合理配置KV Cache大小和量化策略，可以在资源受限环境下实现最佳性能。