PyTorch内存优化的10种策略总结:在有限资源环境下高效训练模型

本文涉及的产品
智能开放搜索 OpenSearch行业算法版,1GB 20LCU 1个月
实时计算 Flink 版,1000CU*H 3个月
实时数仓Hologres,5000CU*H 100GB 3个月
简介: 在大规模深度学习模型训练中,GPU内存容量常成为瓶颈,特别是在训练大型语言模型和视觉Transformer时。本文系统介绍了多种内存优化策略,包括混合精度训练、低精度训练(如BF16)、梯度检查点、梯度累积、张量分片与分布式训练、

在大规模深度学习模型训练过程中,GPU内存容量往往成为制约因素,尤其是在训练大型语言模型(LLM)和视觉Transformer等现代架构时。由于大多数研究者和开发者无法使用配备海量GPU内存的高端计算集群,因此掌握有效的内存优化技术变得尤为关键。本文将系统性地介绍多种内存优化策略,这些技术组合应用可使模型训练的内存消耗降低近20倍,同时不会损害模型性能和预测准确率。以下大部分技术可以相互结合,以获得更显著的内存效率提升。

1、自动混合精度训练

混合精度训练是降低内存占用的基础且高效的方法,它充分利用16位(FP16)和32位(FP32)浮点格式的优势。

混合精度训练的核心思想是在大部分计算中使用较低精度执行数学运算,从而减少内存带宽和存储需求,同时在计算的关键环节保持必要的精度。通过对激活值和梯度采用FP16格式,这些张量的内存占用可减少约50%。然而某些特定的层或操作仍需要FP32格式以避免数值不稳定。PyTorch对自动混合精度(AMP)的原生支持大大简化了实现过程。

混合精度训练低精度训练 有本质区别

关于混合精度训练是否会影响模型准确率的问题,答案是。混合精度训练通过精心设计的计算流程保持了计算精度。

混合精度训练原理

混合精度训练通过结合16位(

FP16

)和32位(

FP32

)浮点格式来保持计算准确性。使用16位精度计算梯度可显著加快计算速度并减少内存消耗,同时维持与32位分辨率相当的结果质量。这种方法在计算资源受限的环境中尤为有效。

"混合精度"一词更准确地描述了这一过程,因为并非所有参数和操作都转换为16位格式。实际训练过程中,32位和16位操作交替执行,形成混合精度计算流程。

如上图所示,该过程首先将权重转换为低精度(

FP16

)以加速计算,然后计算梯度,接着将梯度转回高精度(

FP32

)以确保数值稳定性,最后使用这些适当缩放的梯度更新原始权重。通过这种方式,混合精度训练可提高训练效率的同时保持网络的整体精度和稳定性。

使用

torch.cuda.amp.autocast()

可轻松实现混合精度训练,示例代码如下:

 import torch
from torch.cuda.amp import autocast, GradScaler

# Assume your model and optimizer have been defined elsewhere.
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()

for data, target in data_loader:
    optimizer.zero_grad()
    # Enable mixed precision
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)

    # Scale the loss and backpropagate
    scaler.scale(loss).backward()
    scaler.step(optimizer)
     scaler.update()

2、低精度训练

除了混合精度训练,我们还可以尝试使用完整的16位低精度格式进行训练。由于16位浮点数的表示范围限制,这种方法可能导致

NaN

值出现。为解决这一问题,研究人员开发了多种专用浮点格式。其中,Brain Floating Point(

BF16

)是Google为此专门开发的一种广受欢迎的格式。与标准

FP16

相比,

BF16

提供了更大的动态范围,能够表示极大和极小的数值,使其更适合于深度学习应用中可能遇到的多样化数值情况。尽管较低精度可能在某些计算中影响精确度或导致舍入误差,但在大多数深度学习应用场景中,这种影响对模型性能的影响极小。

虽然

BF16

最初是为TPU设计的,但现在大多数现代GPU(Nvidia Ampere架构及更高版本)也支持这种格式。可以通过以下方法检查GPU是否支持

BF16

格式:

 import torch
 print(torch.cuda.is_bf16_supported())  # should print True

3、梯度检查点

即便采用混合精度和低精度训练,大型模型在前向传播过程中产生的大量中间张量仍会消耗大量内存。梯度检查点(Gradient Checkpointing)技术通过在前向传播过程中选择性地仅存储部分中间结果来解决这一问题。在反向传播过程中,系统会重新计算缺失的中间值,这虽然增加了计算成本,但可以显著降低内存需求。

通过战略性地选择需要设置检查点的层,可以通过动态重新计算激活值而非存储它们来减少内存使用。对于具有深层架构的模型,中间激活值通常占据了内存消耗的主要部分,此时这种权衡尤为有效。以下是梯度检查点的实现示例:

 import torch
from torch.utils.checkpoint import checkpoint

def checkpointed_segment(input_tensor):
    # This function represents a portion of your model
    # which will be recomputed during the backward pass.
    # You can create a custom forward pass for this segment.
    return model_segment(input_tensor)

# Instead of a conventional forward pass, wrap the segment with checkpoint.
 output = checkpoint(checkpointed_segment, input_tensor)

采用此方法,在多数情况下可将激活值所需的内存减少40-50%。尽管反向传播现在包含额外的计算开销,但当GPU内存成为限制因素时,这种权衡通常是合理的。

4、使用梯度累积降低批量大小

在尝试上述方法后,一个自然的问题是:

为何不直接减小批量大小?

虽然这确实是最直接的方法,但通常使用较小批量大小会导致预测性能下降。简单减小批量大小虽然能显著降低内存消耗,但往往会对模型准确率产生不良影响。

如何在这两者之间取得平衡?

梯度累积(Gradient Accumulation)正是为解决这一问题而设计的技术。它允许在训练过程中虚拟增加批量大小,其核心原理是为较小的批量计算梯度,并在多次迭代中累积这些梯度(通常通过求和或平均),而不是在每个批次后立即更新模型权重。一旦累积的梯度达到目标"虚拟"批量大小,才使用这些累积的梯度更新模型参数。

然而需要注意,这种技术的主要缺点是显著增加了训练时间。

5、张量分片和分布式训练

对于即使应用上述优化后仍无法在单个GPU上容纳的超大模型,完全分片数据并行(Fully Sharded Data Parallel, FSDP)技术提供了解决方案。FSDP将模型参数、梯度和优化器状态分片到多个GPU上,这不仅使得训练超大模型成为可能,还能通过更合理地分配通信开销提高训练效率。

FSDP不是在每个GPU上维护完整的模型副本,而是将模型参数分配到多个可用设备上。在执行前向或反向传播时,系统仅将相关分片加载到内存中。这种分片机制显著降低了单个设备的内存需求,与前述技术结合使用,在某些情况下可实现高达10倍的内存降低效果。

FSDP可通过以下方式实现:

 import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

# Initialize your model and ensure it is on the correct device.
model = MyLargeModel().cuda()

# Wrap the model in FSDP for sharded training across GPUs.
 fsdp_model = FSDP(model)

6、高效的数据加载

内存优化中常被忽视的一个方面是数据加载效率。虽然大部分优化关注点集中在模型内部结构和计算过程,但低效的数据处理同样可能造成不必要的瓶颈,影响内存利用和计算速度。作为经验法则,当处理数据加载器时,应始终启用

Pinned Memory

和配置适当的

Multiple Workers

,如下所示:

 from torch.utils.data import DataLoader

# Create your dataset instance and then the DataLoader with pinned memory enabled.
train_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,      # Adjust based on your CPU capabilities
    pin_memory=True     # Enables faster host-to-device transfers
 )

7、使用原地操作

在处理张量时,如果不谨慎管理,每个操作都可能创建新的张量对象。原地操作(In-place Operations)通过直接修改现有张量而非分配新张量,有助于减少内存碎片和总体内存占用。这种方式减少了临时内存分配,在迭代训练循环中尤为重要。示例如下:

 import torch

 x = torch.randn(100, 100, device='cuda')
 y = torch.randn(100, 100, device='cuda')

 # Using in-place addition
 x.add_(y)  # Here x is modified directly instead of creating a new tensor

8、激活和参数卸载

对于极大规模模型,即使应用了所有上述技术,由于大量中间激活值的存在,仍可能达到GPU内存限制。激活和参数卸载(Activation and Parameter Offloading)技术通过将部分中间数据移动到CPU内存,为GPU内存提供额外的缓解。

这种方法通过战略性地将部分激活值和/或参数临时卸载到主机内存(CPU),仅在GPU内存中保留关键计算所需的数据。虽然DeepSpeed、Fabric等专用框架可自动管理这种数据移动,但也可以按如下方式实现自定义卸载逻辑:

 def offload_activation(tensor):
    # Move tensor to CPU to save GPU memory
    return tensor.cpu()

def process_batch(data):
    # Offload some activations explicitly
    intermediate = model.layer1(data)
    intermediate = offload_activation(intermediate)
    intermediate = intermediate.cuda()  # Move back when needed
    output = model.layer2(intermediate)
     return output

9、使用更精简的优化器

各种优化器在内存消耗方面存在显著差异。例如,广泛使用的Adam优化器为每个模型参数维护两个额外状态参数(动量和方差),这意味着更多的内存消耗。将Adam替换为无状态优化器(如SGD)可将参数数量减少近2/3,这在处理LLM等大型模型时尤为重要。

标准SGD的缺点是收敛特性较差。为弥补这一点,可引入余弦退火学习率调度器以实现更好的收敛效果。实现示例:

 # instead of this
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)

# use this
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
num_steps = NUM_EPOCHS * len(train_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
             optimizer, T_max=num_steps)

这种优化可在保持模型准确率达到约97%(取决于具体应用)的同时,显著改善峰值内存消耗。

10、进阶优化技术

除上述基础技术外,以下高级策略可进一步优化GPU内存使用,充分发挥硬件潜能:

内存分析和缓存管理

精确测量是有效优化的前提。PyTorch提供了多种实用工具用于监控GPU内存使用情况:

 import torch

 # print a detailed report of current GPU memory usage and fragmentation
 print(torch.cuda.memory_summary(device=None, abbreviated=False))

 # free up cached memory that's no longer needed by PyTorch
 torch.cuda.empty_cache()

使用TorchScript进行JIT编译

PyTorch的即时编译器(JIT)能够将Python模型转换为经过优化的、可序列化的TorchScript程序。这种转换通过优化内核启动和减少运行时开销,可带来内存和性能的双重提升:

 import torch

 # Suppose `model` is an instance of your PyTorch network.
 scripted_model = torch.jit.script(model)

 # Now, you can run the scripted model just like before.
 output = scripted_model(input_tensor)

这种编译方式可显著优化模型运行效率。

自定义内核融合

编译的另一项重要优势是能够将多个操作融合到单个计算内核中。内核融合有助于减少内存读写操作,提高总体计算吞吐量:

使用torch.compile()进行动态内存分配

进一步利用编译技术,JIT编译器可通过编译时优化改进动态内存分配效率。结合跟踪和计算图优化技术,这种方法可在大型模型和Transformer架构中实现更显著的内存和性能优化。

总结

在GPU和云计算资源成本高昂的环境下,最大化利用现有计算资源至关重要。对于希望在有限计算资源条件下训练或微调大型模型(如LLM或视觉Transformer)的研究者和开发者而言,掌握上述优化技术尤为重要。本文介绍的这些策略代表了研究人员和专业人士在资源受限条件下进行高效模型训练的常用方法。

https://avoid.overfit.cn/post/dc61dc9f03fc45f48dba26c21a276bce

作者:Sahib Dhanjal

目录
相关文章
|
1月前
|
机器学习/深度学习 算法 数据可视化
近端策略优化算法PPO的核心概念和PyTorch实现详解
本文深入解析了近端策略优化(PPO)算法的核心原理,并基于PyTorch框架实现了完整的强化学习训练流程。通过Lunar Lander环境展示了算法的全过程,涵盖环境交互、优势函数计算、策略更新等关键模块。内容理论与实践结合,适合希望掌握PPO算法及其实现的读者。
280 2
近端策略优化算法PPO的核心概念和PyTorch实现详解
|
6月前
|
存储 分布式计算 监控
阿里云服务器实例经济型e、通用算力型u1、计算型c8i、通用型g8i、内存型r8i详解与选择策略
在阿里云现在的活动中,可选的云服务器实例规格主要有经济型e、通用算力型u1、计算型c8i、通用型g8i、内存型r8i实例,虽然阿里云在活动中提供了多种不同规格的云服务器实例,以满足不同用户和应用场景的需求。但是有的用户并不清楚他们的性能如何,应该如何选择。本文将详细介绍阿里云服务器中的经济型e、通用算力型u1、计算型c8i、通用型g8i、内存型r8i实例的性能、适用场景及选择参考,帮助用户根据自身需求做出更加精准的选择。
|
2月前
|
存储 人工智能 自然语言处理
AI代理内存消耗过大?9种优化策略对比分析
在AI代理系统中,多代理协作虽能提升整体准确性,但真正决定性能的关键因素之一是**内存管理**。随着对话深度和长度的增加,内存消耗呈指数级增长,主要源于历史上下文、工具调用记录、数据库查询结果等组件的持续积累。本文深入探讨了从基础到高级的九种内存优化技术,涵盖顺序存储、滑动窗口、摘要型内存、基于检索的系统、内存增强变换器、分层优化、图形化记忆网络、压缩整合策略以及类操作系统内存管理。通过统一框架下的代码实现与性能评估,分析了每种技术的适用场景与局限性,为构建高效、可扩展的AI代理系统提供了系统性的优化路径和技术参考。
154 4
AI代理内存消耗过大?9种优化策略对比分析
|
5月前
|
缓存 并行计算 PyTorch
PyTorch CUDA内存管理优化:深度理解GPU资源分配与缓存机制
本文深入探讨了PyTorch中GPU内存管理的核心机制,特别是CUDA缓存分配器的作用与优化策略。文章分析了常见的“CUDA out of memory”问题及其成因,并通过实际案例(如Llama 1B模型训练)展示了内存分配模式。PyTorch的缓存分配器通过内存池化、延迟释放和碎片化优化等技术,显著提升了内存使用效率,减少了系统调用开销。此外,文章还介绍了高级优化方法,包括混合精度训练、梯度检查点技术及自定义内存分配器配置。这些策略有助于开发者在有限硬件资源下实现更高性能的深度学习模型训练与推理。
1001 0
|
1月前
|
机器学习/深度学习 监控 安全
解密虚拟化弹性内存:五大核心技术与实施策略
本文深入解析虚拟化环境中实现内存弹性管理的五大核心技术与实施策略。内容涵盖内存架构演进、关键技术原理、性能优化方法及典型问题解决方案,助力提升虚拟机密度与资源利用率。
107 0
|
1月前
|
边缘计算 算法 Java
Java 绿色计算与性能优化:从内存管理到能耗降低的全方位优化策略与实践技巧
本文探讨了Java绿色计算与性能优化的技术方案和应用实例。文章从JVM调优(包括垃圾回收器选择、内存管理和并发优化)、代码优化(数据结构选择、对象创建和I/O操作优化)等方面提出优化策略,并结合电商平台、社交平台和智能工厂的实际案例,展示了通过Java新特性提升性能、降低能耗的显著效果。最终指出,综合运用这些优化方法不仅能提高系统性能,还能实现绿色计算目标,为企业节省成本并符合环保要求。
80 0
|
弹性计算 安全 数据库
【转】云服务器虚拟化内存优化指南:提升性能的7个关键策略
作为云计算服务核心组件,虚拟化内存管理直接影响业务系统性能表现。本文详解了内存优化方案与技术实践,助您降低30%资源浪费。
82 0
【转】云服务器虚拟化内存优化指南:提升性能的7个关键策略
|
4月前
|
缓存 Linux 数据安全/隐私保护
Linux环境下如何通过手动调用drop_caches命令释放内存
总的来说,记录住“drop_caches” 命令并理解其含义,可以让你在日常使用Linux的过程中更加娴熟和自如。
1008 23
|
5月前
|
机器学习/深度学习 数据可视化 机器人
比扩散策略更高效的生成模型:流匹配的理论基础与Pytorch代码实现
扩散模型和流匹配是生成高分辨率数据(如图像和机器人轨迹)的先进技术。扩散模型通过逐步去噪生成数据,其代表应用Stable Diffusion已扩展至机器人学领域形成“扩散策略”。流匹配作为更通用的方法,通过学习时间依赖的速度场将噪声转化为目标分布,适用于图像生成和机器人轨迹生成,且通常以较少资源实现更快生成。 本文深入解析流匹配在图像生成中的应用,核心思想是将图像视为随机变量的实现,并通过速度场将源分布转换为目标分布。文中提供了一维模型训练实例,展示了如何用神经网络学习速度场,以及使用最大均值差异(MMD)改进训练效果。与扩散模型相比,流匹配结构简单,资源需求低,适合多模态分布生成。
358 13
比扩散策略更高效的生成模型:流匹配的理论基础与Pytorch代码实现
|
6月前
|
缓存 监控 算法
JVM简介—2.垃圾回收器和内存分配策略
本文介绍了Java垃圾回收机制的多个方面,包括垃圾回收概述、对象存活判断、引用类型介绍、垃圾收集算法、垃圾收集器设计、具体垃圾回收器详情、Stop The World现象、内存分配与回收策略、新生代配置演示、内存泄漏和溢出问题以及JDK提供的相关工具。
JVM简介—2.垃圾回收器和内存分配策略

推荐镜像

更多