大模型推理框架 RTP-LLM 架构解析_AI阅读总结 — 包阅AI

我们首先对RTP-LLM，该系统由多个层次组成，从最外层的InferenceApp到最内层的AsyncModel，每一层都有特定的职责。

最外层是InferenceApp，其核心是一个基于FastAPI的异步服务。这一层负责接受外部的 HTTP 请求，并将请求传递给内部的推理服务器。

InferenceServer是实际处理推理请求的服务器，内部有多个InferenceWorker负责具体的任务。每个InferenceWorker包含一个Pipeline，包含了用于处理请求的异步模型AsyncModel。

AsyncModel是整个系统的核心组件，负责具体的推理任务。它包含两个主要部分：Scheduler和ModelExecutor。

当 HTTP 请求到达时，系统会依次进行以下处理：

1. Prompt 构建：根据请求内容，构建推理所需的 prompt。

2. Tokenize：将 prompt 转换为模型可以理解的 token。

处理完请求后，系统会将任务通过model.enqueue放入队列中，准备进行后续处理。

Decoder Engine是异步推理的核心，包含两个部分：Scheduler和ModelExecutor。

Scheduler

Scheduler负责调度任务，它会从waiting_streams中选择需要处理的任务，并安排执行。在后台，Scheduler会不断运行step函数：

1. Scheduler.schedule获取 batch_query。

2. ModelExecutor.process执行推理任务。

ModelExecutor

ModelExecutor负责实际的模型执行过程：

1. Process BatchQuery：处理 batch_query。

2. Token Embedding：将 img token 转换为 img embedding。

3. CUDA 预测：将输入数据送入 CUDA 引擎进行预测。

Decoder函数会调用step函数，将输入数据放入 batch_query，然后开始等待step函数完成推理，并将结果返回。

该框架在模型的加载过程中，通常包含两个核心部分：模型的权重和配置文件。本小节将通过一个具体案例——添加对新模型（如OPT模型）的支持，来详细探讨如何加载模型的配置文件。

首先给出一个最终的配置加载的实现如下, 具体逻辑为将模型的配置对应到框架的配置上, 因为每个模型的配置文件的变量的名字都不尽相同，比如表示transformer模型多头注意力机制的head num的变量名，有的模型可能叫num_attention_heads(opt)，有的模型可能叫n_head(starcoder)，所以需要对应到框架统一的变量体系下。

classOpt(GPT):    """Opt"""    @classmethod    def _create_config(cls, ckpt_path: str):        offset = 2        config_dict = get_config_from_path(ckpt_path)        config = GptInitModelParameters(            head_num=config_dict['num_attention_heads'],            size_per_head=config_dict['hidden_size'] // config_dict['num_attention_heads'],            layer_num=config_dict.get('num_hidden_layers', 12),            vocab_size=config_dict['vocab_size'],            max_seq_len=config_dict['max_position_embeddings'] + offset        )        config.layernorm_type = 'pre_layernorm'        config.norm_type = "layernorm"        config.has_post_decoder_layernorm = True        config.hidden_size = config_dict['hidden_size']        config.inter_size = config_dict["ffn_dim"]        config.has_positional_encoding = True        config.activation_type = 'relu'        config.add_special_tokens = True        config.special_tokens.eos_token_id = config_dict.get('eos_token_id', 2)        config.special_tokens.pad_token_id = config_dict.get('pad_token_id', 1)        config.special_tokens.bos_token_id = config_dict.get('bos_token_id', 2)        config.head_num_kv = config.head_num        return config

在这个框架中，常用的config变量名及其对应的含义如下：





























在加载配置文件后，框架会根据这些参数去加载权重等。例如，如果has_positional_encoding为True，框架会初始化nn.embedding(config.max_seq_len, config.hidden_dim)。当然，还有很多其他操作，这里就不展开了。如果要添加新模型，可以通过对比框架参数和Hugging Face模型文件下的config.json参数来进行编写。
1）在编写完新模型的支持代码并加载了模型文件配置之后, 需要确认一下模型是不是按照参数的正常含义去使用的。OPT模型就不是，它有着自己独特的位置编码（Positional Encoding）实现方式, 无论是position_ids还是nn.embedding，它都统一增加了一个偏移量，让人好一通debug，一直以为是模型权重没加载对，没想到是位置编码出了问题。


因此，当我们要为一个大模型添加新的支持时，要注意他的独特之处，比如位置编码的实现方式、权重的初始化方法、或者是某些特定的层结构。
classOPTLearnedPositionalEmbedding(nn.Embedding):    """    This module learns positional embeddings up to a fixed maximum size.    """
    def __init__(self, num_embeddings: int, embedding_dim: int):                        self.offset = 2        super().__init__(num_embeddings + self.offset, embedding_dim)
    def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0):        """`input_ids_shape` is expected to be [bsz x seqlen]."""        attention_mask = attention_mask.long()
                positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1
                positions = positions[:, past_key_values_length:]
        return super().forward(positions + self.offset)


2）inter_size是一个必须传递的参数，因为它在底层CUDA GEMM算子的执行中起着关键作用。如果忽略了这个参数，底层CUDA GEMM算子将会报错。
在本小节中，我们以相同的案例（添加对新模型如的支持），来展开如何加载模型的权重。
首先给出一个最终的模型权重加载的实现如下, 和模型的配置是差不多的，具体逻辑就是将模型的权重参数名对应到框架的权重参数名上, 因为每个模型的权重参数名定义每个都不一样，所以需要对应到框架统一的变量体系下。
classOptWeightInfo(ModelDeployWeightInfo):        def _get_weight_info(self):        layer_weights = [                        WeightInfo(W.pre_ln_gamma, [CkptWeightInfo('model.decoder.layers.{i}.self_attn_layer_norm.weight', identity)], identity),            WeightInfo(W.pre_ln_beta, [CkptWeightInfo('model.decoder.layers.{i}.self_attn_layer_norm.bias', identity)], identity),
                        WeightInfo(W.attn_qkv_w, [                CkptWeightInfo('model.decoder.layers.{i}.self_attn.q_proj.weight', identity),                CkptWeightInfo('model.decoder.layers.{i}.self_attn.k_proj.weight', identity),                CkptWeightInfo('model.decoder.layers.{i}.self_attn.v_proj.weight', identity),                ], functools.partial(merge_qkv_hf)),
            WeightInfo(W.attn_qkv_b,[                CkptWeightInfo('model.decoder.layers.{i}.self_attn.q_proj.bias', identity),                CkptWeightInfo('model.decoder.layers.{i}.self_attn.k_proj.bias', identity),                CkptWeightInfo('model.decoder.layers.{i}.self_attn.v_proj.bias', identity),                ], functools.partial(merge_qkv_b)),                        WeightInfo(W.attn_o_w,[CkptWeightInfo('model.decoder.layers.{i}.self_attn.out_proj.weight', identity)], transpose),            WeightInfo(W.attn_o_b,[CkptWeightInfo('model.decoder.layers.{i}.self_attn.out_proj.bias', identity)], identity),
                        WeightInfo(W.post_ln_gamma,[CkptWeightInfo('model.decoder.layers.{i}.final_layer_norm.weight', identity)], identity),            WeightInfo(W.post_ln_beta,[CkptWeightInfo('model.decoder.layers.{i}.final_layer_norm.bias', identity)], identity),                                    WeightInfo(W.ffn_w3,[CkptWeightInfo('model.decoder.layers.{i}.fc1.weight', identity)], transpose),            WeightInfo(W.ffn_b3,[CkptWeightInfo('model.decoder.layers.{i}.fc1.bias', identity)], identity),            WeightInfo(W.ffn_w2,[CkptWeightInfo('model.decoder.layers.{i}.fc2.weight', identity)], transpose),            WeightInfo(W.ffn_b2,[CkptWeightInfo('model.decoder.layers.{i}.fc2.bias', identity)], identity),        ]                weights = [            WeightInfo(W.embedding, [CkptWeightInfo('model.decoder.embed_tokens.weight', concat_0)], identity),            WeightInfo(W.positional_embedding,[CkptWeightInfo('model.decoder.embed_positions.weight', identity)], identity),            WeightInfo(W.final_ln_gamma,[CkptWeightInfo('model.decoder.final_layer_norm.weight', identity)], identity),            WeightInfo(W.final_ln_beta,[CkptWeightInfo('model.decoder.final_layer_norm.bias', identity)], identity),            WeightInfo(W.lm_head,[CkptWeightInfo('lm_head.weight', identity)], identity),        ]                return ModelWeightInfo(layer_weights=layer_weights, weights=weights)

class OPT_125M(GPT):
    @staticmethod    def get_weight_cls():        return OptWeightInfo




框架中对于权重信息分为了layer weights(对应的是decoder layer的权重参数)和weights(对应的是decoder layer之外的权重参数，比如lm_head、token embedding部分 ), 上图中是框架的参数和模型结构的一个对应，首先我们需要看一下pytorch.bin.index里面的参数名称，然后和框架中的参数名对应上，比如：

WeightInfo(W.pre_ln_gamma, [CkptWeightInfo('model.decoder.layers.{i}.self_attn_layer_norm.weight', identity)], identity),


这里W.pre_ln_gamma表示的是Attention和FFN层之前的Layer Norm层，而在OPT模型中，参数名称为model.decoder.layers.

WeightInfo(W.embedding, [CkptWeightInfo('model.decoder.embed_tokens.weight', concat_0)], identity),



这里W.embedding表示的是token的embedding层，在OPT模型中参数名model.decoder.embed_tokens.weight。


https://github.com/alibaba/rtp-llm
https://www.zhihu.com/people/66-34-28-46-64/posts
https://link.zhihu.com/?target=https%3A//github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py%23L566