LlamaLike 模型代码结构与功能实现分析

llama_like.h 和 llama_like.cpp 文件实现了 InferLLM 框架中对类 Llama 架构模型（如 Baichuan 和 Llama2）的支持。这些文件定义了 LlamaLikeGraph 类，该类继承自基础的 Graph 类，并实现了特定于类 Llama 架构模型的加载和计算图构建逻辑。

1. llama_like.h 文件分析

1.1 文件结构

llama_like.h 文件结构简洁，主要包含以下几个部分：

1. 头文件引入：

   #include "core/graph.h"
   

   只引入了基础的图类头文件。

2. 命名空间声明：

   namespace inferllm {
   

   所有代码都在 inferllm 命名空间中。

3. LlamaLikeGraph 类定义：

   class LlamaLikeGraph : public Graph {
       using Graph::Graph;
   
   public:
       void set_weights_alias() override;
       void construct_llm() override;
       void load_param(
               std::shared_ptr<InputFile> fin, LlmParams& param,
               std::shared_ptr<Vocab> vocab) override;
   };
   

   定义了 LlamaLikeGraph 类，继承自 Graph 类，并重写了三个关键方法：

   • set_weights_alias：设置权重别名映射
   • construct_llm：构建模型计算图
   • load_param：加载模型参数

2. llama_like.cpp 文件分析

2.1 头文件引入

#include "llama_like.h"
#include "chatGLM.h"

引入了 llama_like.h 和 chatGLM.h 头文件，后者可能是为了使用其中定义的一些数据结构（如 Header）。

2.2 权重别名设置

void LlamaLikeGraph::set_weights_alias() {
    m_weights_name_aliases.clear();
    // clang-format off
    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"},
            {"model.layers.x.self_attn.o_proj.weight", "layers.x.attention.wo.weight"},
            {"model.layers.x.post_attention_layernorm.weight", "layers.x.ffn.norm.weight"},
            {"model.layers.x.mlp.up_proj.weight", "layers.x.ffn.w3.weight"},
            {"model.layers.x.mlp.down_proj.weight", "layers.x.ffn.w2.weight"},
            {"model.layers.x.mlp.gate_proj.weight", "layers.x.ffn.w1.weight"},
            {"model.norm.weight", "head.norm.weight"},
            {"lm_head.weight", "head.output.weight"},
    };
    // clang-format on
}

这个方法建立了原始权重名称和内部使用的权重名称之间的映射关系。与 GGML Llama 和 ChatGLM 不同，LlamaLike 模型使用分离的 Q、K、V 权重，而不是融合的 QKV 权重。这反映了 Llama2 和 Baichuan 等模型的架构特点。

2.3 参数加载

void LlamaLikeGraph::load_param(
        std::shared_ptr<InputFile> fin, LlmParams& param,
        std::shared_ptr<Vocab> vocab) {
    Header header;
    // load model header
    fin->read_raw((char*)&header.param_offset, sizeof(header.param_offset));
    fin->read_raw((char*)&header.param_length, sizeof(header.param_length));
    fin->read_raw((char*)&header.vocab_offset, sizeof(header.vocab_offset));
    fin->read_raw((char*)&header.vocab_length, sizeof(header.vocab_length));
    fin->read_raw((char*)&header.tensor_offset, sizeof(header.tensor_offset));

    fin->seek(header.param_offset);
    // load param
    fin->read_raw((char*)&param.n_embd, sizeof(param.n_embd));
    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_mult, sizeof(param.n_mult));
    fin->read_raw((char*)&param.n_vocab, sizeof(param.n_vocab));
    m_param = param;

    // load vocab
    fin->seek(header.vocab_offset);
    vocab->load_vocab(fin, param.n_vocab);

    INFER_LOG("total vocab length = %d\n", param.n_vocab);
    fin->seek(header.tensor_offset);
}

这个方法从模型文件中加载参数和词汇表，包括：

1. 读取模型头部信息：

   • 读取参数、词汇表和张量的偏移量和长度

2. 加载模型参数：

   • 读取嵌入维度、头数、层数、前馈网络倍数和词汇表大小

3. 加载词汇表：

   • 根据词汇表偏移量读取词汇表

4. 输出日志：

   • 输出词汇表大小的日志信息

5. 设置文件指针：

   • 将文件指针移动到张量数据的起始位置

2.4 计算图构建

void LlamaLikeGraph::construct_llm() {
    uint32_t embd = m_param.n_embd;
    uint32_t ffn_size = m_param.n_mult;
    uint32_t head = m_param.n_head;
    uint32_t ctx = m_param.n_ctx;
    uint32_t n_vocab = m_param.n_vocab;
    uint32_t rot = embd / head;

    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, /*mul*/ true, /*bias*/ false,
                                                  /*rms*/ true, /*eps*/ 1e-6);
        //! attentin
        auto attention_output = add_module<AttentionModule<LlamaAttention>>(
                this, norm_out_attention, embd, head, rot, ctx, model_config(),
                device(), name + ".attention", i, false, false, RotMode::ModelRotHalf);
        //! 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, /*mul*/ true, /*bias*/ false,
                                /*rms*/ true, /*eps*/ 1e-6);
        //! feed forward
        auto ffn_output = add_module<LlamaFFNModule>(
                this, ffn_norm_out, embd, ffn_size, 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");
}

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

1. 初始化参数：

   • 从模型参数中获取嵌入维度、前馈网络大小、头数等参数
   • 计算旋转位置编码的维度 rot = embd / head

2. 创建输入张量和嵌入层：

   • 创建输入张量
   • 添加嵌入层，将 token ID 转换为嵌入向量

3. 循环添加 Transformer 层：
   每个 Transformer 层包括：

   • 注意力层前的 RMSNorm（rms = true，bias = false，eps = 1e-6）
   • 注意力层 (LlamaAttention)，使用 ModelRotHalf 旋转模式
   • 残差连接
   • 前馈网络前的 RMSNorm
   • 前馈网络 (LlamaFFNModule)
   • 残差连接

4. 添加输出层：

   • 添加 HeadModule，将隐藏状态转换为词汇表概率分布

3. LlamaLike 模型的特点

通过分析 llama_like.h 和 llama_like.cpp 文件，可以看出 LlamaLike 模型的以下特点：

3.1 模型架构

1. Pre-LayerNorm 架构：

   • 在注意力层和前馈网络前进行层归一化
   • 使用 RMSNorm（rms = true，没有偏置项 bias = false）
   • 使用较小的 epsilon 值 (eps = 1e-6)

2. 分离的 Q、K、V 权重：

   • 使用分离的 Q、K、V 权重，而不是融合的 QKV 权重
   • 权重别名映射反映了这一点：

     {"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"},
     

3. 旋转位置编码：

   • 使用旋转位置编码 (RoPE)，通过 rot = embd / head 计算旋转维度
   • 使用 ModelRotHalf 旋转模式：RotMode::ModelRotHalf

4. SwiGLU 激活函数：

   • 在前馈网络中使用 SwiGLU 激活函数，通过 LlamaFFNModule 实现
   • 使用三个权重矩阵：w1、w2 和 w3

5. 无偏置项：

   • 层归一化不使用偏置项 (bias = false)
   • 注意力层不使用偏置项 (false 参数)

3.2 与 GGML Llama 的比较

LlamaLike 模型与 GGML Llama 模型有以下区别：

1. 权重结构：

   • LlamaLike 使用分离的 Q、K、V 权重
   • GGML Llama 可能使用融合的 QKV 权重

2. 权重命名：

   • LlamaLike 使用 model.embed_tokens.weight 等命名
   • GGML Llama 使用不同的命名约定

3. 旋转模式：

   • LlamaLike 明确指定了 ModelRotHalf 旋转模式
   • GGML Llama 可能使用不同的旋转模式

4. 模型格式：

   • LlamaLike 使用自定义格式
   • GGML Llama 使用 GGML 格式

3.3 与 ChatGLM 系列的比较

LlamaLike 模型与 ChatGLM 系列模型有以下区别：

1. 层归一化：

   • LlamaLike 使用 RMSNorm（rms = true，bias = false）
   • ChatGLM 使用传统的层归一化（rms = false，bias = true）

2. 注意力机制：

   • LlamaLike 使用 LlamaAttention
   • ChatGLM 使用 GlmAttention 或 Glm2MultiQueryAttention

3. 前馈网络：

   • LlamaLike 使用 LlamaFFNModule
   • ChatGLM 使用 GlmFFNModule

4. 残差连接：

   • LlamaLike 使用普通的残差连接
   • ChatGLM 使用带缩放因子的残差连接

5. 特殊 token 处理：

   • LlamaLike 模型没有实现 post_tokenize 方法，可能不需要特殊的 token 处理
   • ChatGLM 系列模型实现了 post_tokenize 方法，在输入序列的开头和结尾添加特殊 token

6. 词汇表大小：

   • LlamaLike 模型的词汇表大小由模型文件决定
   • ChatGLM 系列模型使用固定的词汇表大小（ChatGLM 为 130528，ChatGLM2/3 为 65024）

4. 模型加载流程

LlamaLike 模型的加载流程包括以下步骤：

1. 读取模型头部信息：

   fin->read_raw((char*)&header.param_offset, sizeof(header.param_offset));
   fin->read_raw((char*)&header.param_length, sizeof(header.param_length));
   fin->read_raw((char*)&header.vocab_offset, sizeof(header.vocab_offset));
   fin->read_raw((char*)&header.vocab_length, sizeof(header.vocab_length));
   fin->read_raw((char*)&header.tensor_offset, sizeof(header.tensor_offset));
   

2. 加载模型参数：

   fin->seek(header.param_offset);
   fin->read_raw((char*)&param.n_embd, sizeof(param.n_embd));
   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_mult, sizeof(param.n_mult));
   fin->read_raw((char*)&param.n_vocab, sizeof(param.n_vocab));
   

3. 加载词汇表：

   fin->seek(header.vocab_offset);
   vocab->load_vocab(fin, param.n_vocab);
   

4. 构建计算图：
   在 Graph::load 方法中调用 construct_llm 构建计算图

5. 收集权重：
   在 Graph::load 方法中调用 collect_weights 收集权重

6. 设置权重别名：
   在 Graph::load 方法中调用 set_weights_alias 设置权重别名

7. 加载权重：
   在 Graph::load 方法中循环读取权重数据，并将其与计算图中的权重关联

5. 计算图执行流程

LlamaLike 模型的计算图执行流程与其他模型类似，包括以下步骤：

1. 准备输入：

   • 将输入文本转换为 token ID
   • 如果需要，进行特殊 token 处理

2. 执行计算图：

   • 从输入到输出执行计算图
   • 每个 Transformer 层按顺序执行

3. 采样：

   • 根据输出概率分布采样下一个 token

4. 生成文本：

   • 将采样得到的 token ID 转换为文本
   • 重复执行计算图和采样，直到生成结束

6. 优化策略

LlamaLike 模型实现了一些优化策略：

1. RMSNorm：

   • 使用 RMSNorm 代替传统的层归一化，减少计算量
   • 不使用偏置项，减少参数量

2. 旋转位置编码：

   • 使用旋转位置编码 (RoPE)，避免了位置嵌入的额外参数
   • 使用 ModelRotHalf 旋转模式，可能是为了优化性能

3. SwiGLU 激活函数：

   • 在前馈网络中使用 SwiGLU 激活函数，提高模型性能

4. 分离的 Q、K、V 权重：

   • 使用分离的 Q、K、V 权重，可能有助于并行计算

7. 扩展性分析

LlamaLike 模型的设计具有良好的扩展性：

1. 支持多种模型：

   • 通过相同的接口支持 Baichuan 和 Llama2 等多种模型
   • 只需修改权重别名映射，就可以支持新的模型变体

2. 模块化设计：

   • 使用模块化设计构建计算图，便于扩展和维护
   • 可以方便地替换或修改特定模块

3. 参数化配置：

   • 模型参数（如嵌入维度、头数、层数等）由模型文件决定
   • 不需要硬编码特定模型的参数

总结

llama_like.h 和 llama_like.cpp 文件实现了 InferLLM 框架中对类 Llama 架构模型（如 Baichuan 和 Llama2）的支持。它们定义了 LlamaLikeGraph 类，该类继承自基础的 Graph 类，并实现了特定于类 Llama 架构模型的加载和计算图构建逻辑。

通过分析这些文件，我们可以了解类 Llama 架构模型的结构特点和实现细节，以及 InferLLM 框架如何支持不同架构的模型。这种设计使得 InferLLM 可以灵活地支持多种模型架构，而不需要修改框架的核心代码。

LlamaLike 模型的实现展示了现代大型语言模型的一些关键技术，如 RMSNorm、旋转位置编码和 SwiGLU 激活函数等。这些技术的应用使得模型在保持高性能的同时，减少了参数量和计算量，提高了推理效率。