前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >DataParallel里为什么会显存不均匀以及如何解决

DataParallel里为什么会显存不均匀以及如何解决

作者头像
zenRRan
发布2023-03-03 14:35:58
1.3K0
发布2023-03-03 14:35:58
举报
文章被收录于专栏:深度学习自然语言处理
作者:台运鹏 (正在寻找internship...) 主页:https://yunpengtai.top

鉴于网上此类教程有不少模糊不清,对原理不得其法,代码也难跑通,故而花了几天细究了一下相关原理和实现,欢迎批评指正!

关于此部分的代码,可以去https://github.com/sherlcok314159/dl-tools查看

「在开始前,我需要特别致谢一下一位挚友,他送了我双显卡的机器来赞助我做个人研究,否则多卡的相关实验就得付费在云平台上跑了,感谢好朋友一路以来的支持,这份恩情值得一辈子铭记!这篇文章作为礼物赠与挚友。」

Why Parallel

我们在两种情况下进行并行化训练[1]:

  1. 「模型一张卡放不下」:我们需要将模型不同的结构放置到不同的GPU上运行,这种情况叫ModelParallel(MP)
  2. 「一张卡的batch size(bs)过小」:有些时候数据的最大长度调的比较高(e.g., 512),可用的bs就很小,较小的bs会导致收敛不稳定,因而将数据分发到多个GPU上进行并行训练,这种情况叫DataParallel(DP)。当然,DP肯定还可以加速训练,常见于大模型的训练中

这里只讲一下DP在pytorch中的原理和相关实现,即DataParallel和DistributedParallel

Data Parallel

实现原理

实现就是循环往复一个过程:数据分发,模型复制,各自前向传播,汇聚输出,计算损失,梯度回传,梯度汇聚更新,可以参见下[2]:

pytorch中部分关键源码[3]截取如下:

代码语言:javascript
复制
def data_parallel(
 module, 
 input, 
 device_ids, 
 output_device=None
):
    if not device_ids:
        return module(input)

    if output_device is None:
        output_device = device_ids[0]

    # 复制模型
    replicas = nn.parallel.replicate(module, device_ids)
    # 拆分数据
    inputs = nn.parallel.scatter(input, device_ids)
    replicas = replicas[:len(inputs)]
    # 各自前向传播
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    # 汇聚输出
    return nn.parallel.gather(outputs, output_device)

代码使用

因为运行时会将数据平均拆分到GPU上,所以我们准备数据的时候, batch size = per_gpu_batch_size * n_gpus

同时,需要注意主GPU需要进行汇聚等操作,因而需要比单卡运行时多留出一些空间

代码语言:javascript
复制
import torch.nn as nn
# device_ids默认所有可使用的设备
# output_device默认cuda:0
net = nn.DataParallel(model, device_ids=[0, 1, 2], 
                      output_device=None, dim=0)
# input_var can be on any device, including CPU
output = net(input_var)

接下来看个更详细的例子[4],需要注意的是被DP包裹之后涉及到模型相关的,需要调用DP.module,比如加载模型

代码语言:javascript
复制
class Model(nn.Module):
    # Our model
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        # for convenience
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())
        return output

bs, input_size, output_size = 6, 8, 10
# define inputs
inputs = torch.randn((bs, input_size)).cuda()
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [6, xxx] -> [2, ...], [2, ...], [2, ...] on 3 GPUs
  model = nn.DataParallel(model)
# 先DataParallel,再cuda
model = model.cuda()
outputs = model(inputs)
print("Outside: input size", inputs.size(),
   "output_size", outputs.size())
# assume 2 GPUS are available
# Let's use 2 GPUs!
#    In Model: input size torch.Size([3, 8]) output size torch.Size([3, 10])
#    In Model: input size torch.Size([3, 8]) output size torch.Size([3, 10])
# Outside: input size torch.Size([6, 8]) output_size torch.Size([6, 10])

# save the model
torch.save(model.module.state_dict(), PATH)
# load again
model.module.load_state_dict(torch.load(PATH))
# do anything you want

如果经常使用huggingface,这里有两个误区需要小心:

代码语言:javascript
复制
# data parallel object has no save_pretrained
model = xxx.from_pretrained(PATH)
model = nn.DataParallel(model).cuda()
model.save_pretrained(NEW_PATH) # error
# 因为model被DP wrap了,得先取出模型 #
model.module.save_pretrained(NEW_PATH)
代码语言:javascript
复制
# HF实现貌似是返回N个loss(N为GPU数量)
# 然后对N个loss取mean
outputs = model(**inputs)
loss, logits = outputs.loss, outputs.logits
loss = loss.mean()
loss.backward()

# 返回的logits是汇聚后的
# HF实现和我们手动算loss有细微差异
# 手动算略好于HF
loss2 = loss_fct(logits, labels)
assert loss != loss2
True

显存不均匀

了解前面的原理后,就会明白为什么会显存不均匀。因为GPU0比其他GPU多了汇聚的工作,得留一些显存,而其他GPU显然是不需要的。那么,解决方案就是让其他GPU的batch size开大点,GPU0维持原状,即不按照默认实现的平分数据

首先我们继承原来的DataParallel(此处参考[5])),这里我们给定第一个GPU的bs就可以,这个是实际的bs而不是乘上梯度后的。假如你想要总的bs为64,梯度累积为2,一共2张GPU,而一张最多只能18,那么保险一点GPU0设置为14,GPU1是18,也就是说你DataLoader每个batch大小是32,gpu0_bsz=14

代码语言:javascript
复制
class BalancedDataParallel(DataParallel):
    def __init__(self, gpu0_bsz, *args, **kwargs):
        self.gpu0_bsz = gpu0_bsz
        super().__init__(*args, **kwargs)

核心代码就在于我们重新分配chunk_sizes,实现思路就是将总的减去第一个GPU的再除以剩下的设备,源码的话有些死板,用的时候不妨参考我的[6]

代码语言:javascript
复制
def scatter(self, inputs, kwargs, device_ids):
    # 不同于源码,获取batch size更加灵活
    # 支持只有kwargs的情况,如model(**inputs)
    if len(inputs) > 0:
        bsz = inputs[0].size(self.dim)
    elif kwargs:
        bsz = list(kwargs.values())[0].size(self.dim)
    else:
        raise ValueError("You must pass inputs to the model!")

    num_dev = len(self.device_ids)
    gpu0_bsz = self.gpu0_bsz
    # 除第一块之外每块GPU的bsz
    bsz_unit = (bsz - gpu0_bsz) // (num_dev - 1)
    if gpu0_bsz < bsz_unit:
        # adapt the chunk sizes
        chunk_sizes = [gpu0_bsz] + [bsz_unit] * (num_dev - 1)
        delta = bsz - sum(chunk_sizes)
        # 补足偏移量
        # 会有显存溢出的风险,因而最好给定的bsz是可以整除的
        # e.g., 总的=52 => bsz_0=16, bsz_1=bsz_2=18
        # 总的=53 => bsz_0=16, bsz_1=19, bsz_2=18
        for i in range(delta):
            chunk_sizes[i + 1] += 1
        if gpu0_bsz == 0:
            chunk_sizes = chunk_sizes[1:]
    else:
        return super().scatter(inputs, kwargs, device_ids)

    return scatter_kwargs(inputs, kwargs, device_ids, chunk_sizes, dim=self.dim)

优缺点

  • 优点:便于操作,理解简单
  • 缺点:GPU分配不均匀;每次更新完都得销毁「线程」(运行程序后会有一个进程,一个进程可以有很多个线程)重新复制模型,因而速度慢

参考资料

[1]

Pytorch中多GPU并行计算教程: https://blog.csdn.net/qq_37541097/article/details/109736159

[2]

DDP: https://www.cnblogs.com/ljwgis/p/15471530.html

[3]

MULTI-GPU EXAMPLES: https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html?highlight=dataparallel

[4]

OPTIONAL: DATA PARALLELISM: https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html

[5]

transformer-xl: https://github.com/kimiyoung/transformer-xl

[6]

Balanced Data Parallel: https://github.com/sherlcok314159/dl-tools/blob/main/balanced_data_parallel/README.md


📝论文解读投稿,让你的文章被更多不同背景、不同方向的人看到,不被石沉大海,或许还能增加不少引用的呦~ 投稿加下面微信备注“投稿”即可。

本文参与 腾讯云自媒体同步曝光计划,分享自微信公众号。
原始发表:2022-12-13,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 深度学习自然语言处理 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • Why Parallel
  • Data Parallel
    • 实现原理
      • 代码使用
        • 显存不均匀
          • 优缺点
            • 参考资料
            相关产品与服务
            批量计算
            批量计算(BatchCompute,Batch)是为有大数据计算业务的企业、科研单位等提供高性价比且易用的计算服务。批量计算 Batch 可以根据用户提供的批处理规模,智能地管理作业和调动其所需的最佳资源。有了 Batch 的帮助,您可以将精力集中在如何分析和处理数据结果上。
            领券
            问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档