高 GPU RAM 使用率在 A100 上用小数据集训练大型语言模型

Question

我在相对较小的数据集上训练大型语言模型时遇到了 GPU RAM 消耗过多的问题。尽管仅使用 200 行数据，但训练过程在 Nvidia A100 GPU 上消耗了大约 40 GB RAM，这对于数据集大小来说似乎不成比例。

环境：

1. 模型：vilsonrodrigues/falcon-7b-instruct-sharded（a 的变体） 具有 70 亿个参数的大型语言模型）
2. 数据集大小：200行
3. GPU：Nvidia A100
4. 框架：PyTorch 和 Google Colab，
5. Hugging Face 的变形金刚库（指定版本）
6. 训练配置：
7. 批量大小：1
   • 梯度累积步数：16
   • 启用混合精度（FP16）

代码：https://colab.research.google.com/drive/1TNra_fwJbQ9M3FsFB1z8sniOxLzkDJfo?usp=sharing

问题： 训练消耗了大约 40 GB 的 GPU RAM，对于较小的数据集和所使用的训练配置来说，这似乎过多。我已经采用了诸如减少批量大小、启用混合精度训练以及使用梯度累积来管理内存使用等策略，但问题仍然存在。

OutOfMemoryError: CUDA out of memory. Tried to allocate 316.00 MiB. GPU 0 has a total capacty of 15.77 GiB of which 240.38 MiB is free. Process 36185 has 15.54 GiB memory in use. Of the allocated memory 15.19 GiB is allocated by PyTorch, and 43.11 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

问题：

• 是否有任何推荐策略可以进一步减少在小数据集上训练大型语言模型时的 GPU RAM 使用量？
• 我的训练设置中是否存在我可能忽视的潜在配置错误或效率低下？
• 有没有办法优化 A100 GPU 对于此类训练任务的利用率，以防止内存使用过多？

Answer 1

GPU RAM 消耗通常仅取决于模型权重和批量大小，而不取决于数据集大小。

在每次训练迭代中，都会加载模型权重并计算批次中所有数据点的输出，无论数据有 1 个还是 100 万个点。您可以尝试使用 1 个数据点来查看是否有任何变化（应该不会）。

您可能尝试减少内存消耗的一些策略是量化和 LoRA https://pytorch.org/blog/finetune-llms/。