ModelImp 类代码结构与功能实现分析

ModelImp 类是 InferLLM 框架中 Model 类的具体实现，负责模型的加载、初始化和推理过程。下面对 model_imp.h 和 model_imp.cpp 的代码结构和功能实现进行详细分析。

1. 辅助函数

DType dtype_from_str(const std::string& str) {
    if (str == "float32" || str == "fp23") {
        return DType::Float32;
    } else if (str == "float16" || str == "fp16") {
        return DType::Float16;
    } else if (str == "int8") {
        return DType::Int8;
    } else if (str == "uint8") {
        return DType::Uint8;
    } else if (str == "int4") {
        return DType::Int4;
    } else if (str == "uint4") {
        return DType::Uint4;
    } else {
        INFER_ASSERT(0, "Unsupported dytpe.");
    }
}

这个辅助函数将字符串表示的数据类型转换为枚举类型 DType，支持多种精度的数据类型，包括 float32、float16、int8、uint8、int4 和 uint4。

2. ModelImp 类

2.1 构造函数

ModelImp::ModelImp(const ModelConfig& config, const std::string& name)
        : m_name(name), m_config(config) {
    uint32_t nr_thread = config.nr_thread;
    std::string device_type = config.device_type;
    if (device_type == "CPU" || device_type == "cpu") {
#if INFER_X86
        m_device = make_unique<CPUDevice>(KernelType::X86, nr_thread);
#elif INFER_ARM
        m_device = make_unique<CPUDevice>(KernelType::Arm, nr_thread);
#else
        m_device = make_unique<CPUDevice>(KernelType::Naive, nr_thread);
#endif
    } else if (
            device_type == "GPU" || device_type == "CUDA" || device_type == "gpu") {
        // if compile with GPU, use GPU, else use CPUDevice
#if ENABLE_GPU
        m_device = make_unique<GPUDevice>(0);
#else
        INFER_ASSERT(0, "GPU is disabled when build, please build with GPU.");
#endif
    }

    UserConfig user_config;
    user_config.compt_type = dtype_from_str(config.compt_type);
    m_graph = Graph::make_graph(user_config, m_device.get(), name);
    m_past = 0;
}

构造函数根据配置创建相应的设备和计算图：

1. 根据 device_type 创建 CPU 或 GPU 设备
   • 对于 CPU，根据编译选项选择 X86、Arm 或 Naive 内核
   • 对于 GPU，检查是否启用了 GPU 支持
2. 创建用户配置，设置计算类型
3. 创建计算图
4. 初始化 m_past 为 0，表示已处理的 token 数量

2.2 load 方法

void ModelImp::load(const std::string& model_path) {
    m_vocab = std::make_shared<Vocab>();
    std::shared_ptr<InputFile> fin =
            std::make_shared<InputFile>(model_path, m_config.enable_mmap);

    m_param.n_ctx = m_config.nr_ctx;
    m_graph->load(fin, m_param, m_vocab);
    m_logist.resize(m_param.n_vocab);
}

load 方法从指定路径加载模型：

1. 创建词汇表对象
2. 创建输入文件对象，支持内存映射
3. 设置上下文长度
4. 调用计算图的 load 方法加载模型
5. 调整 logits 向量的大小为词汇表大小

2.3 init 方法

void init(
        uint32_t top_k, float top_p, float temp, float repeat_penalty,
        int repeat_last_n, int32_t seed, int32_t end_token) {
    m_top_k = top_k;
    m_top_p = top_p;
    m_temp = temp;
    m_repeat_penalty = repeat_penalty;
    m_repeat_last_n = repeat_last_n;
    m_end_token = end_token;
    for (uint32_t i = 0; i < m_repeat_last_n; i++) {
        m_last_queue.push_back(0);
    }
    m_rng = std::mt19937(seed);
}

init 方法初始化模型的生成参数：

1. 设置采样参数：top_k、top_p、温度、重复惩罚等
2. 设置结束 token
3. 初始化最近生成的 token 队列
4. 使用指定的种子初始化随机数生成器

2.4 prefill 方法

void ModelImp::prefill(const std::string& promote) {
    auto tokens = tokenize(promote, true);
    m_graph->post_tokenize(tokens);
    for (auto token : tokens) {
        m_last_queue.push_back(token);
        m_last_queue.pop_front();
    }
    m_graph->execute(tokens, m_logist, m_past, true);
    m_past = tokens.size();
}

prefill 方法处理初始提示文本：

1. 将文本分词为 token 序列
2. 调用计算图的 post_tokenize 方法进行后处理
3. 更新最近生成的 token 队列
4. 执行计算图，进行预填充（prefill=true）
5. 更新已处理的 token 数量

2.5 decode 方法

std::string ModelImp::decode(const std::string& user_input, int& token) {
    auto tokens = tokenize(user_input, false);
    m_graph->post_tokenize(tokens);
    for (auto token : tokens) {
        m_last_queue.push_back(token);
        m_last_queue.pop_front();
    }
    m_graph->execute(tokens, m_logist, m_past, false);
    sample_and_update();
    m_past += tokens.size();
    token = m_pre_token;
    return m_vocab->id_to_token[m_pre_token].tok;
}

decode 方法处理用户输入并生成下一个 token：

1. 将用户输入分词为 token 序列
2. 调用计算图的 post_tokenize 方法进行后处理
3. 更新最近生成的 token 队列
4. 执行计算图，生成 logits
5. 采样下一个 token 并更新状态
6. 更新已处理的 token 数量
7. 返回生成的 token 对应的文本

2.6 decode_iter 方法

std::string ModelImp::decode_iter(int& token) {
    auto start = m_timer.get_time();
    m_graph->execute({m_pre_token}, m_logist, m_past);
    auto end = m_timer.get_time();
    m_time_cost += end - start;
    sample_and_update();
    m_past++;
    token = m_pre_token;
    return m_vocab->id_to_token[m_pre_token].tok;
}

decode_iter 方法迭代生成下一个 token：

1. 记录开始时间
2. 执行计算图，使用上一个生成的 token 作为输入
3. 记录结束时间并更新总计算时间
4. 采样下一个 token 并更新状态
5. 更新已处理的 token 数量
6. 返回生成的 token 对应的文本

2.7 sample_and_update 方法

int32_t ModelImp::sample_and_update() {
    // sample the next token
   auto token = llama_sample_top_p_top_k(
           *m_vocab, m_logist.data(), m_last_queue, m_repeat_penalty, m_top_k, m_top_p,
           m_temp, m_rng);
    // update the last queue
    m_last_queue.push_back(token);
    m_last_queue.pop_front();
    m_pre_token = token;
    if (token == m_end_token) {
        m_device->deactive();
    }
    return token;
}

sample_and_update 方法采样下一个 token 并更新状态：

1. 使用 top-p 和 top-k 采样方法从 logits 中采样下一个 token
2. 更新最近生成的 token 队列
3. 更新上一个生成的 token
4. 如果生成的 token 是结束 token，则停用设备
5. 返回生成的 token

2.8 tokenize 方法

std::vector<Vocab::Id> ModelImp::tokenize(const std::string& text, bool bos) {
    std::vector<Vocab::Id> res;
    std::vector<int> score;
    std::vector<Vocab::Id> prev;
    int len = text.length();

    score.resize(len + 1);
    prev.resize(len + 1);

    // Forward pass
    for (int i = 0; i < len; i++) {
        int max_len = std::min(len - i, MAX_TOKEN_LEN);
        for (int sub_len = 1; sub_len <= len - i; sub_len++) {
            auto sub = text.substr(i, sub_len);
            auto token = m_vocab->token_to_id.find(sub);
            if (token != m_vocab->token_to_id.end()) {
                int token_score = sub.length() * sub.length();
                int local_score = score[i] + token_score;
                int next = i + sub_len;
                if (score[next] < local_score) {
                    score[next] = local_score;
                    prev[next] = (*token).second;
                }
            }
        }
    }

    // Backward pass
    int i = len;
    while (i > 0) {
        Vocab::Id token_id = prev[i];
        if (token_id == 0) {
            // TODO: Return error or something more meaningful
            printf("failed to tokenize string!\n");
            break;
        }
        res.push_back(token_id);
        auto token = m_vocab->id_to_token[token_id].tok;
        i -= token.length();
    }

    if (bos) {
        res.push_back(1);  // TODO: replace with vocab.bos
    }

    // Pieces are in reverse order so correct that
    std::reverse(res.begin(), res.end());

    return res;
}

tokenize 方法将文本分词为 token 序列：

1. 使用动态规划算法找到最优的分词方式
   • 前向传递：计算每个位置的最优分数和对应的 token
   • 后向传递：从末尾开始，根据前向传递的结果构建 token 序列
2. 如果需要，添加 BOS（Beginning of Sequence）token
3. 反转 token 序列，使其按正确的顺序排列
4. 返回 token 序列

这个分词算法使用了贪婪的方法，尽量选择长度更长的 token，以减少序列长度。

2.9 decode_summary 方法

std::string ModelImp::decode_summary() const {
    std::string ret = "Run Model Summary:\n";
    ret += "Total Model Compute Time:   " + std::to_string(m_time_cost) + "s\n";
    ret += "Total Model Compute Token:  " + std::to_string(m_past) + "\n";
    ret += "Average Token Compute Time: " +
           std::to_string(m_time_cost * 1000 / m_past) + "ms\n";
    ret += "Average Token Generation Speed: " +
           std::to_string(m_past / m_time_cost) + "token/s\n";
    return ret;
}

decode_summary 方法返回模型运行的摘要信息：

1. 总计算时间
2. 总计算 token 数量
3. 平均 token 计算时间
4. 平均 token 生成速度

3. 工作流程

ModelImp 类的典型工作流程如下：

1. 创建和初始化：

   ModelImp model_imp(config, "llama");
   model_imp.load("path/to/model.bin");
   model_imp.init(40, 0.9f, 0.8f, 1.1f, 64, 42, -1);
   

2. 预填充：

   model_imp.prefill("你好，请介绍一下自己。");
   

3. 迭代生成：

   int token;
   std::string result;
   while (true) {
       std::string next = model_imp.decode_iter(token);
       if (token == end_token) break;
       result += next;
   }
   

4. 获取摘要：

   std::string summary = model_imp.decode_summary();
   

4. 关键数据结构

ModelImp 类包含以下关键数据成员：

• 设备和计算图：

  • m_device：计算设备（CPU 或 GPU）
  • m_graph：计算图，负责模型的结构和执行

• 模型参数：

  • m_param：LLM 模型参数
  • m_config：用户配置
  • m_vocab：词汇表

• 生成参数：

  • m_top_k、m_top_p、m_temp：采样参数
  • m_repeat_penalty、m_repeat_last_n：重复惩罚参数
  • m_end_token：结束 token

• 状态变量：

  • m_past：已处理的 token 数量
  • m_pre_token：上一个生成的 token
  • m_last_queue：最近生成的 token 队列，用于重复惩罚
  • m_logist：模型输出的 logits，表示下一个 token 的概率分布

• 性能统计：

  • m_timer：计时器，用于测量计算时间
  • m_time_cost：总计算时间

5. 采样算法

ModelImp 类使用 llama_sample_top_p_top_k 函数进行采样，这是一个结合了 top-k 和 top-p（nucleus sampling）的采样方法：

int32_t llama_sample_top_p_top_k(
        const Vocab& vocab, float* logits, const std::list<int32_t>& last_n_tokens,
        float repeat_penalty, int top_k, float top_p, float temp, std::mt19937& rng) {
    // 应用重复惩罚
    for (auto token_id : last_n_tokens) {
        logits[token_id] /= repeat_penalty;
    }

    // 应用温度
    for (int i = 0; i < vocab.id_to_token.size(); i++) {
        logits[i] /= temp;
    }

    // 计算 softmax
    float max_logit = -INFINITY;
    for (int i = 0; i < vocab.id_to_token.size(); i++) {
        max_logit = std::max(max_logit, logits[i]);
    }
    
    float sum = 0.0f;
    for (int i = 0; i < vocab.id_to_token.size(); i++) {
        logits[i] = expf(logits[i] - max_logit);
        sum += logits[i];
    }
    
    for (int i = 0; i < vocab.id_to_token.size(); i++) {
        logits[i] /= sum;
    }

    // 找出 top-k 的 token
    std::vector<std::pair<float, int>> candidates;
    for (int i = 0; i < vocab.id_to_token.size(); i++) {
        candidates.push_back({logits[i], i});
    }
    
    std::partial_sort(
            candidates.begin(), candidates.begin() + top_k, candidates.end(),
            [](const std::pair<float, int>& a, const std::pair<float, int>& b) {
                return a.first > b.first;
            });
    
    candidates.resize(top_k);

    // 应用 top-p（nucleus sampling）
    float cumsum = 0.0f;
    for (int i = 0; i < candidates.size(); i++) {
        cumsum += candidates[i].first;
        if (cumsum >= top_p) {
            candidates.resize(i + 1);
            break;
        }
    }

    // 按概率采样
    std::vector<float> probs;
    for (auto& candidate : candidates) {
        probs.push_back(candidate.first);
    }
    
    std::discrete_distribution<> dist(probs.begin(), probs.end());
    return candidates[dist(rng)].second;
}

这个采样算法的步骤如下：

1. 应用重复惩罚，降低最近生成的 token 的概率
2. 应用温度参数，控制分布的平滑程度
3. 计算 softmax，将 logits 转换为概率
4. 选择概率最高的 top-k 个 token
5. 应用 top-p 采样，只保留累积概率达到 p 的 token
6. 按概率采样最终的 token

6. 分词算法

ModelImp 类使用动态规划算法进行分词，尽量选择长度更长的 token：

std::vector<Vocab::Id> ModelImp::tokenize(const std::string& text, bool bos) {
    std::vector<Vocab::Id> res;
    std::vector<int> score;
    std::vector<Vocab::Id> prev;
    int len = text.length();

    score.resize(len + 1);
    prev.resize(len + 1);

    // 前向传递：计算每个位置的最优分数和对应的 token
    for (int i = 0; i < len; i++) {
        int max_len = std::min(len - i, MAX_TOKEN_LEN);
        for (int sub_len = 1; sub_len <= len - i; sub_len++) {
            auto sub = text.substr(i, sub_len);
            auto token = m_vocab->token_to_id.find(sub);
            if (token != m_vocab->token_to_id.end()) {
                int token_score = sub.length() * sub.length();
                int local_score = score[i] + token_score;
                int next = i + sub_len;
                if (score[next] < local_score) {
                    score[next] = local_score;
                    prev[next] = (*token).second;
                }
            }
        }
    }

    // 后向传递：从末尾开始，根据前向传递的结果构建 token 序列
    int i = len;
    while (i > 0) {
        Vocab::Id token_id = prev[i];
        if (token_id == 0) {
            printf("failed to tokenize string!\n");
            break;
        }
        res.push_back(token_id);
        auto token = m_vocab->id_to_token[token_id].tok;
        i -= token.length();
    }

    if (bos) {
        res.push_back(1);  // BOS token
    }

    // 反转 token 序列，使其按正确的顺序排列
    std::reverse(res.begin(), res.end());

    return res;
}

这个分词算法的特点是：

1. 使用动态规划找到最优的分词方式
2. 优先选择长度更长的 token，以减少序列长度
3. 支持添加 BOS（Beginning of Sequence）token

7. 优化策略

ModelImp 类实现了几种优化策略：

1. 预填充模式：

   m_graph->execute(tokens, m_logist, m_past, true);
   

   预填充模式可以一次性处理多个 token，提高效率。

2. 延迟加载：

   m_graph->load(fin, m_param, m_vocab);
   

   模型权重采用延迟加载策略，只在需要时才从文件读取。

3. 性能统计：

   auto start = m_timer.get_time();
   // ...
   auto end = m_timer.get_time();
   m_time_cost += end - start;
   

   记录计算时间，用于性能分析和优化。

4. 设备管理：

   if (token == m_end_token) {
       m_device->deactive();
   }
   

   在生成结束时停用设备，释放资源。

8. 与其他组件的交互

ModelImp 类与框架的其他组件有密切的交互：

1. 与 Device 的交互：

   m_device = make_unique<CPUDevice>(KernelType::X86, nr_thread);
   

   创建并管理计算设备，负责内存分配和计算执行。

2. 与 Graph 的交互：

   m_graph = Graph::make_graph(user_config, m_device.get(), name);
   m_graph->load(fin, m_param, m_vocab);
   m_graph->execute(tokens, m_logist, m_past, prefill);
   

   创建并管理计算图，负责模型的结构和执行。

3. 与 Vocab 的交互：

   m_vocab = std::make_shared<Vocab>();
   auto tokens = tokenize(promote, true);
   

   创建并管理词汇表，负责文本和 token 的转换。

总结

ModelImp 类是 InferLLM 框架的核心实现类，它负责模型的加载、初始化和推理过程。它通过与 Device、Graph 和 Vocab 等组件的交互，实现了高效的大型语言模型推理。它支持多种优化策略，如预填充模式、延迟加载和性能统计，以提高推理效率。它还实现了先进的采样算法和分词算法，以提高生成文本的质量。

通过分析 model_imp.h 和 model_imp.cpp，我们可以看到 InferLLM 框架的设计思想和实现细节，这对于理解和使用框架非常有帮助。