如何在 Android 端使用 Memory Mapping (mmap) 技术防止大模型推理引发 OOM

在移动端部署 AI 模型（如 LLM、大参数量 CNN）时，最常见的痛点就是 OOM (Out Of Memory)。传统的加载方式会将整个模型文件读取到内存堆（Heap）中，这对于动辄几百 MB 甚至 GB 级的模型来说是致命的。

本文将介绍如何通过 Memory Mapping (mmap) 技术，实现“零拷贝”加载，有效缓解 Android 系统的内存压力。

为什么传统加载会 OOM？

常规的读取方式使用 fread 或 std::ifstream，这需要经过以下步骤：
1. 在内核缓冲区读取文件。
2. 将数据拷贝到用户态内存空间（Malloc 分配的堆内存）。
3. 推理框架解析内存中的 Buffer。

这种方式会导致物理内存被瞬间占满，触发 Android 的 Low Memory Killer (LMK)。

Memory Mapping (mmap) 的优势

mmap 是一种 Linux 系统调用，它将文件直接映射到进程的 虚拟地址空间，而不是物理内存。
– 延迟加载：只有在推理时真正访问到某一块权重时，操作系统才会触发缺页异常（Page Fault）将其加载进物理内存。
– 减少拷贝：消除了内核空间到用户空间的拷贝过程。
– 自动回收：当系统内存紧张时，操作系统可以自动释放未使用的映射页，而无需交换分区。

实战代码：在 Android NDK 中使用 mmap

如果你正在开发底层推理引擎，可以使用如下方式实现模型加载：

#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

void* map_model_file(const char* path, size_t& out_size) {
    int fd = open(path, O_RDONLY);
    if (fd == -1) return nullptr;

    struct stat sb;
    fstat(fd, &sb);
    out_size = sb.st_size;

    // 使用 MAP_SHARED 进行只读映射
    void* addr = mmap(NULL, out_size, PROT_READ, MAP_SHARED, fd, 0);
    close(fd); // 映射成功后可关闭 fd

    if (addr == MAP_FAILED) return nullptr;
    return addr;
}

// 使用完毕后记得释放
// munmap(addr, out_size);

推理框架中的配置

现代主流推理框架都已经内置了对 mmap 的支持，你只需要配置正确的开关：

1. TensorFlow Lite (TFLite)

TFLite 默认推荐使用 Model::BuildFromFile，其内部就是基于 mmap 实现的：

// 正确姿势：不要手动读取 buffer，直接传路径
std::unique_ptr<tflite::FlatBufferModel> model =
    tflite::FlatBufferModel::BuildFromFile(\"/sdcard/model.tflite\");

2. MNN (Alibaba)

在 MNN 中，创建 Session 时可以通过配置 Memory_Mapping 来启用：

MNN::ScheduleConfig config;
config.type = MNN_FORWARD_CPU;
// MNN 默认在 Android 端加载模型时会尝试使用 mmap
auto session = net->createSession(config);

进阶优化建议

1. 权重对齐：为了让 mmap 发挥最大效能，模型权重的偏移量（Offset）应当与磁盘扇区/页面大小（通常 4KB）对齐。使用 flatc 或 MNN 的转换工具时，通常会自动处理。
2. 配合内存模式：在推理时，设置推理引擎为 Low Memory 模式，通常会减少中间张量（Intermediate Tensors）的驻留时间。
3. 多进程隔离：对于超大型模型，建议在独立的 Remote Service 进程中运行推理，即使崩溃也不会导致主界面卡死或进程重启。

总结

通过 mmap 技术，我们可以将大模型加载的物理内存占用从“模型总大小”降低到“当前活跃计算层大小”，这是在 Android 端跑通大模型的必经之路。