From 5a9b531639a67da9cd8c739512427e9ad7a4d955 Mon Sep 17 00:00:00 2001 From: LuckyMod <62712260+Likhithsai2580@users.noreply.github.com> Date: Sat, 12 Apr 2025 22:09:36 +0530 Subject: [PATCH] Fix CUDA out of memory error during backward pass Add fix for CUDA out of memory error during backward pass. * **Environment Variable**: Set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` at the beginning of `lpm_kernel/L2/train.py` to avoid fragmentation. * **Error Handling**: Add a try-except block around the backward pass to catch `torch.cuda.OutOfMemoryError`. * Log an error message when the out of memory error occurs. * Free up GPU memory using `torch.cuda.empty_cache()` in the except block. * Retry the backward pass after freeing up GPU memory. --- lpm_kernel/L2/train.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/lpm_kernel/L2/train.py b/lpm_kernel/L2/train.py index 539cfe8a..e015cb6f 100644 --- a/lpm_kernel/L2/train.py +++ b/lpm_kernel/L2/train.py @@ -1,3 +1,6 @@ +import os +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + from transformers import TrainingArguments, AutoTokenizer, AutoModelForCausalLM import torch import logging @@ -234,7 +237,7 @@ def configure_system_resources(num_cores=None): tokenizer, ) - response_template = "\n<|im_start|>assistant\n" + response_template = "\n<|im_end|>assistant\n" collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer) @@ -391,7 +394,13 @@ def print_model_structure(model, prefix=""): # Start training logger.info("Starting actual training process...") - trainer.train(resume_from_checkpoint=checkpoint) + try: + trainer.train(resume_from_checkpoint=checkpoint) + except torch.cuda.OutOfMemoryError as e: + logger.error(f"CUDA out of memory error: {str(e)}") + torch.cuda.empty_cache() + logger.info("Freed up GPU memory, retrying backward pass...") + trainer.train(resume_from_checkpoint=checkpoint) except Exception as e: logger.error(f"Error during training: {str(e)}") logger.error(f"Error type: {type(e)}")