PaddlePaddle 加载 ERNIE 开源模型时 Place 设置为 GPU 但显存碎片化的预分配方案

在 PaddlePaddle 中加载 ERNIE 等开源模型并指定 GPU 设备时，显存碎片化通常不会在加载权重那一刻出现，而是在之后训练或推理的若干步中慢慢显现：每次前向反向都会创建和释放大量中间 Tensor，分配器如果按需从 GPU 申请新显存，池内就会出现许多不连续的空闲块。到最后即使整卡还有剩余显存，也可能因找不到连续大块内存而抛出 CUDA out of memory。预分配方案就是在加载模型之前调整 PaddlePaddle 的显存分配策略和池大小，让分配器一次性预留一块比较规整的显存池，并在后续复用空闲块，而不是频繁向驱动申请。

显存碎片化通常在模型加载后多次迭代时出现，核心解决思路是设置 FLAGS_allocator_strategy 为 naive_best_fit，并配合 FLAGS_initial_gpu_memory_in_mb 指定初始显存池大小。该方案能降低 OOM 概率，但会增加峰值显存占用，需要根据实际模型和设备内存调整，并观察是否真的消除了阶梯式显存增长。

先判断碎片化还是真的显存不足

在改动配置前，需要区分两种不同情况。一种是真的模型加批数据超出显存容量，另一种是显存池碎片太多。先用下面的方式观察：

启动训练后，用 nvidia-smi 每隔几秒记录一次显存占用；如果占用在稳定后仍持续小幅上涨，或者每次上涨后下降不到原来的水位，可能是碎片积累。
把批量尺寸缩小到原来的 1/2 或 1/4，如果程序能正常运行且显存峰值明显下降，说明原始设置可能已经接近上限；这时调整分配策略只能有限改善。
在 Python 里打印 PaddlePaddle 自身的显存状态，例如 paddle.device.cuda.memory_used() 与 paddle.device.cuda.memory_reserved()。如果 reserved 远大于 used，并且 reserved 增长曲线呈现很多小台阶，大概率是分配器在反复申请小块显存。

这一步的目的是确定预分配是否真的需要。如果只是显存不足，预分配反而可能让问题更早暴露。

选择分配策略并设置显存池参数

PaddlePaddle 的 GPU 显存分配器主要有两种常见策略。默认或常用的是 auto_growth，即先申请一小块，不够再追加；另一种是 naive_best_fit，会在显存池中维护空闲块列表，优先复用足够大的空闲块。对于加载 ERNIE 这类包含大量固定参数的模型，建议优先使用 naive_best_fit，并在加载前设置初始池大小，让分配器以预分配方式工作。

可以用环境变量和 paddle.set_flags 双重设置，避免某些版本下单个入口不生效。以下参数是相对通用的：

参数	作用	建议设置方向
`FLAGS_allocator_strategy`	选择显存分配策略	设为 `naive_best_fit`
`FLAGS_initial_gpu_memory_in_mb`	初始显存池大小	设为模型权重加典型批次前向所需显存的 1.1 到 1.3 倍，先用原有配置试跑一轮再确定
`FLAGS_max_gpu_memory_in_mb`	显存池上限	设为当前 GPU 总显存减去 300-800MB 保留给其他进程和驱动，避免直接占满导致其他任务失败
`FLAGS_reallocate_gpu_memory_in_mb`	每次扩展池的最小增量	可设置成 1024 或 2048，减少后续小步追加造成的碎片

注意，这些环境变量要在 PaddlePaddle 创建 GPU Place 和 Executor 之前生效。若在 Python 脚本里设置，必须放在 import paddle 之后、创建 Place 之前；如果已经创建了 Place，可以重新启动进程让配置生效。

在加载 ERNIE 前应用的推荐代码

下面这段代码可以作为统一的入口片段，放在加载 ERNIE 模型之前。数字需要根据实际显卡和模型替换，不能直接照抄。

PaddlePaddle 加载 ERNIE 开源模型时 Place 设置为 GPU 但显存碎片化的预分配方案

import os
import paddle

# 方式一：环境变量，先于 paddle.set_device 生效
os.environ['FLAGS_allocator_strategy'] = 'naive_best_fit'
os.environ['FLAGS_initial_gpu_memory_in_mb'] = '4096'    # 按模型和批次调整
os.environ['FLAGS_max_gpu_memory_in_mb'] = '8192'        # 留出系统和其他进程空间
os.environ['FLAGS_reallocate_gpu_memory_in_mb'] = '1024'

# 方式二：paddle.set_flags，部分版本需要在底层初始化前调用
paddle.set_flags({
    'FLAGS_allocator_strategy': 'naive_best_fit',
})

place = paddle.set_device('gpu')

# 然后再加载 ERNIE 模型
model = ErnieModel.from_pretrained(...)

如果是在命令行启动脚本，也可以直接写成：FLAGS_allocator_strategy=naive_best_fit FLAGS_initial_gpu_memory_in_mb=4096 python train.py。这种方式适合在 Docker 或 Kubernetes 环境中通过环境变量注入，不需要改代码。

需要特别确认的是，PaddlePaddle 不同版本对 FLAGS_allocator_strategy 的取值可能略有差异，例如部分版本接受 budget_best_fit。如果你在运行日志看到 Invalid enum value，先通过 paddle.flags 查看当前版本支持的取值，再选择策略。

验证预分配是否生效与风险边界

设置完成后，运行同样的训练或推理任务，并做以下验证：

启动后立即用 nvidia-smi 观察显存占用，如果初始就一次性分配了接近预设池大小的显存，说明预分配生效；如果仍然从很小逐步上涨，则配置可能没有在正确的时间点生效。
连续跑若干个训练 step，观察 reserved 显存曲线是否保持平缓；如果中途还是出现台阶式增长，说明仍有超出预设池的分配发生，可尝试调大 FLAGS_initial_gpu_memory_in_mb 或 FLAGS_reallocate_gpu_memory_in_mb。
把分配策略改回默认值，在同样批量下跑同一模型，如果默认策略出现 OOM 而预分配策略正常，则说明碎片化确实被缓解；如果两种策略都在同一批量崩溃，则是单卡容量不足。

风险方面，预分配会提高进程启动时的显存占用，因此多进程或多卡环境下要预留足够空间，否则其他进程可能无法启动。其次，naive_best_fit 在分配和释放频繁时搜索空闲块的代价更高，可能会增加少量 CPU 时间；不过对于加载 ERNIE 后相对稳定的训练任务，影响通常不明显。最后，FLAGS_max_gpu_memory_in_mb 设置过小会直接导致 OOM，设置过大则可能让驱动没有足够空闲内存供显示和库使用，建议先保留 300MB 以上余量，再结合具体环境微调。

如果调整后仍然出现碎片化，还可以考虑同一批次内减少动态 shape、尽量复用中间 Tensor，或在 ERNIE 模型后接 FP16 混合精度来缩小单个 Tensor 体积。这些动作和预分配策略并不冲突，可以一起配合使用，但不能用预分配掩盖显存真实不足的问题。