ver 4.35.2 transformers.Trainer breaks CUDA AMP support
Author: haixpham
Pull request: no

### System Info

torch 2.1.1
transformers 4.35.2
CUDA 12.1

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

In the latest version of transformers (4.35.2), in `Trainer.__init__()`, there is no option for `use_cuda_amp`, and consequently, in `Trainer.autocast_smart_context_manager()`, `torch.cuda.amp.autocast()` is not invoked, which will lead to runtime error during loss backward.

This is the `__init__()` part of version 4.35.2, which does not enable autocast for both fp16 and bf16:
```
# Mixed precision setup
self.use_apex = False
self.use_cpu_amp = False

# Mixed precision setup for SageMaker Model Parallel
if is_sagemaker_mp_enabled():
    # BF16 + model parallelism in SageMaker: currently not supported, raise an error
    if args.bf16:
        raise ValueError("SageMaker Model Parallelism does not support BF16 yet. Please use FP16 instead ")

    if IS_SAGEMAKER_MP_POST_1_10:
        # When there's mismatch between SMP config and trainer argument, use SMP config as truth
        if args.fp16 != smp.state.cfg.fp16:
            logger.warning(
                f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                f"but FP16 provided in trainer argument is {args.fp16}, "
                f"setting to {smp.state.cfg.fp16}"
            )
            args.fp16 = smp.state.cfg.fp16
    else:
        # smp < 1.10 does not support fp16 in trainer.
        if hasattr(smp.state.cfg, "fp16"):
            logger.warning(
                f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                "but SageMaker Model Parallelism < 1.10 does not support FP16 in trainer."
            )
if (args.fp16 or args.bf16) and args.half_precision_backend == "auto":
    if args.device == torch.device("cpu"):
        if args.fp16:
            raise ValueError("Tried to use `fp16` but it is not supported on cpu")
        else:
            args.half_precision_backend = "cpu_amp"
    logger.info(f"Using {args.half_precision_backend} half precision backend")

if (args.fp16 or args.bf16) and not (self.is_deepspeed_enabled or is_sagemaker_mp_enabled()):
    # deepspeed and SageMaker Model Parallel manage their own half precision
    if args.half_precision_backend == "cpu_amp":
        self.use_cpu_amp = True
        self.amp_dtype = torch.bfloat16
    elif args.half_precision_backend == "apex":
        if not is_apex_available():
            raise ImportError(
                "Using FP16 with APEX but APEX is not installed, please refer to"
                " https://www.github.com/nvidia/apex."
            )
        self.use_apex = True
```

### Expected behavior

the same part of version 4.28.1 (which I know it works):
```
# Mixed precision setup
self.use_apex = False
self.use_cuda_amp = False
self.use_cpu_amp = False

# Mixed precision setup for SageMaker Model Parallel
if is_sagemaker_mp_enabled():
    # BF16 + model parallelism in SageMaker: currently not supported, raise an error
    if args.bf16:
        raise ValueError("SageMaker Model Parallelism does not support BF16 yet. Please use FP16 instead ")

    if IS_SAGEMAKER_MP_POST_1_10:
        # When there's mismatch between SMP config and trainer argument, use SMP config as truth
        if args.fp16 != smp.state.cfg.fp16:
            logger.warning(
                f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16},"
                f"but FP16 provided in trainer argument is {args.fp16},"
                f"setting to {smp.state.cfg.fp16}"
            )
            args.fp16 = smp.state.cfg.fp16
    else:
        # smp < 1.10 does not support fp16 in trainer.
        if hasattr(smp.state.cfg, "fp16"):
            logger.warning(
                f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                "but SageMaker Model Parallelism < 1.10 does not support FP16 in trainer."
            )

if args.fp16 or args.bf16:
    if args.half_precision_backend == "auto":
        if args.device == torch.device("cpu"):
            if args.fp16:
                raise ValueError("Tried to use `fp16` but it is not supported on cpu")
            elif _is_native_cpu_amp_available:
                args.half_precision_backend = "cpu_amp"
            else:
                raise ValueError("Tried to use cpu amp but native cpu amp is not available")
        else:
            args.half_precision_backend = "cuda_amp"

    logger.info(f"Using {args.half_precision_backend} half precision backend")

#the following part is no longer needed because of the switch to accelerator, but add it just in case 
self.do_grad_scaling = False
if (args.fp16 or args.bf16) and not (args.deepspeed or is_sagemaker_mp_enabled()):
    # deepspeed and SageMaker Model Parallel manage their own half precision
    if args.half_precision_backend == "cuda_amp":
        self.use_cuda_amp = True
        self.amp_dtype = torch.float16 if args.fp16 else torch.bfloat16
        #  bf16 does not need grad scaling
        self.do_grad_scaling = self.amp_dtype == torch.float16
        if self.do_grad_scaling:
            if self.sharded_ddp is not None:
                self.scaler = ShardedGradScaler()
            elif self.fsdp is not None:
                from torch.distributed.fsdp.sharded_grad_scaler import (
                    ShardedGradScaler as FSDPShardedGradScaler,
                )

                self.scaler = FSDPShardedGradScaler()
            elif is_torch_tpu_available():
                from torch_xla.amp import GradScaler

                self.scaler = GradScaler()
            else:
                self.scaler = torch.cuda.amp.GradScaler()
    elif args.half_precision_backend == "cpu_amp":
        self.use_cpu_amp = True
        self.amp_dtype = torch.bfloat16
    else:
        if not is_apex_available():
            raise ImportError(
                "Using FP16 with APEX but APEX is not installed, please refer to"
                " https://www.github.com/nvidia/apex."
            )
        self.use_apex = True
```

My workaround is to derive Trainer and add the 4.28.1 part of code above after calling `super().__init__()`, and override `autocast_smart_context_manager()` with the same method in version 4.34.1.

Comments:
cc @muellerzr 

Yep I see what happened with https://github.com/huggingface/transformers/pull/25702/files, it removed CUDA amp entirely for just CPU AMP. I'll work on this

cc @statelesshz 


If I'm not mistaken, the native AMP (Automatic Mixed Precision) has now been entirely migrated to Accelerate, see: https://github.com/huggingface/transformers/blob/fe41647afc98c7b5ce6c93a5f7cf57314bcd56f6/src/transformers/training_args.py#L1555-L1564

I wonder if it’s not picking it up in time as an env variable potentially. 

@muellerzr is this issue resolved?

cc @SunMarc as it's pretty important if still failing! 

@haixpham can you provide me with a minimal reproducer please? Setting `fp16=True` and `bf16=True` I could not recreate this issue on a minimal example:

```python
import evaluate
import numpy as np
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)

# Constants
model_name = "distilbert-base-uncased"
dataset_name = "imdb"
metric = "accuracy"

# AutoModel requires the label mapping
id2label = {0: "NEGATIVE", 1: "POSITIVE"}
label2id = {v: k for k, v in id2label.items()}

# Load dataset
print(f"Downloading dataset ({dataset_name})")
dataset = load_dataset(dataset_name)

# Tokenize the dataset
tokenizer = AutoTokenizer.from_pretrained(model_name)


def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True)


print(f"Tokenizing dataset for {model_name}...")
tokenized_dataset = dataset.map(tokenize_function, batched=True)

# Create an efficient collator which dynamically pads
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# Handle computation of our metrics
print(f"Loading metric ({metric})...")
accuracy = evaluate.load(metric)


def compute_metrics(evaluation_preds):
    predictions, labels = evaluation_preds
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)


print(f"Instantiating model ({model_name})...")
model = AutoModelForSequenceClassification.from_pretrained(
    model_name, num_labels=2, id2label=id2label, label2id=label2id
)

# Define the hyperparameters in the TrainingArguments
print("Creating training arguments (weights are stored at `results/sequence_classification`)...")
training_args = TrainingArguments(
    output_dir="results/sequence_classification",  # Where weights are stored
    learning_rate=2e-5,  # The learning rate during training
    per_device_train_batch_size=16,  # Number of samples per batch during training
    max_steps=5,  # How many iterations through the dataloaders should be done
    weight_decay=0.01,  # Regularization penalization
    fp16=True,
)

# Create the `Trainer`, passing in the model and arguments
# the datasets to train on, how the data should be collated,
# and the method for computing our metrics
print("Creating `Trainer`...")
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# Initiate training
print("Training...")
trainer.train()
```

@muellerzr The problem lied in the version mismatch between `transformers` and `accelerate`. I upgraded accelerate to 0.25.0 and the problem was gone.

hi,
I have experienced type issues with flash-attention and Llama models:
`RuntimeError: FlashAttention only support fp16 and bf16 data type`

which led me to dig into the Trainer (and this issue). I have a very light custom trainer, and use `bf16=True` (or `fp16=True`), but somehow it seems that the CUDA AMP is not properly called. When I add myself `with torch.cuda.amp.autocast(dtype=torch.bfloat16):` in the `compute_loss` (before the forward), everything runs smoothly. 
  
As pointed out by @haixpham, I cannot see anything related to `use_cuda_amp` in the Trainer (for instance in `autocast_smart_context_manager()`), which I find surprising (and in the end I was not surprised it breaks). But @statelesshz pointed out that this aspect has been migrated to Accelerate. It's not super clear to me how all these work together, and how the initial issue "disappeared"; I would appreciate some clarifications :)

Some details => I am using 
* torch==2.1.1 w/ CUDA11.8
* transformers==4.40.2 (latest)
* accelerate==0.30.1(latest)

Maybe https://github.com/Dao-AILab/flash-attention/issues/822, the issue is not with transformers I think

hi @ArthurZucker 
Thx for the pointer! I am a bit confused, bc if I add the right context manager before the forward, it works, so I add the feeling there was somehow an issue similar to the one above happening. Actually, I am not an expert with the Trainer, but I am having a hard time understanding how things can work with the current version of `autocast_smart_context_manager()`, bc `torch.cuda.amp.autocast()` is not invoked anymore. 

@thibault-formal can you give me a full reproducer? I just did a FFT in BF16 with flash attention on llama-3-8B over the weekend without trouble.

@ArthurZucker I will try to prepare something (my code is part of a complex pipeline so I need to make a dummy reproducer).