显存优化综述

简介: 显存优化综述

动态图和静态图


动态图是边计算边搭建计算图,中间结果一目了然。


静态图是先搭建计算图,最后喂入输入,直接出结果。


因为静态图在编译的时候进行了一些优化,比如改变了代码的计算过程,以利于GPU更好的并行计算,那么想debug看中间结果就很困难了,有时候看到的和期望的可能会有些出入。另外就是每次当我们搭建完一个动态计算图,然后在反向传播结束之后,整个计算图就在内存中被释放了。如果想再次使用的话,必须从头再搭一遍。


静态图显存优化


对于大网络结构的训练场景,在静态图上的显存优化主要可以分为三个方向:


静态显存分配。由于获得了整张计算图,所以可以去分析每一个 tensor 和每个算子的生命周期。对于生命周期没有重叠的算子,它们是可以共享显存的。


梯度检查点(用计算换显存)。设置一些梯度检查点,剩下的中间结果就先释放掉,如果将来在反向传播的过程中发现前向结果不在显存中,就找到最近的梯度检查点,恢复出被释放的 tensor。


内存交换(用带宽换显存)。把暂时不用的数据从 GPU 上交换到 CPU 上,到了需要的时候,再把它交换回来。


动态图显存优化


动态图无法提前获得全局的计算图信息。因为无法得到每个 tensor 的生命周期,所以静态显存分配不再可用


  1. 1.用计算换显存,也就是动态图版的 Sublinear 显存优化;

  2. 2.用带宽换显存,在 GPU 和 CPU 之间交换内容。
  3. 19.png
  4. 如上图所示,交换耗时比计算耗时高出很多,因此用带宽换显存不合理。
  5. 20.gif
  6. 如上图所示,在前向传播中(第一行从左到右),蓝色圆圈表示模型的中间计算结果开始占用显存。一直到前向传播完成,第一行完全变为蓝色圆圈,前面计算所占用的显存都不能释放。


等到反向传播开始(第二行从右到左),随着梯度的计算与完成应用,前向传播保留在显存中的张量才可以释放。


很明显,如果要降低显存占用,就要拿前向传播保存的中间计算结果开刀,这也正是 MegEngine 动态图显存优化的主要方向。

21.gif

如上为梯度检查点技术原理示意,前向传播中第三个点为检查点,它会一直保存在显存中。第四个点在完成计算后即可释放显存,在反向传播中如果需要第四个点的值,可以从第三个点重新计算出第四个点的值。


DTR


问题


大部分场景下,网络训练的中间特征占用了绝大部分得到显存,模型权重这是占据了极小的一部分。


这些中间特征在反向传播的算子中还需要用到,所以不好被释放。


解决办法


引入LRU cache的机制,选择代价最低的tensor进行释放,在需要用到的时候进行重计算。


重计算的开销越小越好


释放的tensor占用显存越大越好


该tensor的访问间隔越长越好


LRU cache:距离上次访问时间间隔最长的,进行释放


另外,DTR 论文中还提出,除了重计算带来的开销之外,其他的额外开销主要用于寻找应该被释放掉的最优 tensor。因为在显存中,tensor 停留的时长是不断在变化的,所以只能在需要释放的时候现场计算最优的 tensor。


对此,论文中提出了两个运行时的优化技巧:


不考虑小的 tensor,当 tensor 大小小于候选集中的 tensor 的平均大小的 1% 时,不加入候选集;


每次在需要释放 tensor 的时候,随机采样 sqrt(N) 个 tensor 进行遍历(N 为目前可释放的 tensor 候选集的大小)

23.png

重计算的局限性


  • 释放的显存不连续,形成的显存碎片,无法容纳新的tensor

例如,新的tensor需要100M的显存,为此释放了两个tensor,但是这两个tensor不是连续的,不能被使用,根据释放机制,就会一直释放下去,直到释放出一段连续可用的显存

24.png

在显存池里的排列并非最优,是按照计算顺序排列的

25.png

26.png

in-place op 失效


重计算实质上把in-place op变成了非in-place op


in-place op:模型权重会被修改利用,以此来节省显存和cache,原地修改会造成后续的值改变。


DTR:模型权重不可被修改,额外申请资源进行计算,生成新的tensor,分散在显存池中,很难形成连续大显存。

27.png

Coop


对显存的排列方式进行优化,一次性找出可以生成足够大的连续空闲显存、并且总代价最低的tensor集合。

28.png

29.png

recomputable in place


对于不再变化的tensor,进行共享存储。

30.png

op-guided allocation


根据op进行显存池中位置的分配,相同的放在一起。(原先是按照计算顺序来的)

31.png

Layout-aware Eviction


将显存池中空闲显存当做代价为0的tensor。


利用滑动窗口找出代价最小的连续tensor,进行释放。

32.png


实验效果


消耗的时间


下图是Coop和同样在OneFlow中实现的DTR、DTE策略在八种不同的网络、不同显存阈值下的对比,可以看到Coop均超过了其它两种方法:


横轴为显存,纵轴为时间倍数

33.png

显存碎片率


同时,Coop 将显存碎片率减少了一个量级(注意 BiLSTM 的 y 轴为对数坐标)

在大部分网络中,显存碎片率最低

34.png

搜索时间


Coop 的搜索过程时间复杂度为 O(N) 而不是 O(N^2)(N 为显存池中的 tensor 个数),在

绝大多数场景下也取得了最快的搜索速度(注意 BiLSTM 和 BERT Large 的 y 轴为对数坐

标)

35.png

目录
相关文章
|
JSON 算法 安全
JWT渗透与攻防
JSON Web Token(JWT)作为一种广泛使用的认证机制,因其简单、有效和灵活性而受到开发者的青睐。然而,随着JWT使用的普及,其安全性问题也逐渐显现。本文将探讨JWT的渗透测试方法,分析常见的安全漏洞,并提供相应的防护措施。
1414 3
JWT渗透与攻防
|
机器学习/深度学习 前端开发 数据可视化
Kimi K2 开源发布:擅长代码与 Agentic 任务!
今天,月之暗面正式发布 Kimi K2 模型,并同步开源。Kimi K2 是一款具备更强代码能力、更擅长通用 Agent 任务的 MoE 架构基础模型,总参数 1T,激活参数 32B。
1950 0
|
运维 关系型数据库 MySQL
宝塔面板 Docker 镜像加速教程
本文介绍了如何为宝塔面板中的 Docker 配置轩辕镜像加速。内容包括:安装 Docker 管理器、获取镜像加速地址、配置加速 URL 以及验证设置等步骤。通过使用免费版或专业版加速地址(如 docker.xuanyuan.me 或 xuanyuan.cloud),可显著提升镜像拉取速度。同时,还提供了日常使用方法和注意事项,帮助用户顺利完成配置并保障 Docker 服务稳定运行。
1103 0
|
人工智能 JSON 网络协议
音乐人狂喜!AbletonMCP:让AI帮你写歌,一句话生成专业编曲,Demo级作品秒出
AbletonMCP 是一个开源项目,通过模型上下文协议(MCP)将 Ableton Live 与 Claude AI 连接,实现 AI 辅助音乐制作,支持创建、修改 MIDI 和音频轨道等操作。
1210 10
音乐人狂喜!AbletonMCP:让AI帮你写歌,一句话生成专业编曲,Demo级作品秒出
|
安全 算法 数据安全/隐私保护
远控安全进阶之战:TeamViewer/ToDesk/向日葵设备安全策略对比
在数字化时代,卓越的远程控制软件需兼顾功能与体验,包括流畅连接、高清画质、低门槛UI设计、毫秒级延迟及多功能性,同时要有独树一帜的远控安全技术,通过前瞻性安全策略阻挡网络风险,为用户打造全方位安全体验,在基础安全防护上不断创新。本文选取了当下热门的TeamViewer、ToDesk、向日葵三款远程控制软件,从软件介绍、远控安全策略等多个维度进行深入对比,以呈现一场精彩的远控软件安全进阶“大比拼”。
|
Web App开发 缓存 安全
让你的win10/win11系统变得不再卡顿,优雅草伊凡整理-长期更新-如何让windows操作系统不用老是重装依然保持流畅运行
如题:让你的win10/win11系统变得不再卡顿,优雅草伊凡整理-长期更新-如何让windows操作系统不用老是重装在不断的更新中依然保持流畅运行 本文长期更新,本次更新2023年11月8日! 很多时候 我们的win10win11系统不管再怎么关闭更新,都会莫名又被打开,莫名的关机和自动更新,总是在你不注意的时候就更新,因此大家也明白微软的厉害了,扯远了,但是频繁的更新会让系统特别卡,传统办法就是重装系统,因而还衍生了很多一件重装工具,但是每一次重装的代价都是巨大的我相信每个同学都明白。
1567 4
|
编解码 人工智能 缓存
轻装出海:一站式微短剧出海指南
轻装出海:一站式微短剧出海指南
1334 0
|
存储 缓存 监控
Linux性能分析工具-perf并生成火焰图
Linux性能分析工具-perf并生成火焰图
1538 2
|
存储 网络协议 Unix
docker的底层原理一:客户端-服务器架构
本文详细解释了Docker的客户端-服务器架构,包括常驻后台的Docker守护进程、通过命令行接口发送请求的Docker客户端、以及它们之间通过Unix socket或网络接口进行的通信。
598 0
|
Kubernetes 安全 调度
深度解读:阿里云全球首发的容器计算服务 ACS 诞生背景、核心技术与应用场景
深度解读:阿里云全球首发的容器计算服务 ACS 诞生背景、核心技术与应用场景
70971 45

热门文章

最新文章