Trainer - add cache clearing and the option for batched eval metrics computation
Author: FoamoftheSea
Pull request: yes

# What does this PR do?

This PR does two things which are necessary for using the Trainer in resource constrained environments (like my RTX-3070Ti machine):
1. Add cache clearing in training and evaluation loops
    - This reduces peak GPU load and prevents CUDA OOM errors when running near capacity.
2. Add Trainer arg `batch_eval_metrics` for batched eval metrics computation.
    - When working with limited RAM, storing all logits across the entire evaluation set may not be feasible. A user working in this condition can pass `True` to `batch_eval_metrics` and construct a `compute_metrics` function which can update average metrics at a batch level to prevent OOM errors with large eval sets. Particularly useful for vision transformers.
    - Previous functionality is unaltered if option is not set to `True`

@muellerzr 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_28769). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

cc @pacman100 and @muellerzr 

Hey everyone, I tried to look at the logs for the failed tests, but I don't see any actionable error reports. Can anyone help me figure out what needs to be done for them to pass?

The main CI is a bit broken because of `pytest` package. Let's wait a bit here

just re-ran the ci, you should actually rebase to main should be alright 

BTW @SunMarc would be nice if you can have a look as well! 

CIs are green after merging main ✔️ 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Commenting to keep the PR fresh. I got super busy the past couple weeks but I will finish this soon.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

>If we can confirm that it does reduce your memory footprint without issue then I believe that's quite alright. 

I will check this today, thank you!

This works amazing!
Here is the RAM consumption before:

<img width="407" alt="image" src="https://github.com/huggingface/transformers/assets/4803565/2df91dbc-b317-475a-bf10-66c0bad8001e">

And after:

<img width="402" alt="image" src="https://github.com/huggingface/transformers/assets/4803565/7036a8ea-5d90-4a4d-a5ea-600a5204d08e">


The only thing - PR needs to be updated to work with `4.40`, because `all_labels` logic has changed since `4.39`. I haven't tried to update PR to 4.40, tested on the its own branch only

@FoamoftheSea if you can resolve the merge conflicts, we can land this 🚀 

I am looking into this, it looks like the conflict is due to a different management of the variables in latest, which actually discards the use of the intermediate variable I was clearing from memory, so I want to double check on how that affects this change.

Also, I'm working on getting a test made for the batching functionality. Should have this ready soon.

The pytests are ready and I've updated the code to work with the latest Trainer updates. I still need to run an AB test to see if the cache clearing is still providing any benefit with the new changes from main, since they seem to have the potential to handle the same issue. I just need to finish setting up a test case. Let's hold off on merging until we have some fresh test results.

The test results demonstrate 2 things:

1. `batch_eval_metrics` is essential when training models with large solution spaces such as semantic segmentation where the typical code path accumulates many large tensors in memory during the eval loop, and batching the calculations avoids this.
2. The CUDA cache clearing keeps maximum GPU utilization lower over time, and allows using a larger eval batch size to expedite the process.

Based on these results, I think that we can justify these changes.

Training details:
- Model: `nvidia/segformer-b0-finetuned-cityscapes-1024-1024`
- Dataset: `Antreas/Cityscapes`

System details:
- Windows 10 Pro
- GPU = NVidia Quadro T2000
- CPU = Intel Xeon 2.80GHz (12 CPUs)
- RAM = 32GB

In the following chart, we can see that the standard code path goes OOM and fails as it tries to store all of the dense logits in memory on either the GPU or the CPU (as in the case of using `eval_accumulation_steps`). Only the run using `batch_eval_metrics` survives the evaluation cycle without going OOM and failing.

Further, we can see there is a great boost in memory efficiency during the train and eval phases using the cache clearing, which leads to very slightly lower iteration speed, but leaves a lot of headroom for using a larger batch size during evaluation to compensate.

![W B Chart 5_5_2024, 6_47_43 PM](https://github.com/huggingface/transformers/assets/50897218/7d3c73ec-6011-47be-9ad4-ebd3545682b1)
![W B Chart 5_5_2024, 6_48_29 PM](https://github.com/huggingface/transformers/assets/50897218/a688085b-44f7-4b62-b536-d729236f787d)
![W B Chart 5_5_2024, 6_48_54 PM](https://github.com/huggingface/transformers/assets/50897218/e3f27de7-dc05-4c45-ac65-ee7f0eabc2fb)
![W B Chart 5_5_2024, 6_49_16 PM](https://github.com/huggingface/transformers/assets/50897218/5b9b1a2f-d8ac-4699-8e01-6c6fcc7d491c)


Awesome, great work @FoamoftheSea!

Approving the PR, @muellerzr feel free to merge at your convenience.

Hello, I've noticed that this pull request seems to slow down the speed when using the trainer, likely due to the frequent use of `torch.cuda.empty_cache()`. Is there a way to optimize this, or could we possibly have the option to choose whether or not to use `torch.cuda.empty_cache()`?"@muellerzr 

Do you have a small benchmark for us @Reveyer  I haven’t noticed this yet when I was investigating another timing issue. But would be happy took into this 

@muellerzr I was able to reproduce this on the official code. Here’s the command I used:
```
CUDA_VISIBLE_DEVICES=0 python examples/pytorch/summarization/run_summarization.py \
    --model_name_or_path facebook/bart-large \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir /tmp/tst-summarization \
    --per_device_train_batch_size=16 \
    --per_device_eval_batch_size=16 \
    --overwrite_output_dir \
    --predict_with_generate \
    --num_train_epochs="10" \
    --seed="42"
```

Here are the results with and without `torch.cuda.empty_cache()` removed from the `trainer`:

- With `torch.cuda.empty_cache()` removed:
  | 30/179450 [00:22<37:48:51,  1.32it/s]
- With `torch.cuda.empty_cache()`:
  | 30/179450 [00:23<39:22:17,  1.27it/s]

The discrepancy is even more significant in my personal code, which utilizes Llama-3:

- With `torch.cuda.empty_cache()` removed:
  | 5/89543 [00:43<216:51:58,  8.72s/it]
- With `torch.cuda.empty_cache()`:
  | 5/89543 [00:51<251:54:05, 10.13s/it]

I hope this helps clarify the issue. Looking forward to hearing your thoughts!

@muellerzr @Reveyer I would test the speed difference on the second training iteration (after first eval round). This change sacrificed some speed on the first training iteration for noticeable increase in speed in the second and onward after that. Something about the eval round clogged up memory and made all subsequent training loops very slow, the cache emptying fixed that at a slight decrease in initial speed.

@FoamoftheSea  Hello, I've identified that these two lines of code are key to the speed differences observed. Could you please test the impact of these two lines on your experiments? Thank you!

https://github.com/huggingface/transformers/blob/2b9e252b16396c926dad0e3c31802b4af8004e93/src/transformers/trainer.py#L3270-L3271

Looks like we also have a CPU leak in here too when using custom evals, so I'll be investigating this today. 

@Reveyer sorry for the late response, the past few weeks have been very busy. I will find some time this week to re-run the experiment with that change.

I ran my test again on the main branch today. The results I got with/without the cache clearing in the training loop are mostly identical, other than that the run with no cache clearing was slightly less memory efficient, and in my case very slightly slower over time, but I think that might just be because my GPU was already warmed up for that training run.

Seeing how the difference here is minimal, I would have no problem with reverting the line in the training loop for cache clearing, since it doesn't seem to cause the memory overflow problem that the eval loop does without it.

![W B Chart 6_19_2024, 8_57_45 PM](https://github.com/huggingface/transformers/assets/50897218/3447e303-a042-463c-8739-5d50f8925864)
![W B Chart 6_19_2024, 8_58_55 PM](https://github.com/huggingface/transformers/assets/50897218/e808d498-24cd-4580-9e38-420e5055aac0)

Here's the script I ran the experiment with:
```python
from collections import Counter

import torch

import numpy as np

from typing import Mapping, Optional, Dict, Set
from datasets import load_dataset
from torchvision.transforms.functional import pil_to_tensor

from transformers import (
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
    Trainer,
    TrainingArguments,
    EvalPrediction,
)

from cityscapesscripts.helpers.labels import id2label

id2trainId = {k: v.trainId for k, v in id2label.items() if k >= 0}
trainId2label = {v.trainId: v.name for k, v in id2label.items() if k >= 0}
conversion_lookup = np.array([id2trainId[i] for i in range(len(id2trainId))])

dataset = load_dataset("Antreas/Cityscapes")
train_dataset, eval_dataset = dataset["train"], dataset["val"]

image_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-1024-1024")
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-1024-1024")

BATCH_EVAL_METRICS = True


class SegformerSemanticSegEvalMetric:
    def __init__(
            self,
            id2label: Dict[int, str],
            ignore_class_ids: Optional[Set[int]] = None,
            reduced_labels: bool = False,
            batch_eval_metrics: bool = True,
    ):
        self.total_area_intersect = Counter()
        self.total_area_union = Counter()
        self.total_label_area = Counter()
        self.ignore_class_ids = ignore_class_ids or set()
        self.reduced_labels = reduced_labels
        self.id2label = id2label
        self.batch_eval_metrics = batch_eval_metrics

    def update(self, logits: torch.FloatTensor, gt_labels: torch.LongTensor):

        # logits_tensor = torch.from_numpy(logits)
        # scale the logits to the size of the label
        logits_tensor = torch.nn.functional.interpolate(
            logits,
            size=gt_labels.shape[-2:],
            mode="bilinear",
            align_corners=False,
        ).argmax(dim=1)

        pred_labels = logits_tensor.detach().cpu().numpy()
        gt_labels = gt_labels.detach().cpu().numpy()

        for class_id in self.id2label.keys():
            if class_id in self.ignore_class_ids:
                continue
            if self.reduced_labels:
                label_id = class_id - 1 if class_id != 0 else 255
            else:
                label_id = class_id
            pred_pixels = pred_labels == label_id
            gt_pixels = gt_labels == label_id
            class_label = self.id2label[class_id]
            self.total_area_intersect.update({class_label: np.sum(np.bitwise_and(pred_pixels, gt_pixels))})
            self.total_area_union.update({class_label: np.sum(np.bitwise_or(pred_pixels, gt_pixels))})
            self.total_label_area.update({class_label: np.sum(gt_pixels)})

    def compute(self):
        accuracies = {f"accuracy_{k}": self.total_area_intersect[k] / self.total_label_area[k] for k in self.total_area_union}
        ious = {f"iou_{k}": self.total_area_intersect[k] / self.total_area_union[k] for k in self.total_area_union}
        metrics = {
            "overall_accuracy": sum(self.total_area_intersect.values()) / sum(self.total_label_area.values()),
            "mean_accuracy": np.mean(list(accuracies.values())),
            "mean_iou": np.mean(list(ious.values())),
        }
        metrics.update(accuracies)
        metrics.update(ious)

        return metrics

    def __call__(self, eval_pred: EvalPrediction, compute_result=False):
        if self.batch_eval_metrics:
            return self._call_batched(eval_pred, compute_result)
        else:
            return self._call_nonbatched(eval_pred)

    def _call_nonbatched(self, eval_pred):
        mious = {}
        with torch.no_grad():
            logits, gt_labels = eval_pred.predictions, eval_pred.label_ids
            logits_tensor = torch.from_numpy(logits)
            # scale the logits to the size of the label
            logits_tensor = torch.nn.functional.interpolate(
                logits_tensor,
                size=gt_labels.shape[-2:],
                mode="bilinear",
                align_corners=False,
            ).argmax(dim=1)

            pred_labels = logits_tensor.detach().cpu().numpy()

            for class_id in self.id2label.keys():
                if class_id in self.ignore_class_ids:
                    continue
                if self.reduced_labels:
                    label_id = class_id - 1 if class_id != 0 else 255
                else:
                    label_id = class_id
                pred_pixels = pred_labels == label_id
                gt_pixels = gt_labels == label_id
                class_label = self.id2label[class_id]
                intersection = np.sum(np.bitwise_and(pred_pixels, gt_pixels))
                union = np.sum(np.bitwise_or(pred_pixels, gt_pixels))
                mious[class_label] = intersection / union

        return np.mean(list(mious.values()))

    def _call_batched(self, eval_pred: EvalPrediction, compute_result: bool = True) -> Optional[dict]:
        with torch.no_grad():
            self.update(eval_pred.predictions, eval_pred.label_ids)
            return self.compute() if compute_result else None


def collate_fn(features: list):

    if not isinstance(features[0], Mapping):
        features = [vars(f) for f in features]
    first = features[0]
    images = None
    semantic_masks = None

    if "semantic_segmentation" in first and first["semantic_segmentation"] is not None:
        semantic_masks = [pil_to_tensor(f["semantic_segmentation"])[0] for f in features]
    if "image" in first and first["image"] is not None:
        images = [pil_to_tensor(f["image"].convert("RGB")) for f in features]

    processed = image_processor(images=images, segmentation_maps=semantic_masks)
    labels = np.array(processed.data["labels"])
    labels = conversion_lookup[labels]

    batch = {
        "pixel_values": torch.Tensor(np.array(processed.data["pixel_values"])),
        "labels": torch.LongTensor(np.array(labels)),
    }

    return batch


training_args = TrainingArguments(
    "segformer-test",
    num_train_epochs=10,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=8,
    save_total_limit=3,
    eval_strategy="steps",
    save_strategy="steps",
    save_steps=300,
    eval_steps=150,
    max_steps=10000,
    logging_steps=1,
    load_best_model_at_end=True,
    push_to_hub=False,
    gradient_checkpointing=False,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    use_cpu=False,
    learning_rate=0.0002,
    batch_eval_metrics=BATCH_EVAL_METRICS,
    remove_unused_columns=False,
    # eval_accumulation_steps=1 if not BATCH_EVAL_METRICS else None,
    run_name="segformer-b0-cityscapes-t4e8-batch-eval-metrics-no-cache-clear-train-loop",
)

trainer = Trainer(
    model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=SegformerSemanticSegEvalMetric(id2label=trainId2label, ignore_class_ids={255}),
    data_collator=collate_fn,
)

trainer.train()

```

Hi @FoamoftheSea, thanks for re-running the experiments ! From your observation, I think it will be better to remove cache clearing in the training loop as other users show a huge increase of training time. Would you like to open a PR ? Otherwise, I'll do it ! 

any solution to this issue? My trainer runs fine during train model but the same crashes when run with eval mode.