Auto-conversion of FP16 tensors to FP32 causes OOM error starting in transformers 4.26.0/1 releases
Author: awskila
Pull request: no

### System Info

```
transformers version: 4.26.0 and newer
platform: Ubuntu 20.04.5 (ec2 instance, 5.15.0-1028-aws)
python version: 3.8.13
other dependencies: boto3 1.26.68, awscli 1.26.68, datasets[s3] 2.5.2, sagemaker 2.132.0 + TrainingCompilerConfig, and xla
juggingface_hub version: 0.12.0
PyTorch version: 1.12.0
Using GPU in script?: yes
GPU:  1 NVIDIA T4 Tensor Core GPU with 16GB of VRAM
CPU: 8 Virtual Cascade Lake P-8259L CPUs
Using distributed or parallel set-up in script?: no
```


### Who can help?

@ArthurZucker @younesbelkada, possibly @sgugger given I'm using Trainer

### Information

- [ ] The official example scripts
- [X] 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

### Background ###
**Model Used:** roberta-base
**Dataset:** SST-2

My team at AWS works on the [SageMaker Training Compiler](https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html) that uses XLA to provide compiler-level optimizations. Since the last two transformers releases, we've noticed that datasets with fp16 tensors are being auto-converted to fp32 tensors. This is despite some of our training scripts specifying the `fp16` hyperparameter to `True`. 

### Problem ###
In some of our models, we were seeing memory footprints double, leading to a lot of OOM errors, especially on ec2 instances with GPUs with 16GB of VRAM. We've gotten around this by fixing the transformers version to 4.25.1 in all of our models, but we're concerned about this because our customers are just going to run `pip install transformers` and wonder why they're seeing OOM errors that they haven't seen before.

### Verifying and Trying to Fix the Problem ###
I've tried setting `fp16` to `False` in `transformers 4.25.1` using the `roberta-base` model with `sst-2` dataset, and the memory footprint was the same as when setting `fp16` to `True` in `transformers >= 4.26.0`. I spent some time trying to bisect the commit that made this change, but I wasn't able to narrow down the commit. 

I thought maybe [`keep_in_fp32_modules`](https://github.com/huggingface/transformers/blob/3b7ed25da9ce8291b9dabd6943f5b928352f4a16/src/transformers/modeling_utils.py#L628-L634) was being enabled under the hood, but I still observed this OOM error. I also tried reverting some of the T5 and RoBERTa-specific commits made in the last release, and wasn't able to restore the old behavior. 

Please let me know what I can do to fix this problem. I've seen this bug on multiple models using 4.26.0, 4.26.1, and 4.27.0.dev0 (using commit [`1efe9c0`](https://github.com/huggingface/transformers/commit/1efe9c0b24064acb3971260f505e46caa94617e9)). All the code needed to reproduce should be below - paste the estimator code into one file and the training script code in another, and modify the training script filename/path as needed.  Thanks a lot for all the work the team does - I'm a long time user but first-time bug reporter.

### Code ###

**Estimator**

```
!pip install --upgrade "boto3==1.26.68" "awscli==1.27.68"
!pip install "datasets[s3]==2.5.2" "torch==1.12.0" "fsspec<=2022.7.1"--upgrade
!pip install "transformers==4.26.0"
!pip install -U "sagemaker==2.132.0" "protobuf"

from datasets import Dataset
from transformers import AutoTokenizer
import os
import pandas as pd
import sagemaker
import botocore

tokenizer_name = 'roberta-base'
!aws s3 sync "s3://sagemaker-sample-files/datasets/text/SST2/" "./"

# download tokenizer
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

# tokenizer helper function
def tokenize(batch):
    return tokenizer(batch['text'], padding='max_length', truncation=True)

# load dataset
test_df = pd.read_csv('sst2.test', sep='delimiter', header=None, engine='python', names=['line'])
train_df = pd.read_csv('sst2.train', sep='delimiter', header=None, engine='python', names=['line'])
test_df[['label', 'text']] = test_df['line'].str.split(' ', 1, expand=True)
train_df[['label', 'text']] = train_df['line'].str.split(' ', 1, expand=True)
test_df.drop('line', axis=1, inplace=True)
train_df.drop('line', axis=1, inplace=True)
test_df['label'] = pd.to_numeric(test_df['label'], downcast='integer')
train_df['label'] = pd.to_numeric(train_df['label'], downcast='integer')
train_dataset = Dataset.from_pandas(train_df)
test_dataset = Dataset.from_pandas(test_df)

# tokenize dataset
train_dataset = train_dataset.map(tokenize, batched=True)
test_dataset = test_dataset.map(tokenize, batched=True)

# set format for pytorch
train_dataset =  train_dataset.rename_column("label", "labels")
train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
test_dataset = test_dataset.rename_column("label", "labels")
test_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])

save_path = "/opt/ml/model"

optimized_estimator = PyTorch(
        entry_point=trainer_script.py, #see below
        source_dir="./scripts",
        instance_type="ml.g4dn.2xlarge",
        instance_count=1,
        role="admin",
        py_version="py38",
        framework_version="1.12.0",
        volume_size=volume_size,
        hyperparameters={
            'epochs': 3,
            'max_steps': 5.
            'train_batch_size': 24,
            'model_name': 'roberta-base',
            'optim': 'adamw_torch_xla',
        } ,
        disable_profiler=True,
        debugger_hook_config=False,
        max_retry_attempts=3,
        compiler_config=TrainingCompilerConfig(debug=True),
        environment={
            "XLA_FLAGS": f"--xla_dump_to={save_path} --xla_dump_hlo_as_text",
            "XLA_SAVE_TENSORS_FILE": f"{save_path}/XLA_SAVE_TENSORS_FILE.hlo",
            "XLA_METRICS_FILE": f"{save_path}/XLA_METRICS_FILE.txt",
            "XLA_SAVE_HLO_FILE": f"{save_path}/XLA_SAVE_HLO_FILE.hlo",
        },
        distribution={'pytorchxla': {'enabled': True}},
    )

    # starting the train job with our uploaded datasets as input
    optimized_estimator.fit({"train": training_input_path, "test": test_input_path}, wait=True)

```

**Training Script**

```
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer, TrainerCallback
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from datasets import load_from_disk
import random
import logging
import sys
import argparse
import os
import torch
import numpy as np

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()
    
if __name__ == "__main__":
    main()

def main():

    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--max_steps", type=int, default=-1)
    parser.add_argument("--train_batch_size", type=int, default=32)
    parser.add_argument("--eval_batch_size", type=int, default=64)
    parser.add_argument("--warmup_steps", type=int, default=500)
    parser.add_argument("--model_name", type=str)
    parser.add_argument("--learning_rate", type=str, default=5e-5)

    # Data, model, and output directories
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--model_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

    args, _ = parser.parse_known_args()
    
    # The original LR was set for a batch of 32. We are scaling learning rate with batch size.
    args.learning_rate = float(args.learning_rate)
    args.learning_rate = float('5e-5')/32*args.train_batch_size

    # Set up logging
    logger = logging.getLogger(__name__)

    logging.basicConfig(
        level=logging.getLevelName("INFO"),
        handlers=[logging.StreamHandler(sys.stdout)],
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # load datasets
    train_dataset = load_from_disk(args.training_dir)
    test_dataset = load_from_disk(args.test_dir)

    # compute metrics function for binary classification
    def compute_metrics(pred):
        labels = pred.label_ids
        preds = pred.predictions.argmax(-1)
        precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="binary")
        acc = accuracy_score(labels, preds)
        return {"accuracy": acc, "f1": f1, "precision": precision, "recall": recall}

    # download model from model hub
    model = AutoModelForSequenceClassification.from_pretrained(args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)

    # define training args
    training_args = TrainingArguments(
        output_dir=args.model_dir,
        num_train_epochs=args.epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.train_batch_size,
        per_device_eval_batch_size=args.eval_batch_size,
        warmup_steps=args.warmup_steps,
        logging_dir=f"{args.output_data_dir}/logs",
        learning_rate=args.learning_rate,
        fp16=True, #note that this is True
        dataloader_drop_last=True,
        disable_tqdm=True,
        evaluation_strategy="no",
        save_strategy="no",
        save_total_limit=1,
        logging_strategy="epoch",
    )
    
    # create Trainer instance
    trainer = Trainer(
        model=model,
        args=training_args,
        compute_metrics=compute_metrics,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        tokenizer=tokenizer,
    )

    # train model
    trainer.train()
```

### Expected behavior

**Expected:** With a GPU with 16GB VRAM, the largest XLA module based on buffer-assignment/disassembly files should be Module 3, and the total memory footprint should be 7.859GB. Training should finish without going out of memory. The buffer assignment file shows a mix of fp16 and fp32 tensors. This is the behavior on transformers versions older than 4.26.0, which I tested on 4.21.0, 4.24.0, and 4.25.1 and observed this behavior.

**Actual:** Module 3 has a total memory footprint of 15.463GB, and will lead to an out-of-memory error on a GPU of 16GB. Looking at the buffer assignment file for module 3, all tensors are fp32. This happens on 4.26.0, 4.26.1, and 4.27.0.dev0

Comments:
May I please know if there's an update on this? 

Please be patient while we try to reproduce your problem. If you can single out the commit you think is responsible, it might help our search.

Sure, not a problem. Thank you for your reply. You're right, I should've included the commits in my initial post. I'll cover them in this comment.

So, I did try a run at using git bisect to narrow down the issue, and I wasn't able to find the correct commit(s) to revert. I think maybe there is more than 1 commit that may have contributed to this bug, or (more likely) I didn't find the right commit. With that being said, here are some commits that I was looking at:

`fe65657de112531a2d5303491f245f9e7534ae8d` - maybe `torch_dtype` is not returning the proper type for my model? It could be defaulting `fp32` instead of `fp16`, whereas the input argument for preprocess function, `dtype` may be `fp16`  due to the override I'm passing in my estimator.

`accad48e5b4a98302ea396b9f15c5f1c987b6f7f` - There's an explicit conversion to `torch.float32` in the `_load_state_dict_into_meta_model` function. Maybe that's what I'm seeing.

`1af4bee8965c80e54c0e21aa8aadd035fd1f4189` - Maybe a bug in `keep_in_fp32_modules`? It could be if `fp16` is set to `True` in an estimator, `use_keep_in_fp32_modules` is being set to `True` when it should be `False`?

You are not using any `torch_dtype` to build your model and the model you are using is saved in float32, so none of those commits should change anything. `fp16=True` is mixed precision training in float16, so the model is still in float32 (there are actually two copies of the model, one in float16 and one in float32).

So `torch_dtype` is not set based on the `fp16` value. Hmmm ok. I realize that setting fp16 is for mixed precision, and doesn't convert fp32 tensors to fp16 or something like that. I was thinking that the fp16 copy is being converted to fp32 at some point, or something like that. I don't know for sure if that's happening. Also I'm not sure what about other commits could've caused the issue, other than it's something committed during the 4.26.0 release


Seems like this is still a problem we're observing. The largest XLA module is doubling in VRAM for any mixed precision training, as fp16 tensors are being converted to fp32. This is with setting `fp16=True`, and the behavior has been different since transformers `v4.26.0`. Is there a special thing we should do for our notebooks to ensure this doesn't happen? We want to use the latest Transformers, but our customers need to avoid the increased memory footprint. Please advise. Thanks!

@awskila can you please share the whole example so we are able to reproduce. We are missing the `requirements.txt` for the libraries you installed in the SageMaker training job also does you "Estimator" code have errors. 

Sure, give me a bit more time. I will share everything very soon.

I experience similar OOM errors when calling `model.generate()` at inference time, after the upgrade from 4.25 to 4.26.

Same issue here, `model.generate()` gave me OOM with 4.26.1, so I rolled back to 4.25.1 to fix it.

This has nothing to do with this issue which relates to mixed precision training, so please open a new one with a reproducer.

Good to hear this isn't just me. I will try and post a full working sample. Have to modify the model code in a way that allows it to stay open source, and I was using a customized sst-2 dataset that had differences from the open-source sst-2 dataset. My apologies for the delays.

As far as how this relates to mixed precision training - I was showing that behavior is not working as it did in a prior release. It's very possible (and likely imo) that the real issue is occurring elsewhere, and I just noticed this with mixed precision training. I don't think the issue is at the `model.generate` level given that code has been [stable since 4.25.1](https://github.com/huggingface/transformers/commit/f270b960d6b755d5355a0f193ba71f47131c86b8).

Essentially my issue was 4.26.0+ mixed precision training had a memory footprint double of <4.25.1 mixed precision training. And the resulting disassembly is similar to <4.25.1 fp32 training. 

I was running training off a pre-trained model, and it looks like this has been seen on inference too per @pgmikhael. It may not be two separate issues, but maybe it's a more general issue. I would welcome a new issue with a reproducer though, as I have been getting delayed for compliance reasons.

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.