LayoutLMv3 Significant Training Slowdown from 4.33.3 -> 4.34.0 and beyond versions
Author: kkoehncke
Pull request: no

### System Info

- `transformers` version: 4.33.3 
- Platform: Linux-6.5.0-1021-azure-x86_64-with-glibc2.35
- Python version: 3.11.5
- Huggingface_hub version: 0.23.1
- Safetensors version: 0.4.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.2.2+cu118 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Single GPU 

[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] torch                     2.2.2+cu118              pypi_0    pypi
[conda] torchaudio                2.2.2+cu118              pypi_0    pypi
[conda] torchmetrics              1.4.0.post0              pypi_0    pypi
[conda] torchvision               0.17.2+cu118             pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi

### Who can help?

@ArthurZucker @NielsRogge @younesbelkada 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

!pip install -q datasets seqeval

```
# Converted from notebook to runnable script from https://github.com/NielsRogge/Transformers-Tutorials/blob/master/LayoutLMv3/Fine_tune_LayoutLMv3_on_FUNSD_(HuggingFace_Trainer).ipynb
import os

import numpy as np
from datasets import Array2D, Array3D, ClassLabel, Features, Sequence, Value, load_dataset, load_metric
from transformers import AutoProcessor, LayoutLMv3ForTokenClassification, Trainer, TrainingArguments
from transformers.data.data_collator import default_data_collator

os.environ["TOKENIZERS_PARALLELISM"] = "false"

METRIC = load_metric("seqeval")

# we'll use the Auto API here - it will load LayoutLMv3Processor behind the scenes,
# based on the checkpoint we provide from the hub
processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=False, use_fast=True)


def prepare_examples(examples):
    images = examples["image"]
    words = examples["tokens"]
    boxes = examples["bboxes"]
    word_labels = examples["ner_tags"]

    encoding = processor(images, words, boxes=boxes, word_labels=word_labels, truncation=True, padding="max_length")

    return encoding


# In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
# unique labels.
def get_label_list(labels):
    unique_labels = set()
    for label in labels:
        unique_labels = unique_labels | set(label)
    label_list = list(unique_labels)
    label_list.sort()
    return label_list


def compute_metrics(p):
    return_entity_level_metrics = False
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = [[label_list[p] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels)]
    true_labels = [[label_list[l] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels)]

    results = METRIC.compute(predictions=true_predictions, references=true_labels)
    if return_entity_level_metrics:
        # Unpack nested dictionaries
        final_results = {}
        for key, value in results.items():
            if isinstance(value, dict):
                for n, v in value.items():
                    final_results[f"{key}_{n}"] = v
            else:
                final_results[key] = value
        return final_results
    else:
        return {
            "precision": results["overall_precision"],
            "recall": results["overall_recall"],
            "f1": results["overall_f1"],
            "accuracy": results["overall_accuracy"],
        }


if __name__ == "__main__":
    # this dataset uses the new Image feature :)
    dataset = load_dataset("nielsr/funsd-layoutlmv3")

    features = dataset["train"].features
    column_names = dataset["train"].column_names

    if isinstance(features["ner_tags"].feature, ClassLabel):
        label_list = features["ner_tags"].feature.names
        # No need to convert the labels since they are already ints.
        id2label = {k: v for k, v in enumerate(label_list)}
        label2id = {v: k for k, v in enumerate(label_list)}
    else:
        label_list = get_label_list(dataset["train"]["ner_tags"])
        id2label = {k: v for k, v in enumerate(label_list)}
        label2id = {v: k for k, v in enumerate(label_list)}
    num_labels = len(label_list)
    # we need to define custom features for `set_format` (used later on) to work properly
    features = Features(
        {
            "pixel_values": Array3D(dtype="float32", shape=(3, 224, 224)),
            "input_ids": Sequence(feature=Value(dtype="int64")),
            "attention_mask": Sequence(Value(dtype="int64")),
            "bbox": Array2D(dtype="int64", shape=(512, 4)),
            "labels": Sequence(feature=Value(dtype="int64")),
        }
    )

    train_dataset = dataset["train"].map(
        prepare_examples,
        batched=True,
        remove_columns=column_names,
        features=features,
    )
    eval_dataset = dataset["test"].map(
        prepare_examples,
        batched=True,
        remove_columns=column_names,
        features=features,
    )
    train_dataset.set_format("torch")
    model = LayoutLMv3ForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", id2label=id2label, label2id=label2id)
    training_args = TrainingArguments(
        output_dir=f"{os.getcwd()}/tmp",
        max_steps=1000,
        per_device_train_batch_size=4,
        per_device_eval_batch_size=4,
        learning_rate=1e-5,
        evaluation_strategy="steps",
        eval_steps=100,
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        report_to=[],
    )
    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=processor,
        data_collator=default_data_collator,
        compute_metrics=compute_metrics,
    )
    trainer.train()
```



### Expected behavior

Hello HuggingFace Team,

My team has been training LayoutLMv3 for quite some time and we have discovered that there is significant slowdown in the training time per epoch between `transformers` version `4.33.3` and `4.34.0` (and all version up to the current up-to-date release). This slowdown observed by us is solely due to the `transformers` package as we've tested on the same infra (Ubuntu 22.04, single A100 80GB GPU, NVIDIA Drivers 535.171.04) and have tried various `torch` packages (`2.x`, all CUDA 11.8, even tried `torch==1.13.1+cu117`) and all things point to `transformers` being the issue.

I've also observed this slowdown not only in our custom training code (w/ custom dataset), but also with the example from Niels on fine-tuning LayoutLMv3 on the FUNSD dataset, making the issue reproducible. 


tqdm sample output when `transformers==4.33.3` running sample train script:
`500/1000 [02:14<01:31]`

tqdm sample output when `transformers==4.34.0` running sample train script:
`6/1000 [00:10<26:41]`


Please let me know if any other information is needed to help re-produce & debug. Thanks! 


Comments:
Hi @kkoehncke !
Thanks for the issue - I suspect: https://github.com/huggingface/transformers/commit/a7e0ed829c398a67a641a401e23dae13e2f8b217 might be the culprit. Can you try to build transformers on that commit & right before that commit (being: ab37b801b14d8b9c3186548e6e118aff623e6aa1 ) and let us know?

@younesbelkada Thank you for the reply!

Followed your ask, and indeed the [a7e0ed8](https://github.com/huggingface/transformers/commit/a7e0ed829c398a67a641a401e23dae13e2f8b217) commit is the culprit. Commit right before [ab37b80](https://github.com/huggingface/transformers/commit/ab37b801b14d8b9c3186548e6e118aff623e6aa1), training speed looks as expected

To reproduce: 
```
> pip uninstall -y transformers
> pip install "git+https://github.com/huggingface/transformers.git@a7e0ed829c398a67a641a401e23dae13e2f8b217"
# The script mentioned above
> python demo_train_script.py  # <--- Very slow training speed

> pip uninstall -y transformers
> pip install "git+https://github.com/huggingface/transformers.git@ab37b801b14d8b9c3186548e6e118aff623e6aa1"
> python demo_train_script.py # <---- Normal training speed
```

I opened a PR to revert the commit: https://github.com/huggingface/transformers/pull/30988.

I tested from the branch with latest `main` code, and everything looks good from training speed perspective.

Let me know if there is anything else that I need to do, thanks!

Thank you for confirming @kkoehncke ! I think we could revert it, wdyt @amyeroberts @ArthurZucker ?

trying to find the actual rootcause, maybe the indexing operation is taking very long `rel_pos = self.rel_pos_bias.weight.t()[rel_pos].permute(0, 3, 1, 2)`

I think debugging the cause of this would be a good learning for all of us as to what is the culprit here! 

From debugging, it appears the indexing operation is the root problem. Interestingly, if you wrap it in a `torch.no_grad`, problem goes away (which matches up with my profiling runs, where the bottleneck of the time was spent during the backward pass)

Do we need the gradients to flow for that specific operation? 

if it's only on the indexing operation it should be safe i think no?

It depends on what it's used for, but yeah if the indices are always the same because of bucketing / independent from inputs, totally

Investigating more now, will post findings in a bit. But I think you guys are right, probably okay to not need to have gradients flow through that operation.

<img width="1635" alt="image" src="https://github.com/huggingface/transformers/assets/8129617/b3755aee-5f5c-48d6-a219-a73e97da2a4e">

Two sample training runs, with and without wrapping the indexing operation in a `torch.no_grad` context.

Upon debugging, you guys were right. It is just a simple indexing operation, indices remain constant throughout training, so can be safely wrapped in a `no_grad` context. 

@ArthurZucker I will close my existing PR to revert the commit (https://github.com/huggingface/transformers/pull/30988) and open a new one to wrap those operations with `torch.no_grad` if that is alright with you.

Thanks for your guys speedy replies and helpful discussion!

Sure, you can keep the PR and just push the corrected changes! 🤗 


Cool, updated the PR (https://github.com/huggingface/transformers/pull/30988) in-place