有关 GLM-4-0414 的 SGLang 推理支持

在智谱的 GLM-4-0414 系列模型发布后,我观察到一个有意思的现象。GLM-4-0414 相较于 GLM-4 修改了模型结构,初期仅支持通过 transformers 库推理,但你去搜索它有哪些推理支持的话,什么也找不到,被大量的垃圾自媒体文章淹没,很多相较于官方模型仓库的 README 来说都是 0 信息量。

有些垃圾自媒体,可能连个 AI 都不如。

说回来 GLM-4-0414,虽然它名字看起来是 GLM-4 的延续,但是仔细看 config.json 你就会发现,GLM-4 的模型结构是 ChatGLMModel,而 GLM-4-0414 的模型结构是 Glm4ForCausalLM。这就导致很多推理框架都要对其进行重新适配。

vllm 可能是最早支持 GLM-4-0414 的开源框架,看起来是智谱员工提的 PR,但是第一版实现有两个 bug,会导致模型加载失败或者输出错误结果。由于对 GLM-Z1-9B-0414 模型在报告中声称大幅优于 DeepSeek-R1-Distill-Qwen-14B 过于好奇,我忍不住自己尝试在 SGLang 里适配了一下 GLM4-4-0414,见 PR#5485

嗯,学到了很多东西,也踩了两个 bug 其中的一个,哭。就说说这两个 bug 吧,其实都是源自于对 vllm/SGLang 库算子的不了解。

在 transformers 库中,很多算子仅实现其算子名表示的朴素的功能,但是 vllm/SGLang 代码库的一些算子,除了其朴素功能以外,往往还通过更多的参数实现对前后计算的融合、简化,导致其实际计算需要深入探究。

BUG 1: 融合 RMSNorm

以比较基础的 LLaMA 为例,transformers 中的 LLaMA DecoderLayer 的实现是这样的(删去了一些无关细节),这个跟大家理解的 Transformer 模型结构伪代码是容易对上的。

def forward(): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, self_attn_weights = self.self_attn() hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) return outputs

但是 vllm/SGLang 中的实现是这样的:

def forward(): # Self Attention if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn(...) # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) return hidden_states, residual

粗看你会有点懵,仔细研究就会发现,SGLang 里面通过给 RMSNorm 传入第二个参数,实现了 Norm 与 Add 的融合。但是这种融合需要调整计算顺序,影响了参数和返回值的类型,并且也影响了最后一次 model.norm() 的计算。

相似的还有 SiluAndMul() 算子。

BUG 2: get_rope 参数

GLM-4-0414 的 config.json 中提供了 partial_rotary_factor 参数,作用于 head_dim 上。有两种应用方法,一种是提前计算好 rotary_dim = int(partial_rotary_factor * self.head_dim),然后把这个参数传进去;另一种是令 rotary_dim = head_dim,然后传入 partial_rotary_factor。

def get_rope( head_size: int, rotary_dim: int, max_position: int, base: int, is_neox_style: bool = True, rope_scaling: Optional[Dict[str, Any]] = None, dtype: Optional[torch.dtype] = None, partial_rotary_factor: float = 1.0, ) -> RotaryEmbedding: ... if partial_rotary_factor < 1.0: rotary_dim = int(rotary_dim * partial_rotary_factor)

我重复犯的这个错误就是将 partial_rotary_factor 应用了两遍,rotary_dim 也计算了,partial_rotary_factor 参数也传了(因为参考了 vllm 实现和旧的 SGLang chatglm 实现,甚至看到别人提交的 bugfix 都不认为这里有错),其实就相当于应用了 partial_rotary_factor * partial_rotary_factor。这个 BUG 的现象就是导致 GLM-4 模型的输出大概率陷入死循环而无法结束。

BTW:transformers 库里的 glm4 代码没有读这个参数,而是硬编码在 apply_rotary_pos_emb() 算子中,所以这不是一个可调参数。

PR#5485 还在 Review 中,有需要的朋友可以使用 https://github.com/solrex/sglang/tree/glm4 这个分支支持 GLM-4-0414 系列模型在 SGLang 的推理。

git clone -b glm4 https://github.com/solrex/sglang.git
cd sglang
pip install -e "python[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python
python3 -m sglang.launch_server --model /workspace/GLM-Z1-9B-0414 --enable-torch-compile --torch-compile-max-bs 128 --cuda-graph-max-bs 128 --host 0.0.0.0 --port 8000

LLM 推理优化 Continuous Batching 及其实现

原理

Continuous Batching 是 LLM 推理优化的一项技术,作为这篇文章的知识背景不再赘述,目前流传最广的参考资料是这篇:《How continuous batching enables 23x throughput in LLM inference while reducing p50 latency》。它也有中文翻译,感兴趣可以搜一下,先看看。

图片来自:Anyscale 博客

虽然这篇资料介绍了它的主要原理,尤其是这张简单易懂的图,但是实现与原理是存在差异的,因为工程实现要解决很多现实问题。

(Re)scheduling 开销

如原文所说,Continuous batching 还有个别名,叫做:batching with iteration-level scheduling,这里的 iteration 就是指一次 decode 计算。也就是说在每次 decode 的迭代过程中,做 batch 的调度调整。

但调度本身不是无代价的,它可能涉及到接收和处理新的输入请求,重新组织输入数据的形状,甚至各种状态的重新初始化,这些都需要消耗 CPU 时间。这也就意味着在这段时间里,GPU 是闲着的,GPU 没有得到充分利用。

所以在实现时,程序并不会真的每个 iteration 都做 scheduling,目前看到有两种做法:

  • 合理间隔调度。比如每 16 次 decode 计算后,检查一下是否有新的输入,以及是否有空闲的槽位,然后对 batch 做一次调度调整。这能够显著降低调度的开销(TGIlmdeployvLLM)。
  • 排队比例调度。比如当前 batch 中有 10 个请求的 decode 正在进行,而排队中有 12 个请求,超过了排队比例 1.2,那么就启动一次调度调整(TGI)。

KV Cache 重读

如果真的像图中那样,每个生成 Token 的 decode iteration 与一个 prompt token 的计算对齐,那 KV Cache 的利用就会很糟糕。因为它们需要在 Global Memory 与 Shared Memory 之间反复搬运,而且每次搬进来以后只被使用一次。

这本质上是 prefill 阶段(prompt 计算)与生成阶段的计算特性不同导致的,prefill 阶段并不适合 one-by-one 的 token 计算。

所以在实现时,程序并不会真的做 prefill 和生成的 token 对齐调度。目前看到的调度方法有三种:

  • 在重调度时,如果有新的请求进来,那么将新请求的 prefill 计算和要进行的 decode 计算做合并(Orca、vLLM-prev)。
  • 在重调度时,如果有新的请求进来,那么先对新的请求做 prefill 计算,然后再合并所有进行中的请求做 decode 计算(TGI、vLLM、lmdeploy)。
  • 先根据 decode 耗时估算出来每次 decode 同时能做 prefill 的 token 数量,在重调度时,如果有新的请求进来,对新请求的 prefill 按照上面的估算进行分块,然后将分块的 prefill 和其它请求的 decode 计算融合在一起,一定程度上减少了 KV Cache 重读的次数,又避免先做 prefill 计算带来的 Token 生成延时增加(Sarathi-Serve+vLLM)。

可调优能力

LLM 推理服务往往用于生产环境,而生产环境面临的情况是复杂多样的。

  • 对于做阅读理解的应用来说,Prompt 可能会非常长,但生成的内容可能会非常短,开发人员可能会更追求吞吐;
  • 对于聊天应用来说,Prompt 可能较短,生成的内容也不会太长,开发人员可能会更追求延迟;
  • 对于创作类应用来说,Prompt 可能很短,生成的内容会更长,开发人员可能会更追求首 Token 延迟。

对 Continuous Batching 实现来说,就要求它调度策略尽量清晰,并且参数可调。所以更灵活的实现,未来可能会更受欢迎。