如何在使用accelerate的连续PyTorch训练阶段中正确管理GPU内存？

Question

我在使用PyTorch进行训练时遇到了一个棘手的问题，GPU内存在连续的训练阶段之间没有被正确释放，导致出现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后会变得像图3那样，但实际上仍然有很多GPU内存被占用。

尽管我明确删除了第一阶段使用的模型和数据加载器，并调用了gc.collect()和torch.cuda.empty_cache()，但GPU内存似乎并没有完全释放。因此，在开始第二个训练阶段时，我遇到了CUDA内存不足的错误。

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

有没有人遇到过类似的问题，或者有什么建议可以确保在训练阶段之间正确释放GPU内存？我考虑在阶段之间完全重启过程，但如果可能的话，我更希望有一个更干净的解决方案。

附注：这个问题是从韩文机器翻译过来的，如果语言有些生硬请见谅。

Answer 1

在多GPU环境中安排连续的工作负载可能会非常麻烦，特别是对于那些有多个阶段依赖的工作。如果你还在为这个问题苦恼，Run:ai是一个可以帮助公司高效管理工作负载并自动在集群上安排它们的替代方案。如果你觉得这可能对你有帮助，我们可以在Linkedin上聊聊。