微软手机做网站服务器吗,哈尔滨手机网页制作,wordpress 注册表单 自定义,抖音代运营费用明细目录 Out of MemoryGradient AccumulationGradient CheckpointingXformersDiffusers的显存优化 Out of Memory
先来说下OOM问题#xff0c;其实也是日常会遇到的情况。模型申请的显存超过了设备实际显存大小#xff0c;则会报错Out of Memory。一般情况下#xff0c;batch … 目录 Out of MemoryGradient AccumulationGradient CheckpointingXformersDiffusers的显存优化 Out of Memory
先来说下OOM问题其实也是日常会遇到的情况。模型申请的显存超过了设备实际显存大小则会报错Out of Memory。一般情况下batch size设置过大不能匹配自己手里的计算设备(GPU、TPU等)显存时会经常触发这个问题。
一般遇到这个情况都会选择调小batch size但是很多模型本身就非常大尤其是预训练模型当道的今天记得19年的时候拿一张1080ti做BERT finetune11G的显存batch size最大也就就能设置成4。但是batch size又是很影响训练效果的超参在很多时候只能在原作者调参得到的那个数值下才能训练出较好的结果。此时有钱就选择加卡不然就只能另辟蹊径来磨一磨手里这张小显存的计算卡了。
正常的训练流程每个step送入1个batch_size的数据计算梯度然后用这个梯度反向传播更新参数接着送入下一个batch_step的数据重复。
Gradient Accumulation
梯度累加 gradient_accumulation 顾名思义就是将多次计算得到的梯度值进行累加然后一次性反向传播进行参数更新。 如图所示假设我们有batch size 256的global-batch在单卡训练显存不足时将其分为多个小的mini-batch64如图分为大小为64的4个mini-batch每个step送入1个mini-batch获得梯度将多次获得的梯度进行累加后再更新参数以次达到模拟单次使用global-batch训练的目的。 自动微分机制和梯度累加实现 由于不同深度学习框架的自动微分机制不同所以实现梯度累加的方式有显著差异。当前业界框架的自动微分机制分为两类
以Tensor为核心的自动微分Tensor可配置requires_grad参数控制是否需要梯度。每个Tensor对象有grad_fn属性用来存储该Tensor参与的反向操作同时还有grad属性存储该Tensor对应的梯度。计算梯度时通过标量loss.backward()实现。由于API接口使用方式很符合反向传播的概念业界多数框架选择此方案如Pytorch、Paddle、Oneflow、MegEngine等。函数式自动微分。将神经网络正向计算视作输入到Loss的计算函数通过函数变换获得反向计算函数并调用反向计算函数来获得梯度。业界采用此方案的框架有Jax、MindSpore此外Tensorflow的GradientTape本质上也可以视作该方案的变种。
自动微分原理都是一致的这两种方案的核心差异点在于是否暴露了自动微分更底层的接口如Pytorch等框架更多定位纯深度学习此时只体现backward更符合目标用户的使用习惯。而Jax、MindSpore则在定位上更加底层Jax直接明言自身为数值计算框架MindSpore则定位为AI科学计算融合计算框架因此函数式自动微分设计更符合框架定位。
由于两种方案的差异也造成了梯度累加实现的差异下面主要以Pytorch为例讲一下梯度累加的实现
# batch accumulation parameter
accum_iter 4 # loop through enumaretad batches
for batch_idx, (inputs, labels) in enumerate(data_loader):# forward pass preds model(inputs)loss criterion(preds, labels)# scale the loss to the mean of the accumulated batch sizeloss loss / accum_iter # backward passloss.backward()# weights updateif ((batch_idx 1) % accum_iter 0) or (batch_idx 1 len(data_loader)):optimizer.step()optimizer.zero_grad()由于Pytorch本身在求完梯度后会自动挂载到Tensor.grad属性上而在没有手动清空即optimizer.zero_grad()的情况下每个step训练求得的梯度会自动累加因此只需要控制梯度清零和参数更新的间隔步数即可。
这里着重强调一个位置loss loss / accum_iter这一行操作的含义及实现。稍微翻看了一下搜索引擎排名靠前的几篇发现误导不少。
首先是注释里说明的含义不知出处但是几乎所有人都备注一句normalize loss to account for batch accumulation中文译作loss正则化。这个地方显然是越写越偏了。
实际上这里就是做了一次求mean的操作。原因是直接累加的accum_iter次梯度值作为一次参数更新的梯度是将梯度值放大了accum_iter倍而Pytorch的参数更新是写在optimizer.step()方法内部无法手动控制因此只能根据链式法则在loss处进行缩放来达到缩放梯度的目的。与常规理解的正则化没有任何关系。
此外还有一个谬误的写法 loss criterion(outputs, labels)loss loss / accumulation_stepsloss.backward()loss通常不会这样累加一般会单独维护一个 total_loss, 且累加之后再做loss.backward()微分结果也是错误的由左式正确的偏导变为loss累加和对w求偏导
Gradient Checkpointing
当我们采用了分布式训练以及混合精度训练都不能降低显存大小的时候比如多语言large模型光词表就有几十万现有的 GPU 资源无法训练一个设备装不下的模型。下面我们介绍一种改善这个问题的技术梯度检查点 gradient_checkpointing
简单的说梯度检查点 gradient_checkpointing 的工作原理是在反向传播时重新计算深度神经网络的中间值而通常情况是在前向传播时存储的)。 这个策略是用时间重新计算这些值两次的时间成本来换空间提前存储这些值的内存成本。
神经网络使用的总内存基本上是两个部分的总和
第一部分是模型使用的静态内存。尽管 PyTorch 模型中内置了一些固定开销但总的来说几乎完全由模型权重决定。而如今在生产中使用的现代深度学习模型的总参数在100万到10亿之间。作为参考一个带 16GB GPU 内存的 NVIDIA T4 的实际限制大约在1-1.5亿个参数之间。第二部分是模型的计算图所占用的动态内存。在训练模式下每次通过神经网络的前向传播都为网络中的每个神经元计算一个激活值这个值随后被存储在所谓的计算图中。必须为批次中的每个单个训练样本存储一个值因此数量会迅速的累积起来。总成本取决于模型大小和批处理大小并设置适用于您的GPU内存的最大批处理大小的限制。
模型的训练需要计算前向传播forward来获得预测结果然后通过反向传播backward来计算梯度以进行参数更新。
一开始前向传播forward计算结果时存储储激活值的原因是在反向传播backword期间计算梯度时需要用到激活值。
现在在forward后不存储激活值在backward需要激活值时重新计算。
梯度检查点gradient_checkpointing是如何起作用的
大型模型在静态和动态方面都很耗资源。首先它们很难适配 GPU而且哪怕你把它们放到了设备上也很难训练因为批处理大小被迫限制的太小而无法收敛。
梯度检查点gradient checkpointing的工作原理是从计算图中省略一些激活值由前向传播产生其中这里的”一些“是指可以只省略模型中的部分激活值折中时间和空间陈天奇在它的论文中Training Deep Nets with Sublinear Memory Cost使用了如下动图的方法即前向传播的时候存一个节点释放一个节点空的那个等需要用的时候再backword的时候重新计算。这减少了计算图使用的内存降低了总体内存压力并允许在处理过程中使用更大的批次大小。 PyTorch 通过torch.utils.checkpoint.checkpoint和torch.utils.checkpoint.checkpoint_sequential提供梯度检查点根据官方文档的 notes它实现了以下功能在前向传播时PyTorch 将保存模型中的每个函数的输入元组。在反向传播过程中对于每个函数输入元组和函数的组合以实时的方式重新计算插入到每个需要它的函数的梯度公式中然后丢弃显存中只保存输入数据和函数。网络计算开销大致相当于每个样本通过模型前向传播开销的两倍。
梯度检查点算法将模型的动态内存开销从 O ( N ) O(N) O(N) n为模型中的层数降低到 O ( N ) O(\sqrt{N}) O(N ) 并通过实验展示了将 ImageNet 的一个变种从 48GB 压缩到了 7GB 内存占用。
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.checkpoint import checkpointclass in_conv(nn.Module):def __init__(self, in_ch, out_ch):super(in_conv, self).__init__()self.op nn.Sequential(nn.Conv2d(in_ch, out_ch, kernel_size3, padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue),nn.Conv2d(out_ch, out_ch, kernel_size3, padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue))def forward(self, x):x self.op(x)return xclass conv3x3(nn.Module):def __init__(self, in_ch, out_ch):super(conv3x3, self).__init__()self.op nn.Sequential(nn.Conv2d(in_ch,out_ch, kernel_size3, padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue),nn.Conv2d(out_ch,out_ch, kernel_size3, padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue))def forward(self, x):x self.op(x)return xclass down_block(nn.Module):def __init__(self, in_ch, out_ch):super(down_block, self).__init__()self.pool nn.MaxPool2d(2, stride2)self.conv conv3x3(in_ch, out_ch)def forward(self, x):x self.pool(x)x self.conv(x)return xclass up_block(nn.Module):def __init__(self, in_ch, out_ch, residualFalse):super(up_block, self).__init__()self.up nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size2, stride2)self.conv conv3x3(in_ch, out_ch)def forward(self, x1, x2):x1 self.up(x1)diffY x2.size()[2] - x1.size()[2]diffX x2.size()[3] - x1.size()[3]x1 F.pad(x1, (diffX // 2, diffX - diffX // 2,diffY // 2, diffY - diffY // 2))x torch.cat([x2, x1], dim1)x self.conv(x)return xclass out_conv(nn.Module):def __init__(self, in_ch, out_ch):super(out_conv, self).__init__()self.conv nn.Conv2d(in_ch, out_ch, kernel_size1)def forward(self, x):x self.conv(x)return xclass UNet(nn.Module):def __init__(self, img_channels, n_classes,use_checkpointFalse):super(UNet, self).__init__()self.inc in_conv(img_channels,64)self.down1 down_block(64, 128)self.down2 down_block(128, 256)self.down3 down_block(256, 512)self.down4 down_block(512, 1024)self.up1 up_block(1024, 512)self.up2 up_block(512, 256)self.up3 up_block(256, 128)self.up4 up_block(128, 64)self.outc out_conv(64, 1)def forward(self, x):def forward(self, x):x Variable(x,requires_gradTrue) if self.use_checkpoint:x1 checkpoint(self.inc,x)x2 checkpoint(self.down1,x1)x3 checkpoint(self.down2,x2)x4 checkpoint(self.down3,x3)x5 checkpoint(self.down4,x4)x checkpoint(self.up1,x5,x4)x checkpoint(self.up2,x,x3)x checkpoint(self.up3,x,x2)x checkpoint(self.up4,x,x1)x checkpoint(self.outc,x)else:x1 self.inc(x)x2 self.down1(x1)x3 self.down2(x2)x4 self.down3(x3)x5 self.down4(x4)x self.up1(x5, x4)x self.up2(x, x3)x self.up3(x, x2)x self.up4(x, x1)x self.outc(x)return x注意第94行必须确保checkpoint的输入输出都声明为require_gradTrue的Variable否则运行时会报如下的错 RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
网络训练高效内存管理——torch.utils.checkpoint的使用
torch.utils.checkpoint 简介 和 简易使用
此外HuggingFace Transformers 也支持 Gradient Checkpoint。 梯度检查点可以通过 PreTrainedModel 实例的 gradient_checkpointing_enable 方法执行。
from transformers import AutoConfig, AutoModel# https://github.com/huggingface/transformers/issues/9919
from torch.utils.checkpoint import checkpoint# initializing model
model_path microsoft/deberta-v3-base
config AutoConfig.from_pretrained(model_path)
model AutoModel.from_pretrained(model_path, configconfig)# gradient checkpointing
if model.supports_gradient_checkpointing:model.gradient_checkpointing_enable()print(fGradient Checkpointing: {model.is_gradient_checkpointing})更进一步看看HF是如何在Bert模型内部实现梯度检查点的
modeling_bert.py中定义了BertEncoder其中forward函数内部判断是否开启了梯度检查点并调用torch.utils.checkpoint.checkpoint实现梯度检查点
class BertEncoder(nn.Module):def forward(args:...for i, layer_module in enumerate(self.layer):...if self.gradient_checkpointing and self.training:if use_cache:logger.warning(use_cacheTrue is incompatible with gradient checkpointing. Setting use_cacheFalse...)use_cache Falsedef create_custom_forward(module):def custom_forward(*inputs):return module(*inputs, past_key_value, output_attentions)return custom_forwardlayer_outputs torch.utils.checkpoint.checkpoint(create_custom_forward(layer_module),hidden_states,attention_mask,layer_head_mask,encoder_hidden_states,encoder_attention_mask,)else:layer_outputs layer_module(hidden_states,attention_mask,layer_head_mask,encoder_hidden_states,encoder_attention_mask,past_key_value,output_attentions,)...Xformers
xformers库仅适用于N卡用于推理和训练中在attention block中执行优化可以提高速度并减少内存消耗。在图像生成领域特点是加速图片生成并降低显存占用代价是输出图像不稳定有可能比不开Xformers略差。
self-attention在Transformer中发挥着重要的作用但在实际应用中也面临着两个挑战
1 complexityself-attention的时间复杂度是 O ( T 2 ⋅ D ) O(T^2·D) O(T2⋅D),在处理长序列问题上存在较大瓶颈。
2 structural priorself-attention没有对输入做输入做任何结构偏差的假设因此对于小数据集可能存在过拟合现象。
xformers在Vanilla Transformer普通transformer的self-attention基础上进行改进 Sparse attention 在计算attention matrix时不会attend 每个token而是仅对部分sparse connection的位置计算。根据确定sparse connection的方法又可以细分为 position-based 和 content-based 两种。 Position-based Sparse Attention Global Attention为了缓解稀疏注意力在中长距离依赖关系上模型能力的退化这一类模式增加了一些全局节点作为节点间信息传播的枢纽。这些全局节点可以关注序列中的所有节点并且整个序列也都会关注这些全局节点。Band Attention这一类attention 也可以成为局部attention或者滑动窗口attention设计的主要思路在于由于大部分的数据都有很强的局部关系特性因此可以限制query只关注附近的一些节点从而做到稀疏化的效果。Dilated Attention这种attention的方法与dilated CNN的方法十分相似可以在不增加计算复杂度的情况下扩大感受野的大小并且通过调整dilation 的间距可以扩展至strided attention.Random Attention为了提高非局部位置之间的联系每个query随机的选择一些位置去attend.Block Local Attention将整个序列划分为几个没有重叠部分的block每个block内部之间做attention. Content-based sparse attention ReformerReformer使用了Locality-sensitive hashing(LSH) 对每个query选择对应的key-value对。Route Transformer相比于reformer该方法采用了k-means聚类的方法对query和key进行聚类每个query 只attend 属于同一类cluster的keys. 聚类中心的向量根据被赋予向量的指数移动平均值来计算同时除以cluster中数目的移动平均值。 Linearized attention 这里可将原始的指数内积的计算形式替换为核函数从而提出新的计算方式 对于自回归的attention而言上式中的累积和项可以通过上一个时间步的结果计算叠加而来因此对于transformer的decoder来说整个计算过程类似于RNN的计算过程。 在Linear Transformer中linear map采用了一种简单的计算方法 这种feature map的目的不是为了近似内积attention但是通过实验证明它和标准的Transformer的结果表象相当。 Query Prototyping和memory Compression除了使用sparse attention或者核函数的linearized attention外另一个改进的思路就是减少query和key-value对的数量。 Query Prototyping如 Informer编码器接收大量的长序列输入绿色序列。采用ProbSparse attention 来代替常规的self-attention。蓝色梯形是自注意力的蒸馏操作来提取主要的attention大大减小了网络规模。层叠加复制副本的操作增加了鲁棒性。解码器接收长序列输入将target中元素填充为零根据特征图的加权注意力组成立即以生成方式预测输出元素橙色系列从而使得Transformer能够更有效的处理长序列数据。 Attention with Compressed Key-Value Memory相比于减少query数量的方法这一类方法的主要特点在于通过减少key-value对的数量来减少复杂度。 Memory Compressed Attention(MCA)通过跨步卷积的方法来减少key和value的数量减少的大小和kernel size k的数值有关这种方法相比于之前提到局部注意力而言增加了对全局上下文的捕捉。Linformer利用线性投影将键和值从长度n投射到一个更小的长度的nk。这也将self attention的复杂性降低到线性。这种方法的缺点是必须假定输入序列的长度因此不能用于自回归的问题上。PoolingFormer将原始的全注意力机制修改为一个两级注意力机制第一级采用滑动窗口注意力机制限制每个词只关注近距离的邻居第二级采用池化注意力机制采用更大的窗口来增加每个token 的感受野同时利用池化操作来压缩键和值向量以减少要参加注意力运算的token数量。 多头机制的改进多头注意力的一大优势在于它能够共同关注来自不同子空间的不同位置的信息。然而目前还没有一种机制可以保证不同的注意头确实捕捉了不同的特征。
除此之外还有很多对Transformer的改进都集成进了xformers。xFormers 软件包需要最新版本的 PyTorch。如果需要使用以前版本的 PyTorch那么我建议 从github souce 安装 xFormers。
xFormers安装后diffusers的unet可以使用enable_xformers_memory_efficient_attention() 为了更快的推理和减少内存消耗
if enable_xformers_memory_efficient_attention:if is_xformers_available():unet.enable_xformers_memory_efficient_attention()else:raise ValueError(xformers is not available. Make sure it is installed correctly)PyTorch 2.0相比1.12改变了很多底层实现,不向后兼容。所以不能在PyTorch 2.0上使用xFormers。但是可以使用这个机制GitHub - Dao-AILab/flash-attention: Fast and memory-efficient exact attention
更具体的直接使用xformers实现attention
import torch
from xformers.ops import memory_efficient_attention, LowerTriangularMaskdevicecuda
batch 4
n_head 8
head_dim 16
seq_len 128q torch.rand(batch, seq_len, n_head, head_dim).to(device)
k torch.rand(batch, seq_len, n_head, head_dim).to(device)
v torch.rand(batch, seq_len, n_head, head_dim).to(device)# 使用 causal 偏置掩码
o memory_efficient_attention(q, k, v, LowerTriangularMask())# 不使用任何偏置掩码
o memory_efficient_attention(q, k, v)memory_efficient_attention 的等效pytorch代码实现用于模型导出。注LowerTriangularMask 之类的掩码请另外手动生成一个等效的 attn_bias Tensor 再用于本函数
# 使用自定义的偏置掩码 attn_bias要求 xformers 版本 大于等于 0.17
## 这里的 from_lento_len 分别代表Decoder的序列长度Encoder的序列长度
from_len seq_len
to_len seq_len
attn_bias torch.rand(batch, n_head, from_len, to_len).to(device)
o memory_efficient_attention(q, k, v, attn_bias)import torch.nn.functional as Fdef memory_efficient_attention_pytorch(query, key, value, attn_biasNone, p0., scaleNone):# query [batch, seq_len, n_head, head_dim]# key [batch, seq_len, n_head, head_dim]# value [batch, seq_len, n_head, head_dim]# attn_bias [batch, n_head, seq_len, seq_len]if scale is None:scale 1 / query.shape[-1] ** 0.5 # BLHC - BHLCquery query.transpose(1, 2)key key.transpose(1, 2)value value.transpose(1, 2)# scale queryquery query * scale# BHLC BHCL - BHLLattn query key.transpose(-2, -1)if attn_bias is not None:attn attn attn_biasattn attn.softmax(-1)attn F.dropout(attn, p)# BHLL BHLC - BHLCout attn value# BHLC - BLHCout out.transpose(1, 2)return outDiffusers的显存优化 vae.enable_vae_slicing()将按顺序对每个prompt对应的隐变量进行解码,而不是同时解码整个batch。可以大大减少GPU内存使用,支持更大的batch size。需要 trade-off 推理时间和显存需求。 vae.enable_vae_tiling()在有限VRAM的情况下处理大尺寸图像。例如,在8GB VRAM中生成4K图像。分块VAE解码器将图像分割成重叠的块,对每个块解码,最后将输出混合生成最终图像。输出图像由于块解码器独立,会有一些块间色调变化,但不会出现明显的块间裂缝。512x512或以下大小的图像不会进行分块。分块VAE适用于大尺寸图像的生成,可以突破VRAM限制。需要 trade-off patch的调色差异和显存需求。 pipeline.enable_sequential_cpu_offload()为进一步节省内存底层使用accelerate可以将权重offload到CPU,仅在执行前向传播时加载到GPU。注意,此方法在子模块级别工作,而不是整个模型级别如每次卸载Unet的一个Conv层。这是最小化内存消耗的最佳方式,但由于过程的迭代特性,推理速度会大大降低。管道的UNet组件会运行多次(和num_inference_steps一样多次);每次UNet的不同子模块会按需要顺序地在CPU和GPU之间加载和卸载,所以内存传输的次数非常大。 pipeline.enable_model_cpu_offload()顺序CPU卸载可以节省大量内存,但会降低推理速度,因为子模块在需要时移动到GPU,在新模块运行时立即返回CPU。全model卸载是sequential卸载另一种替代方案它按整个模型而不是每个模型的组成模块来移动到GPU如每次卸载整个Unet。这对推理时间几乎没有影响在这种场景下,pipeline的主要组件(通常是:文本编码器、UNet和VAE)中只有一个处于GPU,其他组件等在CPU。像UNet这样运行多次迭代的组件会一直留在GPU,直到不再需要为止。
import torch
from diffusers import StableDiffusionPipelinepipe StableDiffusionPipeline.from_pretrained(/data1/huggingface/StableDiffusion/stable-diffusion-v1-5, torch_dtypetorch.float16
)
pipe pipe.to(cuda)
pipe.enable_model_cpu_offload() # new diffusers
pipe.unet.enable_xformers_memory_efficient_attention()
pipe.enable_vae_slicing()
pipe.enable_vae_tiling() # new diffusers
pipe.enable_attention_slicing()prompt a photo of an astronaut riding a horse on mars
images pipe([prompt] * 32).images # sample 32 images
