[Python] Is there any fix to resolve CUDA Out of Memory Error?

Discussão em 'Python' iniciado por Stack, Abril 24, 2025.

  1. Stack

    Stack Membro Participativo

    I am trying to train a model for spell check errors and corrections using T5-small model in colab notebook using T4 GPU. Initially I started with 10k reviews but gradually reduced it to 2k since I was receiving CUDA out of memory error. Now with 2k reviews I'm able to train the model and store it in my drive but then when evaluating the model it is giving the same memory error again. I have tried restarting the session and evaluate again but still no fix found.

    Error I'm receiving is :


    CUDA out of memory. Tried to allocate 6.99 GiB. GPU 0 has a total capacity of 14.74 GiB of which 6.95 GiB is free. Process 79643 has 7.79 GiB memory in use. Of the allocated memory 7.59 GiB is allocated by PyTorch, and 76.39 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
    ​

    I have tried increasing gradient_accumulation_steps but still the issue is not resolved below is my code (all functions are defined in above so there are no errors over there)

    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
    from datasets import load_dataset

    # Load dataset
    dataset = load_dataset("csv", data_files="elaborate_cosmetic_reviews.csv")

    # load the pre-trained Tokenizer and Model
    fine_model_name = "./spellcheck_model"
    tokenizer = T5Tokenizer.from_pretrained(fine_model_name, local_files_only=True)
    model = T5ForConditionalGeneration.from_pretrained(fine_model_name, local_files_only=True)

    # Define training arguments
    training_args = TrainingArguments(
    output_dir="./spellcheck_model_base",
    per_device_train_batch_size=2,
    report_to="none",
    fp16=True,
    gradient_accumulation_steps = 256,
    )

    # Create data collator
    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

    trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    eval_dataset=tokenized_train, # Use tokenized_eval for test set
    )
    import torch
    torch.cuda.empty_cache()

    # Evaluate on the training set
    print("Training set Evaluation: ")
    train_metrics = trainer.evaluate(eval_dataset=tokenized_train)
    print(train_metrics)```

    Continue reading...

Compartilhe esta Página