端侧推理的内存映射（mmap）：在Android设备上利用分页机制加载大模型权重 - 智猿学院-前后端，数据库，人工智能，云计算等领域前沿技术讲座

Android 端侧推理：基于内存映射的分页加载大模型权重

大家好，今天我们来聊聊如何在 Android 设备上进行端侧推理，特别是针对那些模型权重体积庞大的情况。我们重点讨论利用内存映射（mmap）和分页机制来解决大模型权重加载的问题。

1. 端侧推理的挑战

在移动设备上进行机器学习推理，相比于服务器端，面临着诸多挑战：

资源限制： 移动设备的内存、CPU、GPU 资源都相对有限。
功耗限制： 推理过程需要尽可能降低功耗，延长电池续航。
模型体积： 深度学习模型的体积越来越大，难以一次性加载到内存中。
启动速度： 应用启动时加载模型，需要尽可能缩短启动时间。

对于大模型而言，一次性加载所有权重数据到内存中，很容易导致内存溢出（OOM）错误，或者显著增加应用启动时间。因此，我们需要一种高效的方式来管理模型权重，按需加载，减少内存占用。

2. 内存映射（mmap）机制

内存映射（Memory Mapping）是一种将文件或设备映射到进程地址空间的技术。通过 mmap，进程可以直接像访问内存一样访问文件内容，而无需显式地进行读写操作。这为我们加载大模型权重提供了很大的便利。

mmap 的工作原理：

创建映射： mmap() 系统调用在进程的虚拟地址空间中创建一个映射区域，将其与指定的文件或设备关联起来。
按需加载： 操作系统内核采用分页机制，只有当进程真正访问映射区域的某个页面时，才会将该页面从磁盘加载到物理内存中。
共享内存： 多个进程可以映射同一个文件，从而实现进程间共享内存。

mmap 的优势：

节省内存： 无需将整个文件加载到内存，只在需要时加载部分页面。
提高效率： 避免了频繁的读写操作，直接通过指针访问文件内容。
简化编程： 可以像访问内存一样访问文件，简化了文件操作的代码。

3. 利用 mmap 加载大模型权重

我们可以利用 mmap 将大模型权重文件映射到内存中，然后按需访问模型权重数据。

步骤如下：

打开模型权重文件： 使用 fopen() 或者 Android NDK 中的 AAssetManager_open() 函数打开模型权重文件。
获取文件大小： 使用 fseek() 和 ftell() 或者 AAsset_getLength() 获取模型权重文件的大小。
创建内存映射： 使用 mmap() 函数创建一个内存映射区域，将文件映射到进程的虚拟地址空间中。
访问模型权重数据： 通过指针访问映射区域中的数据，就像访问内存一样。
释放内存映射： 使用 munmap() 函数释放内存映射区域。

C/C++ 代码示例：

#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>
#include <android/asset_manager.h>
#include <android/asset_manager_jni.h>
#include <android/log.h>

#define LOG_TAG "MMapExample"
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, LOG_TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)

// JNI 函数，用于从 assets 目录加载模型权重
extern "C" JNIEXPORT jlong JNICALL
Java_com_example_mmapexample_MMapHelper_loadModel(JNIEnv *env, jobject thiz, jobject assetManager, jstring assetName) {
    const char *name = env->GetStringUTFChars(assetName, nullptr);
    if (name == nullptr) {
        LOGE("Failed to get asset name");
        return 0;
    }

    AAssetManager* mgr = AAssetManager_fromJava(env, assetManager);
    if (mgr == nullptr) {
        LOGE("Failed to get asset manager");
        env->ReleaseStringUTFChars(assetName, name);
        return 0;
    }

    AAsset* asset = AAssetManager_open(mgr, name, AASSET_MODE_UNKNOWN);
    if (asset == nullptr) {
        LOGE("Failed to open asset: %s", name);
        env->ReleaseStringUTFChars(assetName, name);
        return 0;
    }

    off_t asset_length = AAsset_getLength(asset);
    if (asset_length <= 0) {
        LOGE("Invalid asset length: %lld", (long long)asset_length);
        AAsset_close(asset);
        env->ReleaseStringUTFChars(assetName, name);
        return 0;
    }

    void* buffer = mmap(nullptr, asset_length, PROT_READ, MAP_PRIVATE, AAsset_getFd(asset), AAsset_getOffset(asset));
    if (buffer == MAP_FAILED) {
        LOGE("mmap failed: %s", strerror(errno));
        AAsset_close(asset);
        env->ReleaseStringUTFChars(assetName, name);
        return 0;
    }

    AAsset_close(asset);
    env->ReleaseStringUTFChars(assetName, name);

    return reinterpret_cast<jlong>(buffer); // 返回映射区域的指针
}

// JNI 函数，用于释放内存映射
extern "C" JNIEXPORT void JNICALL
Java_com_example_mmapexample_MMapHelper_unloadModel(JNIEnv *env, jobject thiz, jlong buffer, jlong assetLength) {
    void* ptr = reinterpret_cast<void*>(buffer);
    if (ptr != nullptr) {
        if (munmap(ptr, assetLength) != 0) {
            LOGE("munmap failed: %s", strerror(errno));
        } else {
            LOGI("Model unloaded successfully");
        }
    }
}

Java 代码示例：

package com.example.mmapexample;

import android.content.Context;
import android.content.res.AssetManager;

public class MMapHelper {

    static {
        System.loadLibrary("mmapexample"); // 加载 native 库
    }

    // 声明 native 方法
    public native long loadModel(AssetManager assetManager, String assetName);
    public native void unloadModel(long buffer, long assetLength);

    private long modelBuffer;
    private long modelSize;

    public MMapHelper() {
    }

    public boolean loadModelFromAssets(Context context, String modelName) {
        AssetManager assetManager = context.getAssets();
        modelBuffer = loadModel(assetManager, modelName);
        if (modelBuffer == 0) {
            return false;
        }

        // 此处需要自己确定模型文件的大小。可以考虑在文件读取后存储。
        // 或者使用文件描述符和文件长度来代替 assetManager
        // modelSize = ...;  // TODO: 获取模型文件大小，需要自己实现
        // 这里为了方便演示，假设大小为 1MB
        modelSize = 1024 * 1024;

        return true;
    }

    public void unloadModel() {
        if (modelBuffer != 0) {
            unloadModel(modelBuffer, modelSize);
            modelBuffer = 0;
            modelSize = 0;
        }
    }

    public long getModelBufferAddress() {
        return modelBuffer;
    }
}

代码解释：

C++ 部分：
- loadModel() 函数：
  - 接收 Java 传递的 AssetManager 和模型文件名。
  - 打开模型文件，获取文件大小。
  - 使用 mmap() 创建内存映射，将模型文件映射到内存中。
  - 返回映射区域的指针给 Java 层。
- unloadModel() 函数：
  - 接收 Java 传递的映射区域指针和文件大小。
  - 使用 munmap() 释放内存映射。
Java 部分：
- MMapHelper 类：
  - 加载 native 库。
  - 声明 native 方法 loadModel() 和 unloadModel()。
  - 提供 loadModelFromAssets() 和 unloadModel() 方法，方便 Java 层调用。
  - getModelBufferAddress()方法返回模型权重在内存中的地址，方便后续推理。

注意事项：

权限： 确保应用具有读取文件权限。
文件描述符： mmap() 需要文件描述符作为参数，可以通过 fileno() 函数从 FILE* 指针获取。或者如示例代码所示，使用AAsset_getFd直接获取文件描述符，对于从assets加载的情况。
错误处理： mmap() 和 munmap() 可能会失败，需要进行错误处理。
同步： 如果多个线程访问同一个映射区域，需要进行同步，避免数据竞争。
64 位系统： 在 64 位系统上，指针类型是 64 位的，需要使用 jlong 类型来传递指针。
获取文件大小： 模型文件的大小需要在创建映射前获取。可以事先将模型大小存储在配置文件中，或者在模型文件本身中存储模型大小信息。

4. 分页机制的优势

mmap 与操作系统内核的分页机制结合，可以实现按需加载模型权重，极大地节省内存。

分页机制的工作原理：

虚拟地址空间： 每个进程都拥有独立的虚拟地址空间。
页面： 虚拟地址空间被划分为固定大小的页面（通常为 4KB）。
页表： 操作系统维护一个页表，用于将虚拟地址映射到物理地址。
按需加载： 当进程访问一个尚未加载到物理内存的页面时，会触发一个缺页中断。操作系统会将该页面从磁盘加载到物理内存中，并更新页表。
页面置换： 当物理内存不足时，操作系统会根据一定的算法（例如 LRU）选择一个页面置换出去，将其写回磁盘。

mmap 与分页机制的结合：

当我们使用 mmap 将模型权重文件映射到内存中时，操作系统并不会立即将整个文件加载到物理内存中。只有当我们真正访问某个页面时，才会触发缺页中断，将该页面加载到物理内存中。这样，我们就可以按需加载模型权重，避免一次性加载所有数据，节省内存。

表格对比：一次性加载 vs. mmap + 分页

特性	一次性加载	mmap + 分页
内存占用	加载整个模型文件，占用大量内存	只加载需要的页面，占用少量内存
加载时间	加载整个模型文件，耗时较长	只加载需要的页面，加载速度快
适用场景	模型体积小，内存充足	模型体积大，内存有限
实现复杂度	简单	稍复杂，需要处理内存映射和分页机制

5. Android NDK 的使用

在 Android 上，我们通常使用 NDK（Native Development Kit）来调用 mmap 函数。NDK 提供了一组 C/C++ 工具链，可以让我们在 Android 应用中使用 native 代码。

使用 NDK 的步骤：

安装 NDK： 在 Android Studio 中安装 NDK。
配置 Gradle： 在 build.gradle 文件中配置 NDK。
编写 native 代码： 使用 C/C++ 编写 native 代码，例如上面的 loadModel() 和 unloadModel() 函数。
编译 native 代码： 使用 NDK 编译 native 代码，生成动态链接库（.so 文件）。
加载 native 库： 在 Java 代码中使用 System.loadLibrary() 加载动态链接库。
调用 native 方法： 在 Java 代码中调用 native 方法。

6. 性能优化

使用 mmap 加载大模型权重可以有效节省内存，但仍然需要注意性能优化：

内存对齐： 确保模型权重数据在内存中对齐，可以提高访问效率。可以使用 posix_memalign() 函数来分配对齐的内存。
预加载： 可以预先加载一些常用的页面，减少推理时的缺页中断次数。
页面大小： 了解 Android 设备的页面大小，可以更好地管理内存。
I/O 优化： 尽量减少 I/O 操作，例如使用缓存、异步 I/O 等。

7. 实际应用场景

mmap 在端侧推理中有广泛的应用场景：

大型语言模型： LLM 模型通常体积巨大，使用 mmap 可以有效降低内存占用。
图像识别模型： 高分辨率图像识别模型也需要大量的模型权重数据。
推荐系统模型： 复杂的推荐系统模型也需要大量的模型权重数据。

例如，我们可以使用 mmap 加载一个大型的 Transformer 模型，然后使用该模型进行文本生成、机器翻译等任务。

8. 其他替代方案

除了 mmap，还有其他一些替代方案可以用于加载大模型权重：

模型量化： 将模型权重从浮点数转换为整数，可以减少模型体积。
模型压缩： 使用剪枝、蒸馏等技术压缩模型，减少模型体积。
模型分割： 将模型分割成多个部分，按需加载。
流式加载： 从网络或本地文件流式加载模型权重。

这些方案各有优缺点，需要根据实际情况选择合适的方案。

9. 一些思考

使用 mmap 加载大模型权重是一种有效的技术，但也存在一些挑战：

代码复杂度： 使用 mmap 需要编写 native 代码，增加了代码复杂度。
平台兼容性： 不同的 Android 设备可能存在差异，需要进行兼容性测试。
安全性： 需要注意内存映射的安全性，避免恶意代码访问模型权重数据。

未来，随着移动设备性能的提升和技术的进步，端侧推理将会变得更加普及和高效。

今天我们主要讨论了使用内存映射（mmap）和分页机制在 Android 设备上加载大模型权重，有效节省内存，提高推理效率，并介绍了 NDK 的使用和性能优化方向。希望这次讲座对大家有所帮助。