Cannot replicate if number of devices (1) is different from 8 in Kaggle TPUs
Author: alisusername
Pull request: no
Labels: name: TPU

name: Accelerate

### System Info

Kaggle TPU 3-8
Python 3.10.13
transformers.__version__ '4.39.1'
accelerate.__version__ '0.28.0'
torch.__version__ '2.2.1+cu121'
torch_xla.__version__ '2.2.0+libtpu'


### Who can help?

trainer: @muellerzr and @pacman100
deepspeed: HF Trainer/Accelerate: @pacman100

### Information

- [ ] The official example scripts
- [X] 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

Hello
using transformers.TrainingArguments returns an error in Kaggle tpu notebooks regardless of given parameters 
I can't figure out what is causing this! 

Example Notebook: https://www.kaggle.com/code/azasazas/transformers/notebook

The error:
```python
E0328 12:14:52.355234628      77 oauth2_credentials.cc:238]            oauth_fetch: UNKNOWN:C-ares status is not ARES_SUCCESS qtype=A name=metadata.google.internal. is_balancer=0: Domain name not found {created_time:"2024-03-28T12:14:52.355218637+00:00", grpc_status:2}
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[7], line 1
----> 1 args = transformers.TrainingArguments(
      2     output_dir="test",
      3     per_device_train_batch_size=8,
      4     logging_steps=10,
      5     gradient_accumulation_steps=1,
      6     num_train_epochs=1,
      7     weight_decay=0.1,
      8     warmup_steps=1_000,
      9     save_total_limit=1,
     10     save_strategy='steps',
     11     lr_scheduler_type="cosine",
     12     learning_rate=5e-4,
     13     tpu_num_cores=8,
     14     save_steps=10000,
     15     bf16=True,
     16     push_to_hub=True,
     17     max_steps=100000,
     18     report_to="none",
     19     hub_strategy="every_save",
     20     ignore_data_skip=False,
     21     resume_from_checkpoint=False,
     22     adafactor=True,
     23 )

File <string>:124, in __init__(self, output_dir, overwrite_output_dir, do_train, do_eval, do_predict, evaluation_strategy, prediction_loss_only, per_device_train_batch_size, per_device_eval_batch_size, per_gpu_train_batch_size, per_gpu_eval_batch_size, gradient_accumulation_steps, eval_accumulation_steps, eval_delay, learning_rate, weight_decay, adam_beta1, adam_beta2, adam_epsilon, max_grad_norm, num_train_epochs, max_steps, lr_scheduler_type, lr_scheduler_kwargs, warmup_ratio, warmup_steps, log_level, log_level_replica, log_on_each_node, logging_dir, logging_strategy, logging_first_step, logging_steps, logging_nan_inf_filter, save_strategy, save_steps, save_total_limit, save_safetensors, save_on_each_node, save_only_model, no_cuda, use_cpu, use_mps_device, seed, data_seed, jit_mode_eval, use_ipex, bf16, fp16, fp16_opt_level, half_precision_backend, bf16_full_eval, fp16_full_eval, tf32, local_rank, ddp_backend, tpu_num_cores, tpu_metrics_debug, debug, dataloader_drop_last, eval_steps, dataloader_num_workers, dataloader_prefetch_factor, past_index, run_name, disable_tqdm, remove_unused_columns, label_names, load_best_model_at_end, metric_for_best_model, greater_is_better, ignore_data_skip, fsdp, fsdp_min_num_params, fsdp_config, fsdp_transformer_layer_cls_to_wrap, accelerator_config, deepspeed, label_smoothing_factor, optim, optim_args, adafactor, group_by_length, length_column_name, report_to, ddp_find_unused_parameters, ddp_bucket_cap_mb, ddp_broadcast_buffers, dataloader_pin_memory, dataloader_persistent_workers, skip_memory_metrics, use_legacy_prediction_loop, push_to_hub, resume_from_checkpoint, hub_model_id, hub_strategy, hub_token, hub_private_repo, hub_always_push, gradient_checkpointing, gradient_checkpointing_kwargs, include_inputs_for_metrics, fp16_backend, push_to_hub_model_id, push_to_hub_organization, push_to_hub_token, mp_parameters, auto_find_batch_size, full_determinism, torchdynamo, ray_scope, ddp_timeout, torch_compile, torch_compile_backend, torch_compile_mode, dispatch_batches, split_batches, include_tokens_per_second, include_num_input_tokens_seen, neftune_noise_alpha, optim_target_modules)

File /usr/local/lib/python3.10/site-packages/transformers/training_args.py:1551, in TrainingArguments.__post_init__(self)
   1545     if version.parse(version.parse(torch.__version__).base_version) == version.parse("2.0.0") and self.fp16:
   1546         raise ValueError("--optim adamw_torch_fused with --fp16 requires PyTorch>2.0")
   1548 if (
   1549     self.framework == "pt"
   1550     and is_torch_available()
-> 1551     and (self.device.type != "cuda")
   1552     and (self.device.type != "npu")
   1553     and (self.device.type != "xpu")
   1554     and (get_xla_device_type(self.device) not in ["GPU", "CUDA"])
   1555     and (self.fp16 or self.fp16_full_eval)
   1556 ):
   1557     raise ValueError(
   1558         "FP16 Mixed precision training with AMP or APEX (`--fp16`) and FP16 half precision evaluation"
   1559         " (`--fp16_full_eval`) can only be used on CUDA or NPU devices or certain XPU devices (with IPEX)."
   1560     )
   1562 if (
   1563     self.framework == "pt"
   1564     and is_torch_available()
   (...)
   1571     and (self.bf16 or self.bf16_full_eval)
   1572 ):

File /usr/local/lib/python3.10/site-packages/transformers/training_args.py:2027, in TrainingArguments.device(self)
   2023 """
   2024 The device used by this process.
   2025 """
   2026 requires_backends(self, ["torch"])
-> 2027 return self._setup_devices

File /usr/local/lib/python3.10/site-packages/transformers/utils/generic.py:63, in cached_property.__get__(self, obj, objtype)
     61 cached = getattr(obj, attr, None)
     62 if cached is None:
---> 63     cached = self.fget(obj)
     64     setattr(obj, attr, cached)
     65 return cached

File /usr/local/lib/python3.10/site-packages/transformers/training_args.py:1963, in TrainingArguments._setup_devices(self)
   1961     self._n_gpu = 1
   1962 else:
-> 1963     self.distributed_state = PartialState(
   1964         backend=self.ddp_backend, timeout=timedelta(seconds=self.ddp_timeout)
   1965     )
   1966     self._n_gpu = 1
   1967 if not is_sagemaker_mp_enabled():

File /usr/local/lib/python3.10/site-packages/accelerate/state.py:163, in PartialState.__init__(self, cpu, **kwargs)
    161 self.distributed_type = DistributedType.XLA
    162 self.device = xm.xla_device()
--> 163 xm.set_replication(self.device, [self.device])
    164 self.num_processes = xm.xrt_world_size()
    165 self.process_index = xm.get_ordinal()

File /usr/local/lib/python3.10/site-packages/torch_xla/core/xla_model.py:284, in set_replication(device, devices)
    282 devices = [str(x) for x in devices]
    283 if devices:
--> 284   replication_devices = xla_replication_devices(devices)
    285   torch_xla._XLAC._xla_set_replication_devices(replication_devices)
    286   devctx.device_index = devices.index(device)

File /usr/local/lib/python3.10/site-packages/torch_xla/core/xla_model.py:252, in xla_replication_devices(local_devices)
    249 kind_devices = get_xla_supported_devices(devkind=device_type)
    250 if len(kind_devices) != len(local_devices):
    251   # Replication can only happen among all devices of one kind.
--> 252   raise RuntimeError(
    253       'Cannot replicate if number of devices ({}) is different from {}'.
    254       format(len(local_devices), len(kind_devices)))
    255 replication_devices = []
    256 for device in torch_xla._XLAC._xla_get_all_devices():

RuntimeError: Cannot replicate if number of devices (1) is different from 8
```


### Expected behavior

I should be able to define the TrainingArguments,__ 

Comments:
I made another notebook in this one I didn't update any package and everything is from the Kaggle default environment
https://www.kaggle.com/code/azasazas/original-kaggle/notebook
I thought maybe this would be helpful

Hello
I have some more information to share I tried to run and train a LLama sample model with transformers pytorch and torch xla last version and the accelerate version 0.27.0 and 0.27.2 and it worked it only showed a warning that:
***
/usr/local/lib/python3.10/site-packages/torch/autograd/__init__.py:266: UserWarning: aten::reshape: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior. This behavior is deprecated and will be removed in a future version of PyTorch. If your operator is differentiable, please ensure you have registered an autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, DispatchKey::CompositeImplicitAutograd). If your operator is not differentiable, or to squash this warning and use the previous behavior, please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd. (Triggered internally at ../torch/csrc/autograd/autograd_not_implemented_fallback.cpp:63.)
  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the back
***
I don't know if it has any effect on training or not but the training process starts and finishes with accelerate versions 0.27.0 and 0.27.2, however, it's not working with accelerate 0.28.0

i have the same issue

Also cc @SunMarc !

Gentle ping @SunMarc 

Hey ! 
It seems to be some sort of regression within accelerate cc @muellerzr as well 
@carlesoctav @Aiyns do you confirm you get the same issue with latest accelerate (0.30.0)?

Hi @Aiyns, @muellerzr did a few changes in xla recently such as this [PR](https://github.com/huggingface/accelerate/pull/2634). Could you try with the latest accelerate ? 

@SunMarc Yeah it works now with:
transformers == 4.41.1
accelerate == 0.30.1
torch == 2.2.1
torch_xla == 2.2.0+libtpu
, Thank you, everybody

Closing this issue since it is solved !