按需分配、任务隔离、灵活控制，飞桨显存分配策略重磅升级

用户1386409

发布于 2020-03-23 16:41:37

8340

发布于 2020-03-23 16:41:37

虽然GPU对深度学习计算有普遍明显的加速作用，但其显存也是有限的（如V100的最大显存值也仅有32G），而深度学习模型的训练和推理往往需要大量的显存，用来支持更大的模型和更大的batch size。如何更高效地利用GPU显存，在一张卡或一台机器上同时承载更多的训练和预测任务，让有限的显存支持多个开发者同时进行实验，执行各自的任务呢？

飞桨v1.7在GPU显存使用策略方面做了如下3点升级：

默认使用显存自增长AutoGrowth策略，根据模型实际占用的显存大小，按需自动分配显存，并且不影响训练速度。在模型实际占用显存不多的情况下，同一张GPU卡可以同时运行多个深度学习训练/预测任务。
支持限制任务的最大显存策略，每个任务只能在限定的显存量下运行，实现同一张GPU卡多个任务间的显存隔离。
默认使用Lazy显存分配方式。只有GPU卡工作时才自动分配显存，实现不同GPU卡上的任务的相互隔离，可以在一台机器上实现更灵活的任务排布。

这三种显存策略在飞桨是如何实现的？下面我们走进飞将框架一探究竟。

AutoGrowth实现显存按需分配，且不影响训练速度

1.7版本之前，飞桨默认是显存预分配策略（缺省比例是可用显存的92%），该策略在实现上是比较高效的，但是预分配比例的设置是一个比较头疼的事情。如果采用92%的缺省配置，可以保证大部分情况下任务成功分配，但麻烦的是启动任务之后，即使模型实际占用显存较小，也无法再启动其他的任务了。因此飞桨v1.7升级为显存自增长按需分配的AutoGrowth作为默认的显存分配策略。

考虑对模型训练速度的影响，如果直接使用cudaMalloc和cudaFree接口进行显存分配和释放，调用的过程非常耗时，会严重影响模型的训练和预测速度。飞桨v1.7对此进行了改进，AutoGrowth显存分配策略通过缓存显存的方式提升显存分配和释放速度，如图1所示。

图1 AutoGrowth显存分配策略示意图

框架内部缓存Memory Cache中缓存了若干显存块block，当用户申请request_size大小的显存时，AutoGrowth策略内部的具体流程如下：

若Memory Cache为空，或者所有block均小于request_size，通过调用cudaMalloc从GPU中申请显存（对应图1中的large_request_size）。
若Memory Cache不为空，且存在大于等或者request_size大小的block，查找到满足条件的最小block，从中分割出request_size大小的显存返回给用户。
显存释放时，释放的显存将存储到Memory Cache中，不再返回给GPU。

可以看到，整个策略设计是非常高效的。通过AutoGrowth策略，实现按需分配显存的同时，也保证了模型训练速度不受影响。

实验观察

下面以ResNet50 batch_size=32单卡训练为例，观察一下，执行AutoGrowth策略后，显存的占用情况。

首先您需要在本地安装飞桨1.7，然后在飞桨GitHub中下载模型库代码，运行如下命令，进入”models/PaddleCV/image_classification”目录。

git clone -b release/1.7 https://github.com/PaddlePaddle/models.git
cd models/PaddleCV/image_classification

执行如下命令启动ResNet50单卡训练任务。

export CUDA_VISIBLE_DEVICES=0
python train.py  --model=ResNet50  --data_dir=./data/ILSVRC2012/  --batch_size=32

data_dir：设置ImageNet数据集的路径。
batch_size：设置batch_size为32。

训练任务启动后，运行nvidia-smi命令，观察GPU显存的占用情况。

运行1个ResNet50训练任务，显存占用约4G。（飞桨1.7之前，运行1个ResNet50训练任务，显存空间就完全被占满。）

运行2个训练任务，显存占用约8G。

运行3个训练任务，显存占用约12G。

运行4个训练任务，显存占用约16G，此时显存完全被占满。

实验证明：使用AutoGrowth策略后，一张16G V100的GPU可以并行4个ResNet50训练任务。那么，AutoGrowth策略的使用会不会影响模型的训练速度呢？

在模型训练初期，由于Memory Cache为空或block数量很少，框架会先从GPU中申请显存。显存释放时会存储在Memory Cache中，因此Memory Cache中的block会不断增加。后续显存请求会越来越多的从Memory Cache中先获取到，因此使用AutoGrowth策略后，训练速度保持不变。

支持限制任务的最大显存策略，实现单卡多任务间的资源隔离

实际应用中，常会遇到多个开发者使用同一个GPU卡进行模型训练的场景，此时需要将GPU卡的显存分为若干份，分给开发者独立使用。飞桨1.7支持自定义每个任务使用的最大显存策略，用户只需要配置几个参数，即可实现同一张GPU卡多个任务间的资源隔离。

例如，若想限定某个任务的最大显存占用量不超过2048MB，代码如下：

export  FLAGS_gpu_memory_limit_mb=2048

环境变量FLAGS_gpu_memory_limit_mb表示限定每个任务的最大显存占用量，为uint64类型的整数，单位为MB。

默认值为0，表示飞桨任务可以使用所有可用的显存资源，不设上限。
FLAGS_gpu_memory_limit_mb > 0，表示飞桨任务仅可使用不超过FLAGS_gpu_memory_limit_mb MB的显存。

若FLAGS_gpu_memory_limit_mb > 0，飞桨框架内部会对GPU显存分配cudaMalloc和释放cudaFree这两个接口进行监控，保证任务占用的显存量不超过用户设定的阈值。

举例来说，用户通过下述代码申请了2G大小的Numpy数组，并拷贝到飞桨的GPU Tensor中。

import paddle.fluid as fluid
import numpy as np

# 申请2G大小的Numpy数组
two_gb_numpy_array = np.ndarray([2, 1024, 1024, 1024], dtype='uint8')

place = fluid.CUDAPlace(0)
t = fluid.Tensor()
t.set(two_gb_numpy_array, place) # 将2G大小的Numpy数组拷贝到GPU上

若未设置FLAGS_gpu_memory_limit_mb，上述飞桨任务可正常运行，任务占用2048MB显存。
若设置了FLAGS_gpu_memory_limit_mb=1024，则会报出显存不足错误，如图2所示。表明任务使用的最大显存量被限定为1024MB。

图2显存不足报错提示

默认LAZY显存分配方式，实现不同卡上训练任务的隔离

下面通过执行一段简单的飞桨训练代码，了解下使用LAZY策略后，显存分配方式的变化。假设用户有2张GPU卡，分别为GPU 0和GPU 1。GPU 1被训练任务占用了16092MB显存，几乎将显存完全占满。

在LAZY模式下，用户仍可以在GPU 0上执行训练任务，如下代码所示。

import paddle.fluid as fluid
import numpy as np

x = fluid.data(name='x', shape=[None, 784], dtype='float32')
fc = fluid.layers.fc(x, size=10)
loss = fluid.layers.reduce_mean(fc) 
sgd = fluid.optimizer.SGD(learning_rate=1e-3)
sgd.minimize(loss)

place = fluid.CUDAPlace(0)   # 使用GPU 0卡进行训练
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())

BATCH_SIZE = 32
BATCH_NUM = 1000000

for batch_id in range(BATCH_NUM):
    x_np = np.random.random([BATCH_SIZE, 784]).astype('float32')
    loss_np, = exe.run(fluid.default_main_program(), 
    feed={x.name: x_np}, fetch_list=[loss])
    print('Batch id {}, loss {}'.format(batch_id, loss_np))

运行nvidia-smi命令，观察GPU 0和GPU 1上的显存占用情况。