训练模型后如何清除 Google Collab 中的 GPU 内存

Question

我目前正在尝试在 Google Collab 中与 Llama-3 LLM 进行交叉验证，但我面临着 GPU 内存在我完成实验之前就耗尽的问题。我的代码如下：

我有这个功能来清除内存：

from GPUtil import showUtilization as gpu_usage
from numba import cuda
import gc
from accelerate import Accelerator

accelerator = Accelerator()

def free_gpu_cache():
  print('Initial GPU Usage')
  gpu_usage()

  gc.collect()
  torch.cuda.empty_cache()
  accelerator.free_memory()

  cuda.select_device(0)
  cuda.close()
  cuda.select_device(0)

  print("GPU Usage after emptying the cache")
  gpu_usage()

这是我正在使用的主要代码：

for train_index, test_index in kf.split(data):
    training_data, test_data = data['text'][train_index], data['text'][test_index]

    training_data = Dataset.from_pandas(training_data.to_frame().reset_index())
    test_data = Dataset.from_pandas(test_data.to_frame().reset_index())

    
    free_gpu_cache()

    # model
    model = AutoModelForCausalLM.from_pretrained(model_name,
                                                device_map='auto',
                                                quantization_config=bnb_config,
                                                token = access_token)

    model.config.use_cache = False
    model.config.pretraining_tp = 1

    trainer = SFTTrainer(
      model = model,
      train_dataset = training_data,
      peft_config = lora_config,
      dataset_text_field = 'text',
      max_seq_length = 30,
      tokenizer = tokenizer,
      args = training_arguments,
      packing = True
    )

    trainer.train()

即使我有清除内存的功能，我在第二次迭代中收到以下错误：

ValueError：某些模块被调度到CPU或磁盘上。确保您有足够的 GPU RAM 来适应量化模型。

我尝试使用上面显示的功能清除内存，但没有成功。

我也尝试过这个：

import gc
gc.collect()
torch.cuda.empty_cache()

但它也不起作用。

我不知道还能做什么

Answer 1

我不完全确定这是否正确，所以您可能必须等待其他人回答。但我相信“刷新谷歌 colabs”ram 不会起作用，因为 colab 通过以下方式赚钱：1. 限制 ram 访问，2. 付费访问更好的 GPU。