C++ 高效加载机器学习模型的完整实现

首先，使用 open() 打开文件，并检查文件大小。

fd = open(path.c_str(), O_RDONLY);
file_size = lseek(fd, 0, SEEK_END);

为了避免将整个文件加载到内存中，我使用 mmap() 将文件映射到虚拟内存。

mapped = mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0);

safetensors 文件中的第一个 uint64_t 用于存储 header 的大小。

uint64_t header_size;
memcpy(&header_size, mapped, sizeof(uint64_t));

char *header_ptr = static_cast<char *>(mapped) + sizeof(uint64_t);

std::string header(header_ptr, header_size);

header 包含每个 tensor 的 metadata 信息。为了解析这些 metadata，我使用了 nlohmann::json。

metadata = nlohmann::json::parse(header);

实际的模型参数位于 header 之后：

data_offset = sizeof(uint64_t) + header_size;

为了方便访问每一层的权重，我创建了一个 Tensor 结构体：

struct Tensor
{
    std::string name;
    std::string dtype;

    std::vector<int64_t> shape;

    uint64_t start;
    uint64_t end;

    char *data;
    size_t size;
};

通过 get_tensor() 函数从 safetensors 文件中获取各个 tensor。关键在于使用 tensor 的数据偏移量，计算出正确的内存位置。

Tensor get_tensor(const std::string &name)
{
    if (!metadata.contains(name))
        throw std::runtime_error("Tensor not found: " + name);

    auto json = metadata[name];

    Tensor tensor;

    tensor.name = name;
    tensor.dtype = json["dtype"];
    tensor.shape = json["shape"].get<std::vector<int64_t>>();

    tensor.start = json["data_offsets"][0].get<uint64_t>();
    tensor.end = json["data_offsets"][1].get<uint64_t>();

    tensor.size = tensor.end - tensor.start;

    tensor.data =
        static_cast<char *>(mapped) + data_offset + tensor.start;

    return tensor;
}

由于我的 PC 内存有限，因此我会避免加载完整的模型权重。权重会继续保留在内存映射文件中，仅在 inference 期间需要时才进行数值转换。需要注意的是，这种方式不会把实际参数复制到内存中，只会创建一个指针，指向内存映射文件中的对应位置。在 inference 期间，可以将原始内存指针转换为正确的数据类型，从而访问 tensor 数据。

embedding_weights = reinterpret_cast<uint16_t *>(tensor.data);
output.resize(embedding_dim);
size_t offset = token_id * embedding_dim;
for (int i = 0; i < embedding_dim; i++)
{
    output[i] = bf16_to_float(embedding_weights[offset + i]);
}

另一种方式是，为了获得更快的 inference 速度，可以将 tensor 复制到已分配的内存中，而不是直接访问内存映射文件。

tensor.data = static_cast<char *>(malloc(tensor.size));
memcpy(
    tensor.data,
    static_cast<char *>(mapped) + data_offset + start,
    tensor.size
);

uint16_t *weights = reinterpret_cast<uint16_t *>(tensor.data);
for (int i = 0; i < size; i++)
{
    float value = bf16_to_float(weights[i]);
}

这样会将 tensor 数据加载到普通的已分配内存中，从而能够在 inference 期间直接访问。

如需采取进一步措施，你可以考虑屏蔽此人和/或举报滥用行为。