Repository: huggingface/accelerate

Issue:
Use FastChat with NPU fine-tune LLM get an AssertionError
### System Info

```Shell
linux：ubuntu 20.04/aarch64
driver：npu-smi 23.0.rc2
CANN：Ascend-cann-toolkit_7.0.0.alpha001_linux-aarch64.run
Python：3.9.18
torch：2.1.0
torch-npu：2.1.0rc1
transformers：4.34.1
accelerate：0.24.0
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

When use NPU devices to fine-tune LLM under version 0.24.0, i got these：
```shell
(py39) root@ascend910b-01:/opt/jq/projects/FastChat# torchrun --nproc_per_node=2 --master_port=20001 fastchat/train/train.py \
>     --model_name_or_path /opt/jq/big_models/llama-2-7b/  \
>     --data_path /opt/jq/nlp_data/evol-instruct-chinese-subset.json \
>     --fp16 True \
>     --output_dir output_llama \
>     --num_train_epochs 2 \
>     --per_device_train_batch_size 8 \
>     --per_device_eval_batch_size 1 \
>     --gradient_accumulation_steps 1 \
>     --evaluation_strategy "no" \
>     --save_strategy "steps" \
>     --save_steps 2000 \
>     --save_total_limit 200 \
>     --learning_rate 5e-5 \
>     --weight_decay 0. \
>     --lr_scheduler_type "cosine" \
>     --logging_steps 1 \
>     --fsdp "full_shard auto_wrap" \
>     --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
>     --model_max_length 512 \
>     --gradient_checkpointing True \
>     --lazy_preprocess True
[2023-10-26 03:57:20,063] torch.distributed.run: [WARNING] 
[2023-10-26 03:57:20,063] torch.distributed.run: [WARNING] *****************************************
[2023-10-26 03:57:20,063] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2023-10-26 03:57:20,063] torch.distributed.run: [WARNING] *****************************************
/root/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/training_args.py:1598: FutureWarning: using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead 
  warnings.warn(
/root/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/training_args.py:1598: FutureWarning: using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead 
  warnings.warn(
Loading checkpoint shards: 100%|███████████████████| 2/2 [00:17<00:00,  8.87s/it]
/opt/jq/projects/FastChat/fastchat/train/train.py:235: ResourceWarning: unclosed file <_io.TextIOWrapper name='/opt/jq/nlp_data/evol-instruct-chinese-subset.json' mode='r' encoding='UTF-8'>
  train_json = json.load(open(data_args.data_path, "r"))
Traceback (most recent call last):
  File "/opt/jq/projects/FastChat/fastchat/train/train.py", line 302, in <module>
    train()
  File "/opt/jq/projects/FastChat/fastchat/train/train.py", line 287, in train
    trainer = Trainer(
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/trainer.py", line 347, in __init__
    self.create_accelerator_and_postprocess()
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/trainer.py", line 3981, in create_accelerator_and_postprocess
    self.accelerator = Accelerator(
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/accelerate/accelerator.py", line 313, in __init__
    FullyShardedDataParallelPlugin() if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" else None
  File "<string>", line 18, in __init__
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/accelerate/utils/dataclasses.py", line 934, in __post_init__
    device = torch.cuda.current_device() if not is_xpu_available() else torch.xpu.current_device()
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/cuda/__init__.py", line 769, in current_device
    _lazy_init()
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/cuda/__init__.py", line 289, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled
/root/miniconda3/envs/py39/lib/python3.9/tempfile.py:821: ResourceWarning: Implicitly cleaning up <TemporaryDirectory '/tmp/tmp1neqk6hl'>
  _warnings.warn(warn_message, ResourceWarning)
[2023-10-26 03:58:00,127] torch.distributed.elastic.multiprocessing.api: [WARNING] Sending process 7682 closing signal SIGTERM
[2023-10-26 03:58:03,045] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: 1) local_rank: 1 (pid: 7683) of binary: /root/miniconda3/envs/py39/bin/python3.9
Traceback (most recent call last):
  File "/root/miniconda3/envs/py39/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/run.py", line 806, in main
    run(args)
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/run.py", line 797, in run
    elastic_launch(
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/root/miniconda3/envs/py39/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
fastchat/train/train.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-10-26_03:58:00
  host      : ascend910b-01
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 7683)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
```


### Expected behavior

Successfully fine-tune LLM using NPU.


Discussion:

Resolving patch:
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -32,7 +32,7 @@
 
 from .constants import FSDP_AUTO_WRAP_POLICY, FSDP_BACKWARD_PREFETCH, FSDP_STATE_DICT_TYPE
 from .environment import str_to_bool
-from .imports import is_xpu_available
+from .imports import is_cuda_available, is_npu_available, is_xpu_available
 from .versions import compare_versions
 
 
@@ -932,7 +932,16 @@ def __post_init__(self):
         self.activation_checkpointing = str_to_bool(os.environ.get(prefix + "ACTIVATION_CHECKPOINTING", "False")) == 1
 
         if self.sync_module_states:
-            device = torch.cuda.current_device() if not is_xpu_available() else torch.xpu.current_device()
+            if is_npu_available():
+                device = torch.npu.current_device()
+            elif is_cuda_available():
+                device = torch.cuda.current_device()
+            elif is_xpu_available():
+                device = torch.xpu.current_device()
+            else:
+                raise RuntimeError(
+                    "There are currently no available devices found, must be one of 'XPU', 'CUDA', or 'NPU'."
+                )
             self.param_init_fn = lambda x: x.to_empty(device=device, recurse=False)
 
     @staticmethod


Test patch: