BERT基础模型.pt权重文件的高效加载与优化实践

1次阅读

没有评论

共计 2153 个字符，预计需要花费 6 分钟才能阅读完成。

在处理 BERT 等大型预训练模型时，原始的.pt 权重文件通常会带来几个明显的性能瓶颈：

内存占用高：基础 BERT 模型的完整权重文件通常超过 400MB，加载时会将所有参数一次性读入内存
加载速度慢 ：传统 torch.load() 方式受限于磁盘 IO 速度，尤其在机械硬盘上可能产生数秒延迟
显存压力大：在多 GPU 环境下，每个进程都会独立加载权重副本，造成显存浪费

原理：将 FP32 权重转换为低精度格式(如 FP16/INT8)
优点：直接减少 50-75% 存储空间，提升加载速度
限制：可能带来轻微精度损失，需测试下游任务表现

原理：按需分块加载权重，而非一次性全量加载
优点：显著降低峰值内存占用
限制：增加代码复杂度，需处理权重依赖关系

原理：在共享存储中维护单一权重副本，各进程通过内存映射访问
优点：消除多进程冗余加载，节省整体资源
限制：需要共享文件系统支持

# FP16 量化示例
from torch import nn, half
def quantize_model(model):
    model.half()  # 转换权重为 FP16
    for layer in model.modules():
        if isinstance(layer, nn.LayerNorm):
            layer.float()  # 归一化层保持 FP32
    return model

import torch
import os

class ChunkedLoader:
    def __init__(self, file_path, chunk_size=50):
        self.weights = torch.load(file_path)
        self.keys = list(self.weights.keys())
        self.chunk_size = chunk_size

    def __iter__(self):
        for i in range(0, len(self.keys), self.chunk_size):
            chunk_keys = self.keys[i:i+self.chunk_size]
            yield {k: self.weights[k] for k in chunk_keys}

# 使用示例
loader = ChunkedLoader("bert_base.pt")
for chunk in loader:
    # 逐块处理权重
    process_chunk(chunk)

# 使用 numpy 的 memmap 作为中间格式
import numpy as np

def save_memmap(weights, output_path):
    # 将所有张量拼接为连续数组
    total_size = sum(w.numel() for w in weights.values())
    arr = np.memmap(output_path, dtype='float32', mode='w+', shape=(total_size,))

    # 保存元数据(记录各张量偏移量)
    meta = {}
    ptr = 0
    for name, tensor in weights.items():
        flat = tensor.cpu().numpy().ravel()
        arr[ptr:ptr+flat.size] = flat
        meta[name] = (ptr, ptr+flat.size)
        ptr += flat.size

    # 保存元数据文件
    torch.save(meta, f"{output_path}.meta")

# 加载时只需读取需要的部分
arr = np.memmap("weights.dat", dtype='float32", mode='r')
meta = torch.load("weights.dat.meta")

def load_tensor(name):
    start, end = meta[name]
    return torch.from_numpy(arr[start:end].copy())