ChatGLM 系列模型代码结构与功能实现分析

InferLLM 框架中的 ChatGLM 系列模型实现包括 ChatGLM、ChatGLM2 和 ChatGLM3 三个版本，它们都继承自基础的 Graph 类，并实现了特定的模型结构和加载逻辑。下面对这些文件的代码结构和功能实现进行详细分析。

1. chatGLM.h 文件分析

chatGLM.h 文件定义了 ChatGLM 系列模型的基础结构和接口：

1.1 基础数据结构

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;
};

这个结构体用于存储模型的参数信息，包括隐藏层大小、注意力头数、层数等。

1.2 ChatGLM 系列类定义

class ChatGLMGraph : 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;
    void post_tokenize(std::vector<Vocab::Id>& input) override;
};

ChatGLMGraph 类继承自 Graph 类，并重写了四个关键方法：

• set_weights_alias：设置权重别名映射
• construct_llm：构建模型计算图
• load_param：加载模型参数
• post_tokenize：对输入 token 进行后处理

ChatGLMGraph2 和 ChatGLMGraph3 类的定义与 ChatGLMGraph 类似，也重写了相同的四个方法。

2. chatGLM.cpp 文件分析

chatGLM.cpp 文件实现了 ChatGLM 模型的具体功能：

2.1 权重别名设置

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
}

这个方法建立了原始权重名称和内部使用的权重名称之间的映射关系，使得模型可以正确加载预训练权重。

2.2 参数加载

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);
}

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

1. 读取模型头部信息
2. 根据偏移量读取模型参数
3. 加载词汇表
4. 设置固定的词汇表大小（130528）
5. 将文件指针移动到张量数据的起始位置

2.3 计算图构建

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++) {
        // ... 构建 Transformer 层
    }
    //! 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)

2.4 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：

1. 在开头添加 token ID 5（开始标记）
2. 在结尾添加 token ID 130001 和 130004（结束标记）

3. chatGLM2.cpp 文件分析

chatGLM2.cpp 文件实现了 ChatGLM2 模型的具体功能：

3.1 权重别名设置

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

ChatGLM2 的权重别名与 ChatGLM 不同，反映了模型架构的变化。

3.2 参数加载

void ChatGLMGraph2::load_param(
        std::shared_ptr<InputFile> fin, LlmParams& param,
        std::shared_ptr<Vocab> vocab) {
    // ... 与 ChatGLM 类似的头部和参数加载
    
    // 加载多查询注意力相关参数
    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));
    
    // ... 加载词汇表
    
    // 设置词汇表大小
    m_param.n_vocab = 65024;
    param.n_vocab = 65024;
}

ChatGLM2 的参数加载与 ChatGLM 类似，但增加了多查询注意力相关参数的加载，并使用不同的词汇表大小（65024）。

3.3 计算图构建

void ChatGLMGraph2::construct_llm() {
    // ... 创建输入张量和嵌入层
    
    int nr_layer = m_param.n_layer;
    for (int i = 0; i < nr_layer; i++) {
        // ... 构建 Transformer 层，使用 Glm2MultiQueryAttention
    }
    
    // ... 添加输出层
}

ChatGLM2 的计算图构建与 ChatGLM 类似，但有以下区别：

1. 使用 Glm2MultiQueryAttention 代替 GlmAttention
2. 使用 RMSNorm 代替传统的层归一化（rms 参数为 true）
3. 不使用偏置项（bias 参数为 false）
4. 残差连接不使用缩放因子

3.4 Token 后处理

void ChatGLMGraph2::post_tokenize(std::vector<Vocab::Id>& input) {
    std::vector<Vocab::Id> res;
    res.push_back(64790);
    res.push_back(64792);
    // add a space in the head
    input.insert(input.begin(), res.begin(), res.end());
}

ChatGLM2 的 token 后处理与 ChatGLM 不同，在开头添加了两个特殊 token（64790 和 64792）。

4. chatGLM3.cpp 文件分析

chatGLM3.cpp 文件实现了 ChatGLM3 模型的具体功能，其实现与 ChatGLM2 非常相似：

4.1 权重别名设置

ChatGLM3 的权重别名与 ChatGLM2 相同，表明它们的模型架构相似。

4.2 参数加载

ChatGLM3 的参数加载与 ChatGLM2 相同，包括多查询注意力相关参数的加载和相同的词汇表大小（65024）。

4.3 计算图构建

ChatGLM3 的计算图构建与 ChatGLM2 相同，使用相同的注意力机制和层归一化方式。

4.4 Token 后处理

ChatGLM3 的 token 后处理与 ChatGLM2 相同，在开头添加相同的特殊 token。

5. ChatGLM 系列模型的演进

通过分析这三个版本的代码，可以看出 ChatGLM 系列模型的演进：

5.1 ChatGLM 到 ChatGLM2 的变化

1. 架构变化：

   • 词嵌入层路径从 transformer.word_embeddings 变为 transformer.embedding.word_embeddings
   • Transformer 层路径从 transformer.layers 变为 transformer.encoder.layers

2. 注意力机制：

   • 从普通注意力 (GlmAttention) 变为多查询注意力 (Glm2MultiQueryAttention)

3. 归一化方式：

   • 从传统的层归一化变为 RMSNorm

4. 残差连接：

   • 去掉了残差连接中的缩放因子

5. 词汇表大小：

   • 从 130528 减少到 65024

6. 特殊 token：

   • 使用不同的特殊 token

5.2 ChatGLM2 到 ChatGLM3 的变化

ChatGLM3 与 ChatGLM2 的代码几乎相同，表明它们的基础架构非常相似。可能的区别在于预训练权重和训练方法，而不是模型架构本身。

6. 总结

ChatGLM 系列模型在 InferLLM 框架中的实现展示了大型语言模型的基本架构和演进过程。通过分析这些代码，我们可以看到：

1. 模块化设计：

   • 使用模块化设计构建计算图，便于扩展和维护

2. 权重映射：

   • 使用权重别名机制，适应不同的权重命名约定

3. 架构演进：

   • 从 ChatGLM 到 ChatGLM2/3，模型架构有明显改进
   • 引入多查询注意力机制，提高效率
   • 使用 RMSNorm 代替传统层归一化，提高稳定性

4. 特殊 token 处理：

   • 不同版本使用不同的特殊 token，反映了模型训练的变化

这些实现为我们理解大型语言模型的架构和推理过程提供了宝贵的参考。