即使 RTX 3090 有着 24GB 的 RAM,使用一块 RTX 3090 依然无法 fp32 精度训练最小号的 LLaMA-6B。
估算模型所需的RAM
首先,需要了解如何根据参数量估计模型大致所需的 RAM,这在实践中有很重要的参考意义。需要通过估算设置 batch_size,设置模型精度,选择微调方法和参数分布方法等。
接下来用 LLaMA-6B 模型为例估算其大致需要的内存。
精度对所需内存的影响:
- fp32 精度,一个参数需要 32 bits, 4 bytes.
- fp16 精度,一个参数需要 16 bits, 2 bytes.
- int8 精度,一个参数需要 8 bits, 1 byte.
模型需要的 RAM 大致分三个部分:
- 模型参数
- 梯度
- 优化器参数
模型参数:等于参数量*每个参数所需内存。
- 对于 fp32,LLaMA-6B 需要 6B*4 bytes = 24GB内存
- 对于 int8,LLaMA-6B 需要 6B*1 byte = 6GB
梯度:同上,等于参数量*每个梯度参数所需内存。
优化器参数:不同的优化器所储存的参数量不同。
对于常用的 AdamW 来说,需要储存两倍的模型参数 (用来储存一阶和二阶momentum)。
- fp32 的 LLaMA-6B,AdamW 需要 6B*8 bytes = 48 GB
- int8 的 LLaMA-6B,AdamW 需要 6B*2 bytes = 12 GB
除此之外,CUDA kernel 也会占据一些 RAM,大概 1.3GB 左右,查看方式如下。
综上,int8 精度的 LLaMA-6B 模型部分大致需要 6GB+6GB+12GB+1.3GB = 25.3GB 左右。
再根据LLaMA的架构(hidden_size = 4096, intermediate_size =11008, num_hidden_layers = 32, context_length = 2048)计算中间变量内存。
每个 instance 需要:
所以一张 A100(80GB RAM)大概可以在 int8 精度;batch_size = 50 的设定下进行全参数训练。
ps:查看消费级显卡的内存和算力:
2023 GPU Benchmark and Graphics Card Comparison Chart
Fp16-mixed precision
混合精度训练的大致思路是在 forward pass 和 gradient computation 的时候使用 fp16 来加速,但是在更新参数时使用 fp32。
用 torch 实现:CUDA Automatic Mixed Precision examples
torch fp16 推理:直接使用 model.half() 将模型转换为fp16.
使用 Huggingface Transformers:在 TrainingArguments 里声明 fp16=True
Int8-bitsandbytes
Int8 是个很极端的数据类型,它最多只能表示 - 128~127 的数字,并且完全没有精度。
为了在训练和 inference 中使用这个数据类型,bitsandbytes 使用了两个方法最大程度地降低了其带来的误差:
- vector-wise quantization
- mixed precision decompasition
Huggingface 在这篇文章中用动图解释了 quantization 的实现
论文:LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale
借助 Huggingface PEFT,使用 int8 训练 opt-6.5B 的完整流程
LoRA
Low-Rank Adaptation 是微调 LLMs 最常用的省内存方法之一。
LoRA 的论文
借助 Huggingface PEFT 框架,使用 LoRA 微调 mt0
Gradient Checkpointing
在 torch 中使用 - 把 model 用一个 customize 的 function 包装一下即可,详见:Explore Gradient-Checkpointing in PyTorch
在 Huggingface Transformers 中使用
Torch FSDP+CPU offload
Fully Sharded Data Paralle(FSDP)和 DeepSpeed 类似,均通过 ZeRO 等分布优化算法,减少内存的占用量。其将模型参数,梯度和优化器状态分布至多个 GPU 上,而非像 DDP 一样,在每个 GPU 上保留完整副本。
CPU offload 则允许在一个 back propagation 中,将参数动态地从 GPU -> CPU, CPU -> GPU 进行转移,从而节省 GPU 内存。
Huggingface 这篇博文解释了 ZeRO 的大致实现方法
借助 torch 实现 FSDP,只需要将 model 用 FSDPwarp 一下;同样,cpu_offload 也只需要一行代码
在这个可以查看 FSDP 支持的模型
在 Huggingface Transformers 中使用 Torch FSDP
根据某些 issue,shard_grad_op(只分布保存 optimizer states 和 gradients)模式可能比 fully_shard 更稳定文章来源:https://www.toymoban.com/news/detail-487577.html
参考
pytorch模型训练之fp16、apm、多GPU模型、梯度检查点(gradient checkpointing)显存优化等
图解网络Offload
Flexgen LLM推理 CPU Offload计算架构到底干了什么事情?文章来源地址https://www.toymoban.com/news/detail-487577.html
到了这里,关于省显存(内存?)的大语言模型(LLMs)训练/微调/推理方法的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!