如何使用加速在连续的 PyTorch 训练阶段之间正确管理 GPU 内存？

Question

我遇到了一个具有挑战性的问题，GPU 内存在 PyTorch 的连续训练阶段之间没有正确释放，导致 CUDA 内存不足错误。

我的项目涉及分两个连续阶段微调模型：

• 首先在 FP（进一步预训练阶段）数据集上，
• 然后在 SFT（监督微调）数据集上。

代码结构如下：

from transformers import AutoModelForCausalLM

# Model and data loader initialization
model = AutoModelForCausalLM.from_pretrained(args.model)  # fig.1
fp_data_loader = data_loader(fp_dataset)
sft_data_loader = data_loader(sft_dataset)

# First phase of training
fp_model, fp_loss = train_loop(model, fp_data_loader) #fig.2
fp_model.module.save_pretrained(checkpoint_dir)

# Attempt to release GPU memory
del model, fp_data_loader,
# del fp_model
# fp_model = AutoModelForCausalLM.from_pretrained(checkpoint_dir)
gc.collect()
torch.cuda.empty_cache()  #fig.3

# Second phase of training
sft_model, sft_loss = train_loop(fp_model, sft_data_loader)

图 1：当模型在 GPU 上时

图2：训练期间

图3：empty_cache之后（捕获问题：忽略gpu0的77762/81920，它是57524/81920）

我的预期是图1的GPU分配会像empty_cache之后一样，但是分配了相当多的GPU内存，如图3所示

尽管显式删除了第一阶段使用的模型和数据加载器并调用了 gc.collect() 和 torch.cuda.empty_cache()，但 GPU 内存似乎并未完全释放。结果，当启动第二个训练阶段时，我遇到了 CUDA 内存不足错误。

我还使用 Accelerator.prepare() 包装我的模型、优化器和数据加载器，以进行混合精度和分布式训练，这可能是相关的。

是否有人遇到过类似的问题或对确保在训练阶段之间正确释放 GPU 内存有建议？我考虑过在阶段之间完全重新启动该过程，但如果可能的话，我更喜欢更干净的解决方案。

PS：这个问题是韩语机器翻译的，如果有尴尬的语言，请见谅。

Answer 1

在多 GPU 环境上调度连续的工作负载可能会非常痛苦，特别是对于具有多阶段依赖性的工作负载。如果您仍然遇到挑战，Run:ai 是一种替代方案，可以帮助公司有效管理工作负载并在集群上自动调度它们。如果您认为有用，我们可以在 Linkedin 上聊天