【AMP实操】解放你的GPU运行内存!在pytorch中使用自动混合精度训练

简介: 【AMP实操】解放你的GPU运行内存!在pytorch中使用自动混合精度训练

前言

  自动混合精度(Automatic Mixed Precision,简称AMP)是一种深度学习加速技术,它通过在训练过程中自动选择合适的数值类型(如半精度浮点数和单精度浮点数)来加速计算,并减少内存占用,从而提高训练速度和模型性能。

精度

半精度

  半精度浮点数(Half-Precision Floating Point)是一种浮点数数据类型,也被称为16位浮点数,它可以表示的范围和精度都比单精度浮点数(32位)低。半精度浮点数使用1个符号位、5个指数位和10个尾数位来表示一个实数。它可以表示的最大正数是2^16-2^10 ≈ 6.55×10^4,最小正数是2^-14 ≈ 6.10×10^-5,可以表示的精度为2^-10 ≈ 9.77×10^-4。

  半精度浮点数在深度学习中被广泛使用,因为它可以显著地降低模型的计算和存储成本,同时仍能保持较高的训练精度。在卷积神经网络(Convolutional Neural Network,CNN)等深度学习模型中,半精度浮点数通常被用来存储权重和梯度,而输入和中间特征映射通常使用单精度浮点数或者更高精度的浮点数来存储。

全精度

  全精度浮点数是一种浮点数数据类型,也被称为双精度浮点数,它可以表示的范围和精度都比半精度浮点数和单精度浮点数更高。全精度浮点数使用1个符号位、11个指数位和52个尾数位来表示一个实数,它可以表示的最大正数是2^1023-2^970 ≈ 1.80×10^308,最小正数是2^-1022 ≈ 2.23×10^-308,可以表示的精度为2^-52 ≈ 2.22×10^-16。

  全精度浮点数在深度学习中也被广泛使用,特别是在需要高精度计算的场景下,例如语音识别、自然语言处理等任务。在深度学习模型中,全精度浮点数通常被用来存储梯度和模型参数,因为它能够提供较高的计算精度和稳定性,避免出现数值溢出或下溢的问题。

混合精度

  全精度浮点数的缺点是它的计算和存储成本较高,需要更多的计算资源和内存空间。在一些场景下,例如大规模的模型训练和部署,全精度浮点数可能会成为性能瓶颈。因此,为了提高计算和存储效率,半精度浮点数和混合精度计算等技术被广泛应用于深度学习中。

综合比较

  1. 表示范围和精度:全精度浮点数比半精度浮点数能够表示更大范围和更高精度的实数,因为它使用更多的位来表示指数和尾数。全精度浮点数可以表示的最大正数和最小正数分别是半精度浮点数的10^4和10^-4倍左右,可以表示的精度是半精度浮点数的10^-7倍左右。
  2. 计算速度:半精度浮点数在计算速度上比全精度浮点数更快,因为它的位数更少,可以使用更简单和更快的运算方式。在一些计算密集型的深度学习任务中,使用半精度浮点数可以提高计算速度和效率。
  3. 存储空间:半精度浮点数比全精度浮点数使用更少的存储空间,因为它的位数更少。在大规模的深度学习任务中,使用半精度浮点数可以节省存储空间,从而减少模型训练和部署的成本和复杂度。
  4. 数值稳定性:半精度浮点数在表示范围和精度上的限制可能会导致数值溢出或下溢的问题,从而降低模型的计算精度和稳定性。全精度浮点数则更能够保证计算精度和稳定性,在一些需要高精度计算的任务中表现更好。

实操

  在PyTorch中,可以使用torch.cuda.amp模块来实现自动混合精度运算。以下是一个使用自动混合精度运算训练一个神经网络的例子:

ini

复制代码

import torch
from torch import nn, optim
from torch.cuda.amp import GradScaler, autocast
# 定义神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
    def forward(self, x):
        x = x.view(-1, 784)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.softmax(self.fc3(x))
        return x
# 加载数据
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize((0.1307,), (0.3081,))
                               ])),
    batch_size=64, shuffle=True)
# 创建模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
# 创建GradScaler对象
scaler = GradScaler()
# 训练模型
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.cuda(), target.cuda()
        
        # 自动混合精度计算
        with autocast():
            output = model(data)
            loss = nn.functional.cross_entropy(output, target)
        
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))


相关实践学习
在云上部署ChatGLM2-6B大模型(GPU版)
ChatGLM2-6B是由智谱AI及清华KEG实验室于2023年6月发布的中英双语对话开源大模型。通过本实验,可以学习如何配置AIGC开发环境,如何部署ChatGLM2-6B大模型。
相关文章
|
7月前
|
机器学习/深度学习 算法 PyTorch
125_训练加速:FlashAttention集成 - 推导注意力优化的独特内存节省
2025年,大型语言模型的训练面临着前所未有的挑战。随着模型参数量和序列长度的不断增加,传统注意力机制的内存瓶颈问题日益突出。FlashAttention作为一种突破性的注意力算法,通过创新的内存访问模式和计算优化,显著提升了训练效率和内存利用。
750 3
|
7月前
|
存储 机器学习/深度学习 PyTorch
119_LLM训练的高效内存管理与优化技术:从ZeRO到Flash Attention
大型语言模型(LLM)的训练面临着前所未有的计算和内存挑战。随着模型规模达到数百亿甚至数千亿参数,高效的内存管理成为训练成功的关键因素之一。2025年,LLM训练的内存优化技术已经取得了显著进展,从ZeRO优化器到Flash Attention等创新技术,为训练超大规模模型提供了可能。
754 159
|
机器学习/深度学习 存储 算法
NoProp:无需反向传播,基于去噪原理的非全局梯度传播神经网络训练,可大幅降低内存消耗
反向传播算法虽是深度学习基石,但面临内存消耗大和并行扩展受限的问题。近期,牛津大学等机构提出NoProp方法,通过扩散模型概念,将训练重塑为分层去噪任务,无需全局前向或反向传播。NoProp包含三种变体(DT、CT、FM),具备低内存占用与高效训练优势,在CIFAR-10等数据集上达到与传统方法相当的性能。其层间解耦特性支持分布式并行训练,为无梯度深度学习提供了新方向。
726 1
NoProp:无需反向传播,基于去噪原理的非全局梯度传播神经网络训练,可大幅降低内存消耗
|
10月前
|
缓存 固态存储 Windows
如何让内存发挥到最大效能?全面优化指南,提升电脑运行体验
电脑内存使用不合理会导致卡顿,本文教你如何优化内存性能。检查内存容量与主板支持上限,考虑升级或调整配置;关闭后台程序、管理浏览器标签、结束异常进程以释放内存;设置虚拟内存、调整视觉效果、定期重启提升效率;必要时增加内存条、选择高频内存、更换固态硬盘。避免盲目清理内存和依赖大内存忽视其他硬件瓶颈。只需合理设置,无需额外花钱,就能显著提升电脑速度。
|
缓存 并行计算 PyTorch
PyTorch CUDA内存管理优化:深度理解GPU资源分配与缓存机制
本文深入探讨了PyTorch中GPU内存管理的核心机制,特别是CUDA缓存分配器的作用与优化策略。文章分析了常见的“CUDA out of memory”问题及其成因,并通过实际案例(如Llama 1B模型训练)展示了内存分配模式。PyTorch的缓存分配器通过内存池化、延迟释放和碎片化优化等技术,显著提升了内存使用效率,减少了系统调用开销。此外,文章还介绍了高级优化方法,包括混合精度训练、梯度检查点技术及自定义内存分配器配置。这些策略有助于开发者在有限硬件资源下实现更高性能的深度学习模型训练与推理。
2251 0
|
机器学习/深度学习 并行计算 PyTorch
英伟达新一代GPU架构(50系列显卡)PyTorch兼容性解决方案
本文记录了在RTX 5070 Ti上运行PyTorch时遇到的CUDA兼容性问题,分析其根源为预编译二进制文件不支持sm_120架构,并提出解决方案:使用PyTorch Nightly版本、更新CUDA工具包至12.8。通过清理环境并安装支持新架构的组件,成功解决兼容性问题。文章总结了深度学习环境中硬件与框架兼容性的关键策略,强调Nightly构建版本和环境一致性的重要性,为开发者提供参考。
8475 64
英伟达新一代GPU架构(50系列显卡)PyTorch兼容性解决方案
|
9月前
|
编解码 Ubuntu Linux
ubuntu系统安装指南:免费且适合老旧电脑,4GB内存也能流畅运行!
点击启动台,找到并点击设置。在设置中,选择语言和区域,再点击管理语言。安装所需的语言包,输入密码进行确认。等待大约2分钟,语言包安装完成后,点击安装语言,选择中文选项。这里有简体和繁体两种选择,根据个人需求进行选择。再次等待2分钟,安装完成后,点击这里,选择中文并应用。然后,将出现的中文拖动到最上面,应用更改并退出设置。最后,重启虚拟机,再次进入系统时,你会发现界面已经变成了中文,而且系统依然保持流畅。Ubuntu系统不仅外观漂亮、干净,而且性能稳定、安全可靠。如果你的电脑内存只有4GB,或者你对Windows系统感到厌倦,那么Ubuntu绝对是一个值得尝试的选择。它不仅办公打印一应俱全,还拥
|
并行计算 PyTorch 算法框架/工具
融合AMD与NVIDIA GPU集群的MLOps:异构计算环境中的分布式训练架构实践
本文探讨了如何通过技术手段混合使用AMD与NVIDIA GPU集群以支持PyTorch分布式训练。面对CUDA与ROCm框架互操作性不足的问题,文章提出利用UCC和UCX等统一通信框架实现高效数据传输,并在异构Kubernetes集群中部署任务。通过解决轻度与强度异构环境下的挑战,如计算能力不平衡、内存容量差异及通信性能优化,文章展示了如何无需重构代码即可充分利用异构硬件资源。尽管存在RDMA验证不足、通信性能次优等局限性,但该方案为最大化GPU资源利用率、降低供应商锁定提供了可行路径。源代码已公开,供读者参考实践。
1264 3
融合AMD与NVIDIA GPU集群的MLOps:异构计算环境中的分布式训练架构实践
|
机器学习/深度学习 人工智能 物联网
MiniMind:2小时训练出你的专属AI!开源轻量级语言模型,个人GPU轻松搞定
MiniMind 是一个开源的超小型语言模型项目,帮助开发者以极低成本从零开始训练自己的语言模型,最小版本仅需25.8M参数,适合在普通个人GPU上快速训练。
2704 10
MiniMind:2小时训练出你的专属AI!开源轻量级语言模型,个人GPU轻松搞定
|
人工智能 负载均衡 调度
COMET:字节跳动开源MoE训练加速神器,单层1.96倍性能提升,节省百万GPU小时
COMET是字节跳动推出的针对Mixture-of-Experts(MoE)模型的优化系统,通过细粒度的计算-通信重叠技术,显著提升分布式训练效率,支持多种并行策略和大规模集群部署。
852 9

热门文章

最新文章