TrainingArguments does not support `mps` device (Mac M1 GPU)

[saattrupdan]

System Info

• transformers version: 4.21.0.dev0
• Platform: macOS-12.4-arm64-arm-64bit
• Python version: 3.8.9
• Huggingface_hub version: 0.8.1
• PyTorch version (GPU?): 1.12.0 (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?: no

Who can help?

@sgugger

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

export TASK_NAME=wnli
python run_glue.py \
  --model_name_or_path bert-base-cased \
  --task_name $TASK_NAME \
  --do_train \
  --do_eval \
  --max_seq_length 128 \
  --per_device_train_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir /tmp/$TASK_NAME/

Expected behavior

When running the Trainer.train on a machine with an MPS GPU, it still just uses the CPU. I expected it to use the MPS GPU. This is supported by torch in the newest version 1.12.0, and we can check if the MPS GPU is available using torch.backends.mps.is_available().

It seems like the issue lies in the TrainingArguments._setup_devices method, which doesn't appear to allow for the case where device = "mps".

[saattrupdan]

A simple hack fixed the issue, by simply overwriting the device attribute of TrainingArguments:

import torch
from transformers import TrainingArguments


class TrainingArgumentsWithMPSSupport(TrainingArguments):

    @property
    def device(self) -> torch.device:
        if torch.cuda.is_available():
            return torch.device("cuda")
        elif torch.backends.mps.is_available():
            return torch.device("mps")
        else:
            return torch.device("cpu")

This at least shows that it might just be the aforementioned _setup_devices that needs changing.

[saattrupdan]

Another observation: Some PyTorch operations have not been implemented in mps and will throw an error. One way to get around that is to set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1, which will fallback to CPU for these operations. It still throws a UserWarning however.

[sgugger]

This is not supported yet, as this has been introduced by PyTorch 1.12, which also breaks all speech models due to a regression there. We will look into the support for Mac M1 GPUs once we officially support PyTorch 1.12 (probably won't be before they do a patch 1.12.1).

[saattrupdan]

@sgugger And it's not possible to add a use_mps flag to TrainingArguments, which just requires PyTorch 1.12.x, alongside a warning of some kind? Or is that too unstable?

[sgugger]

I have no idea, since we haven't tried and tested it out yet. And as I said our whole CI is constrained by PyTorch < 1.12 right now, so until that pin is dropped we can't test the integration :-). You can certainly try it on your own fork in the meantime!

[miguelwon]

I'm seeing this odd behavior. I'm trying a code from here, adapted with @saattrupdan solution. It runs but the problem is that I'm getting very different results when using cpu and mps. With device_type = "cpu" I get the expected results (f1=0.92) but when using device_type = "mps" I'm getting a very low f1 (~0.3), likely as a result a random guess.

device_type = "mps"
device = torch.device(device_type)
# Tokenizer
from transformers import AutoTokenizer
model_ckpt = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
def tokenize(batch):
    return tokenizer(batch["text"],padding = True,truncation = True)

# Load Data
from datasets import list_datasets
from datasets import load_dataset
emotions = load_dataset("emotion")
emotions_encoded = emotions.map(tokenize,batched = True,batch_size = None)

#Model
from transformers import AutoModelForSequenceClassification
num_labels = 6
model = (AutoModelForSequenceClassification.from_pretrained(model_ckpt,num_labels = num_labels).to(device))

#Metric
from sklearn.metrics import accuracy_score,f1_score
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(1)
    f1 = f1_score(labels,preds,average = "weighted")
    acc = accuracy_score(labels,preds)
    return {"accuracy":acc,"f1":f1}

#Train
from transformers import Trainer,TrainingArguments
class TrainingArgumentsWithMPSSupport(TrainingArguments):
    @property
    def device(self) -> torch.device:
        if device_type == "mps":
            return torch.device("mps")
        else:
            return torch.device("cpu")
batch_size = 64
loggin_steps = len(emotions_encoded["train"])
model_name = f"{model_ckpt}-finetuned-emotion"
train_args = TrainingArgumentsWithMPSSupport(output_dir = model_name,
                               num_train_epochs = 2,
                               learning_rate = 2e-5,
                               per_device_train_batch_size = batch_size,
                               per_device_eval_batch_size = batch_size,
                               weight_decay = 0.01,
                               evaluation_strategy = "epoch",
                               disable_tqdm = False,
                               logging_steps = loggin_steps,
                               push_to_hub = False,
                               log_level = "error"
                              )
trainer = Trainer(model = model,args = train_args,
                  compute_metrics = compute_metrics,
                  train_dataset = emotions_encoded["train"],
                  eval_dataset = emotions_encoded["validation"],
                  tokenizer = tokenizer)

print("Trainner device:",trainer.args.device)

trainer.train()

[sgugger]

We've also observed a drop in metrics when training, see this issue.

[julien-c]

Now that PyTorch 1.12.1 is out I think we should reopen this issue! cc @pacman100

[julien-c]

Note that on the inference side, pipelines now support device="mps" since #18494

[saattrupdan]

@julien-c That's great to hear! In my own scripts I've used this implementation, just tweaking the TrainingArguments._setup_devices method.

I also guess that the no_cuda training argument has to either be changed to no_gpu, if the current functionality should be preserved, or otherwise the handling of this keyword needs to be changed in the method (potentially adding a no_mps argument as well then, but I'm not sure if that's desirable).

I can open a PR if needed 🙂