算力不够也能玩大模型

@daviiid

已于 2024-05-24 18:02:27 修改

阅读量838

点赞数 27

分类专栏： AI 文章标签：人工智能神经网络深度学习 gpt-3

于 2024-05-22 12:04:22 首次发布

本文链接：https://blog.csdn.net/wb3533366/article/details/136347996

版权

AI 专栏收录该内容

13 篇文章 0 订阅

订阅专栏

背景

这个问题在当下甚至之后的一段时间应该都是比较常见的，毕竟数据量、模型参数量增加的速度远超过算力提升速度，再者不是每个组织或者个人都拥有那么多的显卡。

本文整理了一些方法，希望能对那些和我一样用不起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

Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limitedhttps://arxiv.org/pdf/1911.13214

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/parallelism GPipe论文精读【论文精读】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Φ

如下图所示：

上图的 $P_{os}$ 、 $P_{os+g}$ 、 $P_{os+g+p}$ 分别对应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/fastmoe DeepSpeed-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

Parameter-Efficient Fine-Tuning for Large Models: A Comprehensive Surveyhttps://arxiv.org/pdf/2403.14608

PEFThttps://github.com/huggingface/peft

开源

目前相关的开源的框架、代码都很多，列出几个有代表性的、强大的：

结语

本文只做简单的介绍，提供一些指引，如果展开，上面提到的每个点都可以写一篇。当然已有不少优秀的博文可以参考，也没必要再造轮子。

所以，下一篇应该会详细记录下1080、 3090上的ds踩坑，填坑的过程。

引用

Composer (mosaicml.com)

Determined AI Documentation — Determined AI Documentation

prigoyal/pytorch_memonger: Experimental ground for optimizing memory of pytorch models (github.com)GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism AI and Memory Wall. (This blogpost has been written in… | by Amir Gholami | riselab | Medium GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism

gradient-checkpointing
大模型涉及到的精度有多少种？FP32、TF32、FP16、BF16、FP8、FP4、NF4、INT8都有什么关联，一文讲清楚

Transformer Engine documentation

@daviiid

关注

27
点赞
踩
14

收藏

觉得还不错? 一键收藏
打赏
0
评论
算力不够也能玩大模型

计算资源不够怎么训练大模型？这个问题在当下甚至之后的一段时间应该都是比较常见的，毕竟数据量、模型参数量增加的速度远超过算力的提升速度，再者不是每个组织或者个人都拥有那么多的显卡。本文整理了一些行之有效的方法，希望能对那些和我一样买不起卡的道友有点帮助。
复制链接

扫一扫