大模型推理框架 RTP-LLM 架构解析
RTP-LLM 是阿里巴巴智能引擎团队推出的大模型推理框架，支持了包括淘宝、天猫、闲鱼、菜鸟、高德、饿了么、AE、Lazada 等多个业务的大模型推理场景。RTP-LLM 与当前广泛使用的多种主流模型兼容，使用高性能的 CUDA kernel, 包括 PagedAttention、FlashAttention、FlashDecoding 等，支持多模态、LoRA、P-Tuning、以及 WeightOnly 动态量化等先进功能，已在众多 LLM 场景中得到实际应用与检验。
本篇文章介绍了 RTP-LLM 的整体架构，并着重分析了模型加载过程中的核心部分：模型的权重和配置文件。本文主要由社区用户 mingming 贡献，特此感谢其对项目的支持。
一、架构介绍
我们首先对 RTP-LLM，该系统由多个层次组成，从最外层的 InferenceApp 到最内层的 AsyncModel，每一层都有特定的职责。
1.1 InferenceApp 和 FastAPI 服务
最外层是 InferenceApp，其核心是一个基于 FastAPI 的异步服务。这一层负责接受外部的 HTTP 请求，并将请求传递给内部的推理服务器。
1.2 InferenceServer 和 InferenceWorker
InferenceServer 是实际处理推理请求的服务器，内部有多个 InferenceWorker 负责具体的任务。每个 InferenceWorker 包含一个 Pipeline，包含了用于处理请求的异步模型 AsyncModel。
二、异步模型 AsyncModel 的工作流程
AsyncModel 是整个系统的核心组件，负责具体的推理任务。它包含两个主要部分：Scheduler 和 ModelExecutor。
2.1 接收请求
当 HTTP 请求到达时，系统会依次进行以下处理：
Prompt 构建：根据请求内容，构建推理所需的 prompt。
Tokenize：将 prompt 转换为模型可以理解的 token。
2.2 任务排队
处理完请求后，系统会将任务通过 model.enqueue 放入队列中，准备进行后续处理。
2.3 Decoder Engine
Decoder Engine 是异步推理的核心，包含两个部分：Scheduler 和 ModelExecutor。
Scheduler
Scheduler 负责调度任务，它会从 waiting_streams 中选择需要处理的任务，并安排执行。在后台，Scheduler 会不断运行 step 函数：
Scheduler.schedule获取 batch_query。ModelExecutor.process执行推理任务。
ModelExecutor
ModelExecutor 负责实际的模型执行过程：
Process BatchQuery：处理 batch_query。
Token Embedding：将 img token 转换为 img embedding。
CUDA 预测：将输入数据送入 CUDA 引擎进行预测。
2.4 执行推理
Decoder 函数会调用 step 函数，将输入数据放入 batch_query，然后开始等待 step 函数完成推理，并将结果返回。
三、模型配置加载
该框架在模型的加载过程中，通常包含两个核心部分：模型的权重和配置文件。本小节将通过一个具体案例——添加对新模型（如 OPT 模型）的支持，来详细探讨如何加载模型的配置文件。
3.1 完整实现
首先给出一个最终的配置加载的实现如下, 具体逻辑为将模型的配置对应到框架的配置上, 因为每个模型的配置文件的变量的名字都不尽相同，比如表示 transformer 模型多头注意力机制的 head num 的变量名，有的模型可能叫 num_attention_heads(opt)，有的模型可能叫 n_head(starcoder)，所以需要对应到框架统一的变量体系下。
3.2 实现解析
在这个框架中，常用的 config 变量名及其对应的含义如下：
head_num: 在 Transformer 模型的自注意力（Self-Attention）机制中，表示注意力机制被分割成的头的数量。每个头负责学习输入数据的不同表示子空间。
head_num_kv: 指定了在自注意力机制中用于处理键（Key）和值（Value）的注意力头的数量。在标准的 Transformer 模型中，查询（Query）、键和值通常使用相同数量的注意力头进行处理。然而，某些模型（如 llama2 70B）可能设计为使用不同数量的头来处理键和值。
size_per_head: 表示每个注意力头的维度大小。通常，模型的总隐藏层维度（hidden_size）被平均分配给所有的头。例如，如果hidden_size是 512 且head_num是 8，那么size_per_head将是 64。
layer_norm_type: 指定了模型中使用的层归一化（Layer Normalization）的类型或位置。例如，'pre_layernorm'表示归一化操作在每个子层（如注意力层或前馈网络层）之前进行；而'post_layernorm'则表示归一化操作在每个子层之后进行。
vocab_size: 表示词表的大小。
max_seq_len: 表示模型能够处理的最长序列长度。
norm_type: 指定模型使用的归一化技术，如“Layernorm”或“RMSNorm”。
has_post_decoder_layernorm: 表示模型是否在解码器的最后添加了额外的层归一化层。
hidden_size: 表示隐藏层的维度。
inter_size: 表示 FFN 中间层的维度。在 LLM 的 FFN 中，会先进行 up_proj，然后进行 down_proj，这里的inter_size就是 up_proj 后的维度。
has_positional_encoding: 表示模型是否使用位置编码，且位置编码是否通过nn.embedding得到
activation_type: 指定激活函数的类型，如 relu、gelu 或 swiglu。
add_special_tokens: 表示是否在模型的词汇表中添加特殊标记，如 EOS（结束标记）、PAD（填充标记）和 BOS（开始标记）。
special_tokens: 包含模型使用的特殊标记的 ID，如eos_token_id、pad_token_id和bos_token_id。
在加载配置文件后，框架会根据这些参数去加载权重等。例如，如果has_positional_encoding为 True，框架会初始化nn.embedding(config.max_seq_len, config.hidden_dim)。当然，还有很多其他操作，这里就不展开了。如果要添加新模型，可以通过对比框架参数和 Hugging Face 模型文件下的config.json参数来进行编写。
3.3 注意事项
1）在编写完新模型的支持代码并加载了模型文件配置之后, 需要确认一下模型是不是按照参数的正常含义去使用的。OPT 模型就不是，它有着自己独特的位置编码（Positional Encoding）实现方式, 无论是position_ids还是nn.embedding，它都统一增加了一个偏移量，让人好一通 debug，一直以为是模型权重没加载对，没想到是位置编码出了问题。
因此，当我们要为一个大模型添加新的支持时，要注意他的独特之处，比如位置编码的实现方式、权重的初始化方法、或者是某些特定的层结构。
2）inter_size是一个必须传递的参数，因为它在底层 CUDA GEMM 算子的执行中起着关键作用。如果忽略了这个参数，底层 CUDA GEMM 算子将会报错。
四、模型权重加载
在本小节中，我们以相同的案例（添加对新模型如的支持），来展开如何加载模型的权重。
4.1 完整实现
首先给出一个最终的模型权重加载的实现如下, 和模型的配置是差不多的，具体逻辑就是将模型的权重参数名对应到框架的权重参数名上, 因为每个模型的权重参数名定义每个都不一样，所以需要对应到框架统一的变量体系下。
框架中对于权重信息分为了 layer weights(对应的是 decoder layer 的权重参数)和 weights(对应的是 decoder layer 之外的权重参数，比如 lm_head、token embedding 部分 ), 上图中是框架的参数和模型结构的一个对应，首先我们需要看一下 pytorch.bin.index 里面的参数名称，然后和框架中的参数名对应上，比如：
这里 W.pre_ln_gamma 表示的是 Attention 和 FFN 层之前的 Layer Norm 层，而在 OPT 模型中，参数名称为 model.decoder.layers.
这里 W.embedding 表示的是 token 的 embedding 层，在 OPT 模型中参数名 model.decoder.embed_tokens.weight。
参考链接
[01] RTP-LLM 项目地址：
https://github.com/alibaba/rtp-llm
[02] 原作者博客：
https://www.zhihu.com/people/66-34-28-46-64/posts
[03] OPT 模型：
版权声明: 本文为 InfoQ 作者【阿里技术】的原创文章。
原文链接:【http://xie.infoq.cn/article/091a2323bde250c83ce174673】。文章转载请联系作者。
专注分享阿里技术的丰富实践和前沿创新。 2022-05-24 加入
阿里技术的官方号，专注分享阿里技术的丰富实践、前沿洞察、技术创新、技术人成长经验。阿里技术，与技术人一起创造成长与成就。
评论