# CUDA out of memory when using Trainer with compute\_metrics

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%3E%5C
**Category:** 🤗Transformers\
**Created:** [December 23, 2020, 4:14pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941 "2020-12-23T16:14:06Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![Randool](https://avatars.discourse-cdn.com/v4/letter/r/edb3f5/32.png) [@Randool](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FRandool)%5C
**Post date:** [December 23, 2020, 4:14pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F1 "2020-12-23T16:14:06Z")

</div>

Recently, I want to fine-tuning Bart-base with Transformers (version 4.1.1). The fine-tuning process is very smooth with `compute_metrics=None` in `Trainer`. However, when I implement a function of computing metrics and offer this function to `Trainer`, I received the `CUDA out of memory` error during the evaluation stage.

I want to try this feature, so my implementation is straightforward.

```auto
def compute_metrics(pred):
    preds = pred.predictions
    labels = pred.label_ids
    print(preds.shape, labels.shape)
    return {
        'loss': 1
    }

```

Because when `compute_metrics=None`, the training process is normal, so I think it can’t be the problem of batch size. But I still tried a smaller batch size, but even if I set the batch size to 1, the situation is the same. I even try `tinier-bart`, but I received the same error at last.

One thing caught my attention 🤔 . When I set a tiny batch size, the memory will not fill up at once, but the occupancy rate has increased until the memory can’t hold more data. That is, the processed data is not released in time. Is there any magic operation to solve this problem?

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [December 23, 2020, 5:51pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F2 "2020-12-23T17:51:11Z")

</div>

When computing metrics inside the `Trainer`, your predictions are all gathered together on the device (GPU/TPU) and only passed back to the CPU at the end (because that operation can be slow). If your dataset is large (or your model outputs large predictions) you can use `eval_accumulation_steps` to set a number of steps after which your predictions are sent back to the CPU (slower but uses less device memory). This should avoid your OOM.

---

<div class="post-metadata">

**Author:** ![Randool](https://avatars.discourse-cdn.com/v4/letter/r/edb3f5/32.png) [@Randool](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FRandool)%5C
**Post date:** [December 24, 2020, 2:41am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F3 "2020-12-24T02:41:02Z")

</div>

I have tried to use `eval_accumulation_steps`, but another problem occurs.

Here is part of my fine-tuning code:

```python
args = TrainingArguments(
    output_dir="exp/bart/results",
    do_train=True,
    do_eval=True,
    evaluation_strategy="steps",
    eval_steps=1000,
    logging_dir="exp/bart/logs",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    gradient_accumulation_steps=2,
    eval_accumulation_steps=1,
)

trainer = Trainer(
    model=bart,
    args=args,
    data_collator=collate_fn,
    train_dataset=train_set,
    eval_dataset=eval_set,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

```

* * *

Besides, the number of lines in the evaluation set is 22,161. When I set `eval_accumulation_steps=1`, I receive:

> MemoryError: Unable to allocate 149. GiB for an array with shape (22162, 36, 50265) and data type float32

It means that the `Trainer` will apply for all the space at once. Do I need to set other parameters to ensure that the `Trainer` applies for the right space every time?

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [January 4, 2021, 12:58pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F4 "2021-01-04T12:58:31Z")

</div>

This error means you are trying to get predictions that just don’t fit in RAM, so there is nothing `Trainer` can do to help. I don’t know which bart models you’re using, but it looks like you have huge logits so you should split your evaluation dataset in small parts or use a custom evaluation loop.

---

<div class="post-metadata">

**Author:** ![NicoHambauer](https://avatars.discourse-cdn.com/v4/letter/n/b9e5f3/32.png) [@NicoHambauer](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FNicoHambauer)%5C
**Post date:** [February 15, 2022, 12:26pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F6 "2022-02-15T12:26:57Z")

</div>

Thank you for this answer! It helped me a lot!

---

<div class="post-metadata">

**Author:** ![IamAdiSri](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/iamadisri/32/20434_2.png) [@IamAdiSri](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FIamAdiSri)%5C
**Post date:** [April 8, 2022, 5:31am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F7 "2022-04-08T05:31:57Z")

</div>

Hi, I’m still struggling with this issue. I’m trying to finetune a Bart model and while I can get it to train, I always run out of memory during the evaluation phase. This does not happen when I don’t use compute\_metrics, so I think there’s an issue there - when I don’t use compute\_metrics I can run batch sizes of up to 16, however on using compute metrics, I can’t even use a batch size of 1 even with eval accumulation.

Could you please explain why compute metrics is so much heavier when I can run training and evaluation without issues otherwise? In your answer above you mentioned that the trainer holds all predictions on the GPU but why is this being done for metrics calculation?

I have used Fairseq for seq2seq tasks with similarly sized models before this, and have never run into this issue before, so I was also wondering if they do metrics computation differently.

---

<div class="post-metadata">

**Author:** ![Dusan](https://avatars.discourse-cdn.com/v4/letter/d/9fc348/32.png) [@Dusan](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FDusan)%5C
**Post date:** [April 11, 2022, 6:51pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F8 "2022-04-11T18:51:58Z")

</div>

I also experience this when including my own compute\_metrics implementation, and it gradually increases GPU memory occupation over time.

Could it be that data structures (tensors I assume) used in our own implementation with each estimation are filling up GPU space and this is overloading our GPU device, and somehow default implementation is using memory garbage collector better? Should we somehow dump memory from variables not being used anymore over time? It seems variables are not being dumped after the compute\_metrics() function is done? @sgugger

---

<div class="post-metadata">

**Author:** ![kashif09](https://avatars.discourse-cdn.com/v4/letter/k/8e8cbc/32.png) [@kashif09](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fkashif09)%5C
**Post date:** [July 25, 2022, 3:26pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F9 "2022-07-25T15:26:40Z")

</div>

i was getting this same error so now i am doing this `eval_accumulation_steps`=16 with per\_device\_eval\_batch\_size=1, but i am now getting error “your ram collapsed bcoz you have used up the available ram” in google colab, any more help will be appreciated, i am using colab pro with 15 GB available GPU and model size 2.12 GB Pegasus,dataset Dialogsum, my train batch size is also 2

---

<div class="post-metadata">

**Author:** ![kashif09](https://avatars.discourse-cdn.com/v4/letter/k/8e8cbc/32.png) [@kashif09](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fkashif09)%5C
**Post date:** [July 25, 2022, 4:04pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F10 "2022-07-25T16:04:29Z")

</div>

were you able to solve your problem?

---

<div class="post-metadata">

**Author:** ![morenolq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/morenolq/32/55542_2.png) [@morenolq](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmorenolq)%5C
**Post date:** [September 18, 2022, 3:26pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F11 "2022-09-18T15:26:48Z")

</div>

I think it should be some leakage somewhere, it seems just not feasible that training works while eval is not.

The problem might be storing everything or large variables before computing the custom `compute_metrics`. Is this something you are planning to investigate?

---

<div class="post-metadata">

**Author:** ![Randool](https://avatars.discourse-cdn.com/v4/letter/r/edb3f5/32.png) [@Randool](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FRandool)%5C
**Post date:** [September 18, 2022, 3:46pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F12 "2022-09-18T15:46:25Z")

</div>

Yes, and I have found that `compute_metrics` will save all the logits before move them from CUDA to memory, just as sgugger said:

> [@sgugger](#):
>
> If your dataset is large (or your model outputs large predictions) you can use `eval_accumulation_steps` to set a number of steps after which your predictions are sent back to the CPU (slower but uses less device memory).

Sorry for forgetting mark this question as SOLVED 🙂

---

<div class="post-metadata">

**Author:** ![morenolq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/morenolq/32/55542_2.png) [@morenolq](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmorenolq)%5C
**Post date:** [September 18, 2022, 4:16pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F13 "2022-09-18T16:16:02Z")

</div>

I don’t think it solves the issue, it only moves it to RAM instead of GPU. The real solution is introduced with `preprocess_logits_for_metrics` function ([here](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Trainer.preprocess_logits_for_metrics)).

I leave here my specific solution (both functions):

```python
def compute_metrics(pred):

    labels_ids = pred.label_ids
    pred_ids = pred.predictions[0]

    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    labels_ids[labels_ids == -100] = tokenizer.pad_token_id
    label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)

    rouge_output = rouge.compute(
        predictions=pred_str,
        references=label_str,
        rouge_types=["rouge1", "rouge2", "rougeL", "rougeLsum"],
    )

    return {
        "R1": round(rouge_output["rouge1"], 4),
        "R2": round(rouge_output["rouge2"], 4),
        "RL": round(rouge_output["rougeL"], 4),
        "RLsum": round(rouge_output["rougeLsum"], 4),
    }

def preprocess_logits_for_metrics(logits, labels):
    """
    Original Trainer may have a memory leak. 
    This is a workaround to avoid storing too many tensors that are not needed.
    """
    pred_ids = torch.argmax(logits[0], dim=-1)
    return pred_ids, labels

```

BTW, proceeding in this way, you may not need to use `eval_accumulation_steps=1` (that slows down the evaluation significantly).

---

<div class="post-metadata">

**Author:** ![Randool](https://avatars.discourse-cdn.com/v4/letter/r/edb3f5/32.png) [@Randool](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FRandool)%5C
**Post date:** [September 19, 2022, 3:41am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F14 "2022-09-19T03:41:36Z")

</div>

Thank you very much! This should be the optimal solution.

---

<div class="post-metadata">

**Author:** ![AndyReas](https://avatars.discourse-cdn.com/v4/letter/a/db5fbb/32.png) [@AndyReas](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FAndyReas)%5C
**Post date:** [November 14, 2022, 4:22pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F15 "2022-11-14T16:22:27Z")

</div>

What is the reason for only using the first element of logits and predictions?

---

<div class="post-metadata">

**Author:** ![AndyReas](https://avatars.discourse-cdn.com/v4/letter/a/db5fbb/32.png) [@AndyReas](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FAndyReas)%5C
**Post date:** [November 22, 2022, 3:08pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F16 "2022-11-22T15:08:19Z")

</div>

I figured it out, maybe. Whatever is returned in preprocess\_logits\_for\_metrics(logits, labels) will be available in pred.predictions inside compute\_metrics(pred). By returning (pred\_ids, labels), pred.predictions becomes a tuple, where the second element is the labels. But the labels will still be available in pred.label\_ids no matter what. So only the predictions need to be returned by preprocess\_logits\_for\_metrics, and in that case, we don’t need to access the predicitons at index 0. As for only performing argmax on logits[0], I still don’t understand. Would that not just take the first element in the batch, and ignore the rest?

Let me know if I’m wrong 🙂

---

<div class="post-metadata">

**Author:** ![abuendia](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/abuendia/32/12728_2.png) [@abuendia](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fabuendia)%5C
**Post date:** [November 25, 2022, 9:38pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F17 "2022-11-25T21:38:23Z")

</div>

I think your understanding is correct on both. `logits[0]` would only take the first sequence in the batch assuming all other training args are left to their defaults – I modified above snippet to remove the subscript.

---

<div class="post-metadata">

**Author:** ![morenolq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/morenolq/32/55542_2.png) [@morenolq](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmorenolq)%5C
**Post date:** [December 5, 2022, 6:00pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F18 "2022-12-05T18:00:50Z")

</div>

Sorry for the late reply, Actually I used `logits[0]` because the version I was using was passing both logits and labels in logits (maybe I was also misinterpreting them). I agree with you, the correct snippet should be with `logits` instead of `logits[0]` in the argmax.

```python
pred_ids = torch.argmax(logits, dim=-1)

```

---

<div class="post-metadata">

**Author:** ![ghh001](https://avatars.discourse-cdn.com/v4/letter/g/bc8723/32.png) [@ghh001](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fghh001)%5C
**Post date:** [April 28, 2023, 3:04am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F19 "2023-04-28T03:04:32Z")

</div>

Yes, I solved the “CUDA out of memory when using Trainer with comput\_metrics” by using your solution. Thank you so much. 🥰

---

<div class="post-metadata">

**Author:** ![seba3y](https://avatars.discourse-cdn.com/v4/letter/s/e68b1a/32.png) [@seba3y](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fseba3y)%5C
**Post date:** [August 26, 2023, 10:35am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F20 "2023-08-26T10:35:53Z")

</div>

in my case, I’m working on speechT5ASR and the logits are tuple of 2 items, the first is the decoder output (logits that i need) and second is the encoder last hidden state. so l work with `logits[0]`

---

<div class="post-metadata">

**Author:** ![averoo](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/averoo/32/58222_2.png) [@averoo](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Faveroo)%5C
**Post date:** [September 26, 2023, 7:22am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941%2F21 "2023-09-26T07:22:28Z")

</div>

Thank you! Fixed my problem.

[Next page](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fcuda-out-of-memory-when-using-trainer-with-compute-metrics%2F2941.md%3Fpage%3D2)
