Make torch xla available on GPU
Author: yitongh
Pull request: yes

# What does this PR do?

Make torch xla available on GPU. Currently, torch xla can be used in a GPU environment, but there are some conflicts between XLA and native PyTorch CUDA when using an environment with torch xla installed. This PR introduces the environment variable USE_TORCH_XLA to address this issue. When USE_TORCH_XLA is set to false, native PyTorch CUDA can be used seamlessly, even if torch xla is installed.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@muellerzr and @pacman100



Comments:
The main changes:

1. Change `is_torch_tpu_available` to `is_torch_xla_available`
2. Change `require_torch_tpu` to `require_torch_xla`
3. Add `USE_TORCH_XLA` to enable or disable torch_xla
4. Fix amp check
5. Move `grad_norm.item()` into `_maybe_log_save_evaluate` to prevent a performance degradation in XLA
6. Copy the xla_fsdp_config to avoid modifying the original config

This PR is related to https://github.com/huggingface/accelerate/pull/2176 and https://github.com/huggingface/accelerate/pull/2467.

@will-cromar Could you please check if this PR has any impact on the TPU environment? Thanks.


The ci only failed in tests/test_modeling_utils.py::ModelUtilsTest::test_use_safetensors. I run this test in master and it also hangs. From pystack, it looks like this issue related to ssl read. stack log: https://gist.github.com/yitongh/34dc9c9f3de79d208533964bd63bb6f5

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29334). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@muellerzr, would you be available to take a look at this PR when you have a moment? Alternatively, if you're not, perhaps you could suggest someone else who might be suited to review it? Many thanks.

Let's make sure we can fix those failing tests though, can you try rebasing from `main`?

@muellerzr  I have rebased from main. I rerun the failing in my machine, both main and this pr passed `test_run_ner_no_trainer` and `test_run_squad_no_trainer`, but failed at `test_run_glue_no_trainer`. It looks like not related to this pr.

```
pytest -s -v examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_ner_no_trainer examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_squad_no_trainer examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_glue_no_trainer
======================================================================================================== test session starts =========================================================================================================
platform linux -- Python 3.10.12, pytest-8.0.0, pluggy-1.4.0 -- /usr/bin/python3.10
cachedir: .pytest_cache
hypothesis profile 'default' -> database=DirectoryBasedExampleDatabase(PosixPath('/root/hyt/github/transformers/.hypothesis/examples'))
rootdir: /root/hyt/github/transformers
configfile: pyproject.toml
plugins: hypothesis-6.98.17, xdist-3.5.0, subtests-0.12.1, anyio-4.3.0, timeout-2.3.1
collected 3 items

examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_ner_no_trainer PASSED
examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_squad_no_trainer PASSED
examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_glue_no_trainer FAILED

============================================================================================================== FAILURES ==============================================================================================================
__________________________________________________________________________________________ ExamplesTestsNoTrainer.test_run_glue_no_trainer ___________________________________________________________________________________________

self = <test_accelerate_examples.ExamplesTestsNoTrainer testMethod=test_run_glue_no_trainer>

    @mock.patch.dict(os.environ, {"WANDB_MODE": "offline", "DVCLIVE_TEST": "true"})
    def test_run_glue_no_trainer(self):
        tmp_dir = self.get_auto_remove_tmp_dir()
        testargs = f"""
            {self.examples_dir}/pytorch/text-classification/run_glue_no_trainer.py
            --model_name_or_path distilbert/distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --seed=42
            --num_warmup_steps=2
            --checkpointing_steps epoch
            --with_tracking
        """.split()

        run_command(self._launch_args + testargs)
        result = get_results(tmp_dir)
>       self.assertGreaterEqual(result["eval_accuracy"], 0.75)
E       AssertionError: 0.6666666666666666 not greater than or equal to 0.75

examples/pytorch/test_accelerate_examples.py:98: AssertionError
========================================================================================================== warnings summary ==========================================================================================================
../../../../usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py:1394
  /usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py:1394: PytestConfigWarning: Unknown config option: doctest_glob

    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
====================================================================================================== short test summary info =======================================================================================================
FAILED examples/pytorch/test_accelerate_examples.py::ExamplesTestsNoTrainer::test_run_glue_no_trainer - AssertionError: 0.6666666666666666 not greater than or equal to 0.75
========================================================================================= 1 failed, 2 passed, 1 warning in 143.07s (0:02:23) =========================================================================================
```

They look to be timeout issues, I'm rerunning the tests now. However if they still fail they were not failing before this 😅 

Tests pass on our CI so looks to be fine