背景
这个问题在当下甚至之后的一段时间应该都是比较常见的,毕竟数据量、模型参数量增加的速度远超过算力提升速度,再者不是每个组织或者个人都拥有那么多的显卡。
本文整理了一些方法,希望能对那些和我一样用不起H100,A100这种高度卡的道友有点帮助。
术语
浮点数
浮点数种类挺多的,主要是range和precision的bit配置差异会影响到模型的性能。
memory wall
最初由 William Wulf 和 Sally Mckee 在 1995 年提出 [25]。“内存墙”问题既涉及内存容量的限制,又涉及内存传输带宽的限制。这需要考虑不同层级的内存数据传输。例如,计算逻辑和片上内存之间的数据传输、计算逻辑和 DRAM 内存之间的数据传输,或跨越不同插槽(socket)上不同处理器的数据传输。在所有这些情况下,数据传输容量和速度都远远落后于硬件 (HW) 的计算能力。
具体而言,旗舰级大型语言模型的规模每两年增长 410 倍(见上图)。类似地,大型推荐系统模型的参数规模也已经达到 O(10) TB 级别。相比之下,加速器 DRAM 内存的扩展速度仅为每两年 2 倍。
早期CV模型,acitivation是主要的显存占用,到了LLM,Adam成了主流优化器以后,Optimizer成了主要的内存占用,如上图所示。
PF-days
PF-day = 10^15 × 24 × 3600 = 8.64 × 10^19 floating point operations.
Bcrit
An Empirical Model of Large-Batch Traininghttps://arxiv.org/pdf/1812.06162.pdf
import numpy as np
def compute_noise_scale(gradients):
"""
计算梯度噪声尺度
Args:
gradients: 梯度列表
Returns:
噪声尺度
"""
# 计算梯度估计的协方差
covariance = np.cov(gradients)
# 计算噪声尺度
noise_scale = np.sqrt(np.trace(covariance))
return noise_scale
def compute_optimal_batchsize(noise_scale, learning_rate):
"""
计算最优batchsize
Args:
noise_scale: 噪声尺度
learning_rate: 学习率
Returns:
最优batchsize
"""
# 计算临界batchsize
b_crit = noise_scale / learning_rate
# 选择最优batchsize
if b_crit < 1:
batch_size = 1
else:
batch_size = int(np.ceil(b_crit))
return batch_size
def main():
# 训练模型
# ...
# 收集梯度
gradients = []
for epoch in range(num_epochs):
# ...
gradients.append(model.gradient)
# 计算噪声尺度
noise_scale = compute_noise_scale(gradients)
# 计算最优batchsize
learning_rate = 0.01
batch_size = compute_optimal_batchsize(noise_scale, learning_rate)
print("最优batchsize:", batch_size)
if __name__ == "__main__":
main()
也有经验的结合搜索把batchsize逐步加大,但总的结论就是:
小的batchsize,梯度方差大,sample efficiency更高,在budge充足的情况下,计算资源利用率低,耗时多;
大的batchsize,梯度方差小,稳定性高,sample efficiency低,在budge充足的情况下,计算资源利用率高,前期收敛速度快,为了达到一致的效果,耗时前小后大;
scaling law
这个是指导大模型设计和训练的最核心的部分了吧?相对来说也比较抽象,直接上结论:
Performance depends strongly on scale, weakly on model shape:
性能强依赖于scale,弱依赖于模型结构:scale由三个要素组成:模型参数量 N(不包含嵌入层),数据集大小 D,以及训练所使用的计算量 C。在合理范围内,性能与其他架构超参数(例如深度 vs. 宽度)的依赖性非常弱。
Smooth power laws:
当不受到其他两个因素限制时,性能与三个扩展因子 N、D、C 呈幂律关系,趋势跨越六个以上的数量级。
Universality of overfitting:
只要我们同时扩展 N 和 D,性能就会可预测地提升。但是,如果固定其中一个因素 (N 或 D) 仅增大另一个因素,性能提升的回报率会逐渐减少。性能损失可预测地取决于 N^0.74/D 的比率,意味着每当我们将模型大小增加 8 倍时,我们只需要将数据量增加大约 5 倍就可以避免性能损失。
Universality of training:
训练曲线遵循可预测的幂律,曲线参数与模型大小大致无关。通过外推训练曲线的早期部分,我们可以大致预测如果训练更长时间将达到的损失值。
Transfer improves with test performance:
当我们用与训练数据分布不同的文本评估模型时,结果与验证集上的结果密切相关,并且损失值会存在一个大致恒定的偏移量。换句话说,转移到不同的分布会带来一个恒定的损失惩罚,但除此之外,模型的性能会与训练集上的性能大致保持一致的提升趋势。
Sample efficiency:
大型模型比小型模型更具采样(拟合)效率,用更少的优化步骤和更少的数据量就能达到相同性能水平。
Convergence is inefficient:
在固定计算预算 C ,而模型大小 N 和可用数据 D 没有其他限制的情况下,训练一个非常大的模型并在远未收敛时停止训练,我们就可以获得最佳性能。因此,最大化计算效率训练将比训练小型模型收敛要高效得多,当数据量随着计算呈 D ∼ C^0.27 的速率非常缓慢地增长。
Optimal batch size:
理想的batchsize仅与损失函数的幂次方大致成正比,并且可以通过测量梯度噪声尺度来确定。
其实大家更想知道,如何快速判断模型是否遵循scaling law,以及如何设计一个遵循scaling law的模型。目前来看,大部分都是从N、D、C三方面出发设计实验进行归纳总结。
技术
混合精度训练amp
大部分人都用过,主要涉及到的就是BF16、FP16、TF32的选择问题。
优先选择BF16和TF32。但是,目前至少部分显卡支持这两种类型(Hopper,Ampere)。
import torch
//是否支持tf32
torch.backends.cuda.matmul.allow_tf32
//是否允许tf32,在PyTorch1.12及更高版本中默认为False
torch.backends.cudnn.allow_tf32
//是否支持bf16
import transformers
transformers.utils.import_utils.is_torch_bf16_gpu_available()
另外还有个FP8 (仅支持NVIDIA Hopper and NVIDIA Ada GPUs):
bitsandbytes
关于低精度量化相关的库
比如:
bitsandbytes.optim.AdamW8bit
bitsandbytes.nn.Linear8bitLt
梯度累积
最常见的方法,经典的yolov1就自带了,当一个batch显存放不下时适用。时间换空间。
流程就是把一个batch分成多个minibatch。逐个对minibatch进行bp,把梯度累积起来,然后进行参数更新。
一份伪代码,假设batch=16x32=512, minibatch=32:
# Initialize parameters and optimizer
theta = initialize_parameters()
optimizer = create_optimizer(theta)
# Set accumulation steps and batch size
accumulation_steps = 16 # Accumulate gradients for 16 steps
batch_size = 32 # Process 32 samples per batch
# Training loop
for epoch in range(num_epochs):
for step in range(num_steps_per_epoch):
# Get the next batch of data
data, labels = get_next_batch(batch_size)
# Forward pass
outputs = model(data)
loss = compute_loss(outputs, labels)
# Backward pass (accumulated gradients)
loss.backward()
if (step + 1) % accumulation_steps == 0:
# Update parameters using accumulated gradients
optimizer.step()
# Reset accumulated gradients
optimizer.zero_grad()
# Check and save checkpoints
if (step + 1) % checkpoint_interval == 0:
save_checkpoint(epoch, step, model)
# Final checkpoint
save_checkpoint(epoch, num_steps_per_epoch, model)
Gradient Checkpointing & Activation Checkpointing
这个技术简单说就是:减少缓存,需要时才计算。
BP在更新梯度的过程中,需要用到激活值和梯度,传统的流程就是会在forward保存激活值,backward逐层计算梯度:
传统的计算流程:
使用Checkpointing:
Training Deep Nets with Sublinear Memory Costhttps://arxiv.org/pdf/1604.06174
Parallel Training
目的是解决:
- 充分利用多路计算资源,缩减训练时间
- 显存资源不够导致的无法训练大batch、以及无法加载模型。
并行技术主要分为DP和MP,其中MP又细分为TP和PP:
- Data Parallel
- Model Parallel
- Tensor Parallel
- Pipeline Parallel
通常上面三种可以结合起来用DP+TP+PP。
更详细生动的讲解可以参考下面链接:
Model Parallelismhttps://huggingface.co/docs/transformers/v4.17.0/en/parallelismGPipe论文精读【论文精读】https://www.bilibili.com/video/BV1v34y1E7zu/
ZeRO(Zero Redundancy Optimizer)
ZeRO-DP是一种更先进的DP算法,单独拎出来是因为当下大名鼎鼎的deepspeed就是在此基础上建立起来的。
假设某个模型,有Φ个参数,基于adam+amp进行优化,则它所要存储的实际参数量为:
- FP16的参数 = 2Φ
- FP16的梯度 = 2Φ
- 优化器状态 = fp32参数 + fp32的momentum + fp32的var = 12Φ
如下图所示:
上图的、、分别对应ZeRO的三个优化stage(deepspeed中每个stage对应不同的可配置参数),相比与baseline,os+g+p优化了63倍。
之后又推出了ZeRO-2、ZeRO-offload、ZeRO-Infinity,一步一步把降显存提升到新高度。
ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelshttps://arxiv.org/pdf/1910.02054
MOE
Mixture of Experts (MoE),即混合专家,以前做过bagging的可能似曾相识,确实,总体的思想其实是很像的,但实现起来更复杂。
MOE由一个expert list和router组成,替换掉了FFN,如下图所示:
MOE的架构有两个好处:
- 可以在不明显增加资源占用的同时扩展模型规模,甚至可以无限扩展。
- 因为只会进行TOP-K个experts的计算,因此,可以减少存力算力开销。
缺点也很明显:
- 资源调度,通信可能是个瓶颈;
- 参数选择比较重要,比如top-k,capacity等;
但总的来说,MOE已经算是大模型的标配了。
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsityhttps://arxiv.org/pdf/2101.03961群魔乱舞:MoE大模型详解https://www.zhihu.com/tardis/zm/art/677638939?source_id=1005
fastmoehttps://github.com/laekov/fastmoeDeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scalehttps://arxiv.org/pdf/2201.05596
PEFT
大部分个人或者公司基本都是微调开源模型,这个比较实际,毕竟做底座真的很难。Parameter-Efficient Fine-Tuning即高效参数微调。包含几个不同的流派:
- Additive Fine-tuning
- Selective Fine-tuning
- Reparameterized Fine-tuning
- Hybrid Fine-tuning
建议阅读理解LoRA的官方代码,跑个demo,对比下FFT。
LoRAhttps://github.com/microsoft/LoRA
PEFThttps://github.com/huggingface/peft
开源
目前相关的开源的框架、代码都很多,列出几个有代表性的、强大的:
deepspeedhttps://www.deepspeed.ai/
ColossalAIhttps://colossalai.org/
Megatron-LMhttps://github.com/NVIDIA/Megatron-LM
horovodhttps://github.com/horovod/horovod
DeepSeek-V2https://github.com/deepseek-ai/DeepSeek-V2
结语
本文只做简单的介绍,提供一些指引,如果展开,上面提到的每个点都可以写一篇。当然已有不少优秀的博文可以参考,也没必要再造轮子。
所以,下一篇应该会详细记录下1080、 3090上的ds踩坑,填坑的过程。
引用
prigoyal/pytorch_memonger: Experimental ground for optimizing memory of pytorch models (github.com)GPipe: Easy Scaling with Micro-Batch Pipeline ParallelismAI and Memory Wall. (This blogpost has been written in… | by Amir Gholami | riselab | MediumGPipe: Easy Scaling with Micro-Batch Pipeline Parallelism
gradient-checkpointing
大模型涉及到的精度有多少种?FP32、TF32、FP16、BF16、FP8、FP4、NF4、INT8都有什么关联,一文讲清楚