使用MNN部署语言模型

155 阅读10分钟

使用MNN部署语言模型

1. 引言

英伟达的论文指出,在"代理型 AI"(Agentic AI)场景中, Small Language Models(SLMs) 足够强大、运算合适且更加经济,因此比大型语言模型(LLMs)更适合作为未来趋势; 当需要通用对话能力时,推荐 异构系统(结合 SLM 与 LLM 的模块化系统),小型语言模型在实际应用中的巨大潜力。

越来越多的实际应用场景需要在终端设备上部署语言模型:

终端AI助手

  • 低延迟响应:本地推理避免网络延迟
  • 离线运行:无需网络连接即可工作
  • 隐私保护:敏感数据不离开本地设备

边缘计算

  • IoT设备部署:在资源受限的嵌入式设备上运行
  • 实时推理:工业控制、自动驾驶等需要实时响应的场景
  • 资源受限环境:电力、带宽、计算资源有限的环境

移动应用

  • Android/iOS原生应用:集成到移动App中
  • 本地知识问答:无需联网的智能问答系统
  • 实时对话系统:语音助手、客服机器人等

面对这些应用需求,阿里开源的MNN推理框架为我们提供了一个出色的解决方案。

本文将基于 mnn-llm 这个使用MNN框架部署大语言模型的实际案例,分析如何使用MNN框架实现大语言模型的终端部署。

2. 核心业务流程梳理

使用MNN部署大语言模型涉及三个主要流程:模型加载流程推理生成流程CLI应用流程

2.1 模型加载流程

模型加载是整个系统的初始化过程:

// 核心流程:配置解析 → MNN模型加载 → 运行时初始化
void Llm::load() {
    // 1. 获取配置文件中的模型路径
    auto model_path = config_->llm_model();
    
    // 2. 使用MNN加载模型文件
    modules_.emplace_back(Module::load({}, {}, model_path.c_str()));
    
    // 3. 初始化推理运行时环境
    init_runtime();
}

关键步骤

  • 从JSON配置文件读取模型文件路径
  • 调用MNN的Module::load()加载.mnn模型文件
  • 配置推理后端(CPU/GPU/NPU)和线程数

2.2 推理生成流程

这是核心的文本生成流程,从用户输入到模型输出:

// 核心流程:文本预处理 → MNN推理 → 结果后处理
void response(const std::string& user_content) {
    // 1. 应用提示词模板
    auto prompt = apply_prompt_template(user_content);
    
    // 2. 文本分词转换为token序列
    auto input_ids = tokenizer_encode(prompt);
    
    // 3. 自回归生成token序列
    auto output_ids = generate(input_ids);
    
    // 4. 将token解码为文本输出
    for(int id : output_ids) {
        std::cout << tokenizer_decode(id);
    }
}

关键步骤

  • 输入文本按提示词模板格式化
  • Tokenizer将文本转换为模型可理解的数字序列
  • 模型自回归生成新的token
  • 将生成的token解码回文本

2.3 CLI应用流程

CLI Demo展示了完整的应用流程:

// 核心流程:参数解析 → 模型初始化 → 业务执行
int main(int argc, const char* argv[]) {
    // 1. 解析命令行参数
    std::string model_dir = argv[1];
    
    // 2. 初始化LLM实例
    std::unique_ptr<Llm> llm(Llm::createLLM(model_dir));
    llm->load();
    
    // 3. 选择执行模式
    if (argc < 3) {
        llm->chat();  // 交互式对话
    } else {
        benchmark(llm.get(), argv[2]);  // 批量测试
    }
}

关键步骤

  • 解析模型目录路径参数
  • 创建并加载LLM实例
  • 根据参数选择交互模式或测试模式

3. 核心技术实现细节

本章深入分析MNN-LLM项目中的关键技术实现,重点关注文本处理、MNN接口使用和数据转换等核心环节。

3.1 文本分词技术实现

Tokenizer的工作机制

Tokenizer负责文本和token序列之间的双向转换:

// include/tokenizer.hpp
class Tokenizer {
public:
    // 文本编码:文本 → token序列
    std::vector<int> encode(const std::string& str) {
        // 1. 文本预处理(标准化、清洗)
        auto processed_text = preprocess(str);
        
        // 2. 应用分词算法(BPE/SentencePiece/WordPiece)
        auto tokens = tokenize(processed_text);
        
        // 3. 转换为数字ID
        std::vector<int> token_ids;
        for (const auto& token : tokens) {
            token_ids.push_back(vocab_[token]);
        }
        return token_ids;
    }
    
    // token解码:数字ID → 文本
    virtual std::string decode(int id) = 0;
    
    // 特殊token判断
    bool is_stop(int token) { return stop_words_.count(token) > 0; }
    bool is_special(int token) { return special_tokens_.count(token) > 0; }
};
多种Tokenizer实现

项目支持主流的分词器类型:

// SentencePiece分词器 (src/tokenizer.cpp:245-280)
class SentencePieceTokenizer : public Tokenizer {
    sentencepiece::SentencePieceProcessor sp_;
    
    std::vector<int> encode(const std::string& str) override {
        std::vector<int> ids;
        sp_.Encode(str, &ids);  // 使用SentencePiece库编码
        return ids;
    }
    
    std::string decode(int id) override {
        return sp_.IdToPiece(id);  // ID转换为piece
    }
};

// TikToken分词器 (用于GPT系列模型)
class TikTokenizer : public Tokenizer {
    tiktoken::Encoding enc_;
    
    std::vector<int> encode(const std::string& str) override {
        return enc_.encode(str);  // 使用TikToken编码
    }
    
    std::string decode(int id) override {
        return enc_.decode_single_token(id);
    }
};

3.2 词嵌入技术实现

嵌入向量的生成过程
// src/llm.cpp中的embedding实现
MNN::Express::VARP embedding(const std::vector<int>& input_ids) {
    // 1. 从token ID查找对应的嵌入向量
    if (disk_embedding_) {
        // 磁盘嵌入:按需从磁盘加载
        auto embedding_ptr = std::make_shared<float>(input_ids.size() * hidden_size_);
        disk_embedding_->disk_embedding_lookup(input_ids, embedding_ptr.get());
        
        // 2. 创建MNN Tensor
        auto embedding_tensor = MNN::Express::_Const(
            embedding_ptr.get(), 
            {(int)input_ids.size(), hidden_size_}, 
            MNN::Express::NHWC
        );
        return embedding_tensor;
    } else {
        // 内存嵌入:直接从权重矩阵查找
        return embedding_weight_[input_ids];  // 简化表示
    }
}
磁盘嵌入优化实现

为了节省内存,大型嵌入矩阵可以存储在磁盘上:

// src/llm.cpp:99-228 DiskEmbedding类
class DiskEmbedding {
private:
    std::unique_ptr<uint8_t[]> weight_;      // 权重数据缓冲区
    std::unique_ptr<uint8_t[]> alpha_;       // 量化参数缓冲区
    int hidden_size_;                        // 嵌入维度
    int quant_bit_;                          // 量化位数(4位或8位)
    
public:
    void disk_embedding_lookup(const std::vector<int>& input_ids, float* dst) {
        for (size_t i = 0; i < input_ids.size(); i++) {
            int token = input_ids[i];
            
            if (quant_bit_ > 0) {
                // 量化模式:从磁盘读取量化数据并反量化
                seek_read(weight_.get(), weight_token_size_, 
                         w_offset_ + token * weight_token_size_);
                
                // 按块反量化
                auto alpha_ptr = reinterpret_cast<float*>(alpha_.get()) 
                               + token * block_num_ * 2;
                for (int n = 0; n < block_num_; n++) {
                    float scale = alpha_ptr[n * 2 + 1];  // 缩放因子
                    float zero = alpha_ptr[n * 2];       // 零点
                    uint8_t* src = weight_.get() + n * (quant_block_ * quant_bit_ / 8);
                    float* dst_ptr = dst + i * hidden_size_ + n * quant_block_;
                    
                    // 4位或8位反量化
                    dequant_(src, dst_ptr, scale, zero, quant_block_);
                }
            } else {
                // bf16模式:直接读取bf16数据
                seek_read(weight_.get(), weight_token_size_, token * weight_token_size_);
                bf16_to_fp32(weight_.get(), dst + i * hidden_size_, hidden_size_);
            }
        }
    }
};

3.3 MNN推理接口使用

MNN模型加载
// src/llm.cpp中的模型加载过程
void Llm::load() {
    // 1. 配置MNN运行时参数
    MNN::ScheduleConfig config;
    config.type = backend_type_convert(config_->backend_type());  // CPU/GPU/NPU
    config.numThread = config_->thread_num();                     // 线程数
    
    // 2. 加载MNN模型文件
    auto model_path = config_->llm_model();
    auto runtime_manager = MNN::Express::ExecutorScope::Current()->getRuntime();
    
    // 3. 创建Module实例
    modules_.emplace_back(MNN::Module::load(
        {"input_ids", "attention_mask", "position_ids"},  // 输入名称
        {"output"},                                        // 输出名称
        model_path.c_str(),                               // 模型路径
        runtime_manager,                                   // 运行时
        &config                                           // 配置
    ));
}
MNN推理执行
// MNN前向推理的实现
MNN::Express::VARP forward(const std::vector<int>& input_ids) {
    // 1. 准备输入数据
    auto embeddings = embedding(input_ids);              // 词嵌入
    auto attention_mask = gen_attention_mask(input_ids); // 注意力掩码
    auto position_ids = gen_position_ids(input_ids);     // 位置编码
    
    // 2. 调用MNN推理
    auto outputs = modules_[0]->onForward({
        {"input_embeddings", embeddings},
        {"attention_mask", attention_mask},
        {"position_ids", position_ids}
    });
    
    // 3. 获取输出logits
    return outputs[0];  // shape: [batch_size, seq_len, vocab_size]
}

3.4 MNN输出处理和转换

Logits到Token的转换
// 从模型输出logits中采样下一个token
int sample_token(MNN::Express::VARP logits, float temperature = 1.0) {
    // 1. 获取最后一个位置的logits
    auto last_logits = logits[logits->getInfo()->dim[1] - 1];  // [vocab_size]
    
    // 2. 应用温度参数
    if (temperature != 1.0) {
        last_logits = last_logits / temperature;
    }
    
    // 3. 计算softmax概率分布
    auto probs = MNN::Express::_Softmax(last_logits, 0);
    
    // 4. 采样策略选择
    if (do_sample_) {
        // 随机采样
        return multinomial_sample(probs);
    } else {
        // 贪婪搜索:选择概率最大的token
        auto max_indices = MNN::Express::_ArgMax(probs, 0);
        return max_indices->readMap<int>()[0];
    }
}
自回归生成循环
// 文本生成的完整流程
std::vector<int> generate(const std::vector<int>& input_ids) {
    std::vector<int> generated_ids = input_ids;
    
    for (int step = 0; step < max_new_tokens_; step++) {
        // 1. MNN模型推理
        auto logits = forward(generated_ids);
        
        // 2. 采样下一个token
        int next_token = sample_token(logits, temperature_);
        
        // 3. 检查停止条件
        if (tokenizer_->is_stop(next_token)) {
            break;
        }
        
        // 4. 添加到序列中
        generated_ids.push_back(next_token);
        
        // 5. 实时输出(流式生成)
        std::cout << tokenizer_->decode(next_token) << std::flush;
    }
    
    return generated_ids;
}

3.5 性能监控实现

// 生成状态统计
struct GenerateState {
    int prompt_len_ = 0;          // 提示词长度
    int gen_seq_len_ = 0;         // 生成序列长度  
    int64_t prefill_us_ = 0;      // 预填充耗时(微秒)
    int64_t decode_us_ = 0;       // 解码耗时(微秒)
    
    // 计算性能指标
    float prefill_speed() const {
        return prompt_len_ / (prefill_us_ / 1e6f);  // tokens/秒
    }
    
    float decode_speed() const {
        return gen_seq_len_ / (decode_us_ / 1e6f);  // tokens/秒
    }
};

// 性能统计实现
void benchmark_performance(const GenerateState& state) {
    printf("=== 性能报告 ===\n");
    printf("提示词长度: %d tokens\n", state.prompt_len_);
    printf("生成长度: %d tokens\n", state.gen_seq_len_);
    printf("预填充速度: %.2f tok/s\n", state.prefill_speed());
    printf("生成速度: %.2f tok/s\n", state.decode_speed());
    printf("总耗时: %.2f ms\n", (state.prefill_us_ + state.decode_us_) / 1000.0);
}

4. 使用MNN的核心价值

通过业务流程分析,MNN框架为LLM部署提供了:

🔧 MNN框架提供

  • 高效的神经网络推理引擎
  • 跨平台的硬件适配(CPU/GPU/NPU)
  • 内存和计算资源优化
  • 模型加载和执行管理

💼 开发者专注

  • 业务逻辑设计(对话管理、用户交互)
  • 数据预处理(分词、格式转换)
  • 应用层优化(缓存策略、性能监控)
  • 用户体验(命令行界面、流式输出)

这种分工让开发者可以专注业务创新,无需关心底层推理引擎的复杂实现。

5. 总结

通过对MNN-LLM项目的深入分析,我们可以看到使用MNN框架部署大语言模型的完整技术路径:

5.1 核心业务流程清晰简洁

项目展示了三个核心流程的简洁实现:

  • 模型加载流程:从配置文件到MNN运行时的标准化初始化过程
  • 推理生成流程:文本预处理 → MNN推理 → 结果后处理的完整链路
  • CLI应用流程:从命令行参数到业务执行的用户友好界面

5.2 技术实现务实高效

在技术实现层面,项目采用了多项实用技术:

  • 多样化分词支持:SentencePiece、TikToken等主流分词器,适配不同模型需求
  • 灵活的嵌入处理:支持内存和磁盘两种嵌入模式,在性能和内存之间找到平衡
  • 标准的MNN接口:直接使用MNN::Module的标准API,降低学习和维护成本
  • 高效的输出转换:从logits到文本的完整转换链路,支持多种采样策略

5.3 MNN框架的实用价值

通过这个Demo项目,我们看到MNN框架的核心价值:

  • 开发效率高:开发者只需关注业务逻辑,无需处理底层推理优化
  • 部署门槛低:统一的接口设计,简化了模型加载和推理过程
  • 性能表现好:内置的量化、缓存等优化技术,适合资源受限环境
  • 平台兼容强:支持多种硬件后端,一套代码多平台部署

5.4 实际应用启示

作为一个实用的Demo项目,MNN-LLM为开发者提供了宝贵的实践参考:

  1. 技术选型:展示了如何选择合适的分词器、嵌入方案和推理策略
  2. 性能优化:通过磁盘嵌入、量化等技术实现内存和性能的平衡
  3. 工程实践:从配置管理到性能监控的完整工程化实现
  4. 用户体验:交互式对话和批量测试两种模式,满足不同使用场景

随着边缘AI和终端部署需求的增长,MNN-LLM这样的轻量级解决方案将发挥越来越重要的作用。对于希望在资源受限环境中部署大语言模型的开发者来说,这个项目提供了一个优秀的技术实现参考。


本文基于 mnn-llm 开源项目源码分析,该项目是使用MNN框架部署大语言模型的优秀实践案例