从‘你好’到完整回复:一步步图解ChatGLM2-6B的推理循环(附代码片段)

发布时间:2026/6/12 16:47:27
从‘你好’到完整回复:一步步图解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大小和量化策略可以在资源受限环境下实现最佳性能。