InferLLM 图结构代码分析

InferLLM 的 graph 目录包含了不同模型架构的计算图实现，主要负责构建和管理大型语言模型的计算流程。下面对其代码结构和功能进行详细分析。

1. 整体架构

graph 目录主要包含以下几个模型的实现：

1. ChatGLM 系列：ChatGLM、ChatGLM2、ChatGLM3
2. Llama 系列：GgmlLlama、LlamaLike (包括 Baichuan 和 Llama2)

所有这些模型都继承自基础的 Graph 类，并实现了特定的模型结构和加载逻辑。

2. 基础组件

2.1 Graph 工厂方法

std::shared_ptr<Graph> Graph::make_graph(
        UserConfig model_config, Device* device, const std::string& name) {
    if (name == "llama") {
        return std::make_shared<GgmlLlamaGraph>(model_config, device, name);
    } else if (name == "chatglm") {
        return std::make_shared<ChatGLMGraph>(model_config, device, name);
    } else if (name == "chatglm2") {
        return std::make_shared<ChatGLMGraph2>(model_config, device, name);
    } else if (name == "chatglm3") {
        return std::make_shared<ChatGLMGraph3>(model_config, device, name);
    } else if (name == "baichuan" || name == "llama2") {
        return std::make_shared<LlamaLikeGraph>(model_config, device, name);
    } else {
        INFER_ASSERT(0, "unsupported model.");
    }
}

这个工厂方法根据模型名称创建相应的图实现。

2.2 模型参数结构

struct Header {
    int param_offset;
    int param_length;
    int vocab_offset;
    int vocab_length;
    int tensor_offset;
};

struct Param {
    int hidden_size;
    int n_heads;
    int n_layers;
    int embd_dim;
    int fc_hidden;
    int vacab_size;
    int multi_query = 0;
    int multi_query_group_num = 1;
};

这些结构体用于存储模型文件的头信息和模型参数。

3. 模型实现分析

3.1 ChatGLM 系列

3.1.1 ChatGLMGraph

ChatGLM 是最早的 ChatGLM 系列模型实现。

权重别名设置：

void ChatGLMGraph::set_weights_alias() {
    m_weights_name_aliases.clear();
    // clang-format off
    m_weights_name_aliases = {
            {"transformer.word_embeddings.weight", "tok_embeddings.weight"},
            {"transformer.layers.x.input_layernorm.weight", "layers.x.attention_norm.weight"},
            // ... 其他权重别名
    };
    // clang-format on
}

这个方法建立了原始权重名称和内部使用的权重名称之间的映射关系。

参数加载：

void ChatGLMGraph::load_param(
        std::shared_ptr<InputFile> fin, LlmParams& param,
        std::shared_ptr<Vocab> vocab) {
    Header header;
    // 加载模型头部信息
    fin->read_raw((char*)&header.param_offset, sizeof(header.param_offset));
    // ... 读取其他头部信息
    
    // 加载参数
    fin->seek(header.param_offset);
    fin->read_raw((char*)&param.n_embd, sizeof(param.n_embd));
    // ... 读取其他参数
    
    // 加载词汇表
    fin->seek(header.vocab_offset);
    vocab->load_vocab(fin, param.n_vocab);
    
    // 设置词汇表大小
    m_param.n_vocab = 130528;
    param.n_vocab = 130528;
    fin->seek(header.tensor_offset);
}

这个方法从模型文件中加载参数和词汇表。

构建计算图：

void ChatGLMGraph::construct_llm() {
    m_input = std::make_shared<Tensor>(device(), name() + ":input");
    std::shared_ptr<Tensor> input = m_input;
    //! embd
    input = add_module<EmbdModule>(
            this, input, m_param.n_embd, m_param.n_vocab, model_config(), device(), "");

    int nr_layer = m_param.n_layer;
    float scale = sqrt(2 * nr_layer);
    for (int i = 0; i < nr_layer; i++) {
        std::string name = "layers." + std::to_string(i);
        //! layer norm
        std::shared_ptr<Tensor> attention_input = input;
        auto norm_out_attention =
                add_one_opr_module<LayerNorm>(
                        this, OpIOs{attention_input}, device(),
                        name + ".attention_norm")
                        ->add_opr(
                                m_param.n_embd, /*mul*/ true, /*bias*/ true,
                                /*rms*/ false);
        //! attentin
        auto attention_output = add_module<AttentionModule<GlmAttention>>(
                this, norm_out_attention, m_param.n_embd, m_param.n_head, m_param.n_rot,
                m_param.n_ctx, model_config(), device(), name + ".attention", i,
                true /*fused_weights*/, true /*bias*/);
        //! add  norm_out_attention * scale + attention_output
        auto add_output = add_one_opr_module<Elemwise>(
                                  this, OpIOs{norm_out_attention, attention_output},
                                  device(), name + ".attention.Elemwise")
                                  ->add_opr(ElemMode::Add, scale);

        std::shared_ptr<Tensor> feed_forward_input = add_output;
        //! layer normal
        auto ffn_norm_out =
                add_one_opr_module<LayerNorm>(
                        this, OpIOs{feed_forward_input}, device(), name + ".ffn_norm")
                        ->add_opr(
                                m_param.n_embd, /*mul*/ true, /*bias*/ true,
                                /*rms*/ false);
        //! feed forward
        auto ffn_output = add_module<GlmFFNModule>(
                this, ffn_norm_out, m_param.n_embd, m_param.n_mult, model_config(),
                device(), name);
        //! add ffn_norm_out * scale + ffn_output
        input = add_one_opr_module<Elemwise>(
                        this, OpIOs{ffn_norm_out, ffn_output}, device(),
                        name + ".ffn.Elemwise")
                        ->add_opr(ElemMode::Add, scale);
    }
    //! the last layer
    m_output = add_module<HeadModule>(
            this, input, m_param.n_embd, m_param.n_vocab, model_config(), device(),
            "head", true);
}

这个方法构建了 ChatGLM 模型的计算图，包括：

1. 创建输入张量
2. 添加嵌入层
3. 循环添加多个 Transformer 层，每层包括：
   • 注意力层前的层归一化
   • 注意力层 (GlmAttention)
   • 残差连接
   • 前馈网络前的层归一化
   • 前馈网络 (GlmFFNModule)
   • 残差连接
4. 添加输出层 (HeadModule)

后处理 token：

void ChatGLMGraph::post_tokenize(std::vector<Vocab::Id>& input) {
    //! the begin token
    input.insert(input.begin(), 5);
    //! the end token
    input.push_back(130001);
    input.push_back(130004);
}

这个方法在 token 序列的开头和结尾添加特殊 token。

3.1.2 ChatGLMGraph2 和 ChatGLMGraph3

ChatGLM2 和 ChatGLM3 的实现与 ChatGLM 类似，但有一些区别：

1. 权重别名不同：反映了模型架构的变化
2. 参数加载增加了多查询注意力相关参数：

   int32_t multi_query;
   fin->read_raw((char*)&multi_query, sizeof(multi_query));
   param.is_multi_query = multi_query > 0;
   fin->read_raw((char*)&param.multi_query_group_num, sizeof(param.multi_query_group_num));
   

3. 构建计算图使用不同的注意力机制：
   • ChatGLM2 使用 Glm2MultiQueryAttention
   • ChatGLM3 可能使用其他变体

3.2 Llama 系列

3.2.1 GgmlLlamaGraph

GgmlLlamaGraph 实现了 GGML 格式的 Llama 模型。

加载方法：

void GgmlLlamaGraph::load(
        std::shared_ptr<InputFile> fin, LlmParams& param,
        std::shared_ptr<Vocab> vocab) {
    // 验证魔数
    uint32_t magic;
    uint32_t version = 0;
    fin->read_raw((char*)&magic, sizeof(magic));
    if (magic != 'ggml') {
        fin->read_raw((char*)&version, sizeof(version));
    }
    LlamaModelType model_type;
    if (magic == 'ggml' && version == 0) {
        model_type = LlamaModelType::LLAMA_FILE_VERSION_GGML;
    } else if (magic == 'ggmf' && version == 1) {
        model_type = LlamaModelType::LLAMA_FILE_VERSION_GGMF_V1;
    } else if (magic == 'ggjt' && version == 1) {
        model_type = LlamaModelType::LLAMA_FILE_VERSION_GGJT_V1;
    } else {
        INFER_ASSERT(0, "unsupported model type.");
    }
    
    // 加载参数
    fin->read_raw((char*)&param.n_vocab, sizeof(param.n_vocab));
    // ... 读取其他参数
    
    // 加载词汇表
    if (model_type == LlamaModelType::LLAMA_FILE_VERSION_GGJT_V1) {
        vocab->load_vocab_with_score(fin, param.n_vocab);
    } else {
        vocab->load_vocab(fin, param.n_vocab);
    }
    
    m_param = param;
    
    // 构建计算图
    construct_llm();
    collect_weights();
    
    // 设置权重别名并加载权重
    set_weights_alias();
    // ... 读取权重
}

这个方法处理了 GGML 格式的模型加载，包括识别不同的 GGML 版本。

构建计算图：

void GgmlLlamaGraph::construct_llm() {
    uint32_t embd = m_param.n_embd;
    uint32_t mult = m_param.n_mult;
    uint32_t head = m_param.n_head;
    uint32_t rot = m_param.n_rot;
    uint32_t ctx = m_param.n_ctx;
    uint32_t n_vocab = m_param.n_vocab;

    size_t nff = ((2 * (4 * embd) / 3 + mult - 1) / mult) * mult;
    m_input = std::make_shared<Tensor>(device(), name() + ":input");
    std::shared_ptr<Tensor> input = m_input;
    //! embd
    input = add_module<EmbdModule>(
            this, input, embd, n_vocab, model_config(), device(), "");

    int nr_layer = m_param.n_layer;
    for (int i = 0; i < nr_layer; i++) {
        std::string name = "layers." + std::to_string(i);
        //! layer norm
        std::shared_ptr<Tensor> attention_input = input;

        auto norm_out_attention = add_one_opr_module<LayerNorm>(
                                          this, OpIOs{attention_input}, device(),
                                          name + ".attention.norm")
                                          ->add_opr(embd);
        //! attentin
        auto attention_output = add_module<AttentionModule<LlamaAttention>>(
                this, norm_out_attention, embd, head, rot, ctx, model_config(),
                device(), name + ".attention", i);
        //! add
        auto add_output = add_one_opr_module<Elemwise>(
                                  this, OpIOs{attention_input, attention_output},
                                  device(), name + ".attention_add")
                                  ->add_opr(ElemMode::Add);

        std::shared_ptr<Tensor> feed_forward_input = add_output;
        //! layer normal
        auto ffn_norm_out =
                add_one_opr_module<LayerNorm>(
                        this, OpIOs{feed_forward_input}, device(), name + ".ffn.norm")
                        ->add_opr(embd);
        //! feed forward
        auto ffn_output = add_module<LlamaFFNModule>(
                this, ffn_norm_out, embd, nff, model_config(), device(), name);
        //! add
        input = add_one_opr_module<Elemwise>(
                        this, OpIOs{feed_forward_input, ffn_output}, device(),
                        name + ".ffn_add")
                        ->add_opr(ElemMode::Add);
    }
    //! the last layer
    m_output = add_module<HeadModule>(
            this, input, embd, n_vocab, model_config(), device(), "head");
}

这个方法构建了 Llama 模型的计算图，结构与 ChatGLM 类似，但使用了 LlamaAttention 和 LlamaFFNModule。

3.2.2 LlamaLikeGraph

LlamaLikeGraph 实现了类 Llama 架构的模型，如 Baichuan 和 Llama2。其实现与 GgmlLlamaGraph 类似，但加载方式和权重别名有所不同。

权重别名设置：

void LlamaLikeGraph::set_weights_alias() {
    m_weights_name_aliases.clear();
    m_weights_name_aliases = {
            {"model.embed_tokens.weight", "tok_embeddings.weight"},
            {"model.layers.x.input_layernorm.weight", "layers.x.attention.norm.weight"},
            {"model.layers.x.self_attn.q_proj.weight", "layers.x.attention.wq.weight"},
            {"model.layers.x.self_attn.k_proj.weight", "layers.x.attention.wk.weight"},
            {"model.layers.x.self_attn.v_proj.weight", "layers.x.attention.wv.weight"},
            // ... 其他权重别名
    };
}

LlamaLikeGraph 使用分离的 Q、K、V 权重，而不是融合的 QKV 权重。

参数加载：

void LlamaLikeGraph::load_param(
        std::shared_ptr<InputFile> fin, LlmParams& param,
        std::shared_ptr<Vocab> vocab) {
    // 加载参数
    fin->read_raw((char*)&param.n_embd, sizeof(param.n_embd));
    fin->read_raw((char*)&param.n_mult, sizeof(param.n_mult));
    fin->read_raw((char*)&param.n_head, sizeof(param.n_head));
    fin->read_raw((char*)&param.n_layer, sizeof(param.n_layer));
    fin->read_raw((char*)&param.n_rot, sizeof(param.n_rot));
    fin->read_raw((char*)&param.n_vocab, sizeof(param.n_vocab));
    
    // 加载词汇表
    vocab->load_vocab(fin, param.n_vocab);
    
    m_param = param;
}

LlamaLikeGraph 的参数加载相对简单，直接读取各个参数值。

4. 模块组织结构

InferLLM 的图结构采用模块化设计，主要包括以下几种模块：

4.1 基础模块

1. EmbdModule：实现词嵌入功能，将 token ID 转换为嵌入向量
2. LayerNorm：实现层归一化，用于稳定网络训练
3. Elemwise：实现元素级操作，如加法、乘法等
4. HeadModule：实现输出层，将隐藏状态转换为词汇表概率分布

4.2 注意力模块

1. AttentionModule：注意力机制的通用模块，可以使用不同的注意力实现
2. GlmAttention：ChatGLM 系列的注意力实现
3. Glm2MultiQueryAttention：ChatGLM2 的多查询注意力实现
4. LlamaAttention：Llama 系列的注意力实现

4.3 前馈网络模块

1. GlmFFNModule：ChatGLM 系列的前馈网络实现
2. LlamaFFNModule：Llama 系列的前馈网络实现

5. 计算图构建流程

所有模型的计算图构建流程大致相同：

1. 创建输入张量：

   m_input = std::make_shared<Tensor>(device(), name() + ":input");
   

2. 添加嵌入层：

   input = add_module<EmbdModule>(
           this, input, m_param.n_embd, m_param.n_vocab, model_config(), device(), "");
   

3. 循环添加 Transformer 层：

   for (int i = 0; i < nr_layer; i++) {
       // 添加注意力层
       // 添加前馈网络层
   }
   

4. 添加输出层：

   m_output = add_module<HeadModule>(
           this, input, m_param.n_embd, m_param.n_vocab, model_config(), device(),
           "head", true);
   

6. 模型加载流程

模型加载流程通常包括以下步骤：

1. 读取模型头部信息：获取参数、词汇表和张量的偏移量
2. 加载模型参数：读取嵌入维度、头数、层数等参数
3. 加载词汇表：读取词汇表信息
4. 构建计算图：根据模型参数构建计算图
5. 收集权重：收集计算图中的所有权重
6. 设置权重别名：建立原始权重名称和内部使用的权重名称之间的映射关系
7. 加载权重：从模型文件中读取权重数据

7. 模型推理流程

模型推理流程通常包括以下步骤：

1. 准备输入：将输入文本转换为 token ID
2. 前处理：添加特殊 token，如开始 token
3. 执行计算图：从输入到输出执行计算图
4. 采样：根据输出概率分布采样下一个 token
5. 后处理：将 token ID 转换为文本

8. 不同模型的特点

8.1 ChatGLM 系列

1. Pre-LayerNorm 架构：先进行层归一化，再进行注意力计算
2. 融合的 QKV 权重：使用一个权重矩阵计算 Q、K、V
3. 残差连接带缩放：残差连接时使用缩放因子 scale = sqrt(2 * nr_layer)
4. 带偏置的层归一化：层归一化使用偏置项

8.2 Llama 系列

1. RMSNorm：使用 RMSNorm 代替传统的层归一化
2. 旋转位置编码：使用旋转位置编码 (RoPE) 代替位置嵌入
3. SwiGLU 激活函数：在前馈网络中使用 SwiGLU 激活函数
4. 无偏置项：大多数层不使用偏置项

总结

InferLLM 的图结构代码实现了多种大型语言模型的计算图构建和执行。它采用模块化设计，支持不同的模型架构，如 ChatGLM 系列和 Llama 系列。通过分析这些代码，我们可以了解不同模型的结构特点和实现细节，以及 InferLLM 框架的设计思想和工作原理。

这种设计使得 InferLLM 可以灵活地支持新的模型架构，只需继承 Graph 类并实现相应的方法，就可以添加新的模型支持。同时，模块化的设计也使得代码更易于维护和扩展。