Add bf16/fp16 support for amp with mps device

[SunMarc]

What does this PR do?

This PR adds MPS mixed-precision autocast support.

• FP16: autocast is supported in pytorch 2.5 but not GradScaler (Enable GradScaler for MPS devices pytorch/pytorch#142397).
• BF16: autocast is supported in pytorch 2.6. Also bf16 dtype is supported for MacOS 14 + . Got it from pytorch source code here.

Draft until we get support for GradScaler with autocast. Right now, support for bf16 ops with mps are still a bit limited but pytorch team is working on improving the coverage.

Feel free to test the PR to try bf16 for now

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[smartliuhw]

@SunMarc Hi Marc, I have tried to install this version of accelerate and pytorch2.6.0 to use trainer on mps device, but got the following error message, could you please help me check it out?

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[6], line 30
     20 trainer = Trainer(
     21     model=model,
     22     args=training_args,
   (...)
     27     compute_metrics=compute_metrics
     28 )
     29 logger.info("Start training")
---> 30 trainer.train()

File ~/miniforge3/envs/IOAI/lib/python3.12/site-packages/transformers/trainer.py:1885, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1883         hf_hub_utils.enable_progress_bars()
   1884 else:
-> 1885     return inner_training_loop(
   1886         args=args,
   1887         resume_from_checkpoint=resume_from_checkpoint,
   1888         trial=trial,
   1889         ignore_keys_for_eval=ignore_keys_for_eval,
   1890     )

File ~/miniforge3/envs/IOAI/lib/python3.12/site-packages/transformers/trainer.py:2216, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   2213     self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
   2215 with self.accelerator.accumulate(model):
-> 2216     tr_loss_step = self.training_step(model, inputs)
   2218 if (
   2219     args.logging_nan_inf_filter
   2220     and not is_torch_xla_available()
   2221     and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
   2222 ):
   2223     # if loss is nan or inf simply add the average of previous logged losses
   2224     tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)

File ~/miniforge3/envs/IOAI/lib/python3.12/site-packages/transformers/trainer.py:3250, in Trainer.training_step(***failed resolving arguments***)
   3248         scaled_loss.backward()
   3249 else:
-> 3250     self.accelerator.backward(loss)
   3252 return loss.detach() / self.args.gradient_accumulation_steps

File ~/GitHub/accelerate/src/accelerate/accelerator.py:2250, in Accelerator.backward(self, loss, **kwargs)
   2248     self.lomo_backward(loss, learning_rate)
   2249 else:
-> 2250     loss.backward(**kwargs)

File ~/miniforge3/envs/IOAI/lib/python3.12/site-packages/torch/_tensor.py:626, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    616 if has_torch_function_unary(self):
    617     return handle_torch_function(
    618         Tensor.backward,
    619         (self,),
   (...)
    624         inputs=inputs,
    625     )
--> 626 torch.autograd.backward(
    627     self, gradient, retain_graph, create_graph, inputs=inputs
    628 )

File ~/miniforge3/envs/IOAI/lib/python3.12/site-packages/torch/autograd/__init__.py:347, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    342     retain_graph = create_graph
    344 # The reason we repeat the same comment below is that
    345 # some Python versions print out the first line of a multi-line function
    346 # calls in the traceback and some print out the last line
--> 347 _engine_run_backward(
    348     tensors,
    349     grad_tensors_,
    350     retain_graph,
    351     create_graph,
    352     inputs,
    353     allow_unreachable=True,
    354     accumulate_grad=True,
    355 )

File ~/miniforge3/envs/IOAI/lib/python3.12/site-packages/torch/autograd/graph.py:823, in _engine_run_backward(t_outputs, *args, **kwargs)
    821     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    822 try:
--> 823     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    824         t_outputs, *args, **kwargs
    825     )  # Calls into the C++ engine to run the backward pass
    826 finally:
    827     if attach_logging_hooks:

RuntimeError: Expected scalar_type == ScalarType::Float || inputTensor.scalar_type() == ScalarType::Int || scalar_type == ScalarType::Bool to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)