Repository: huggingface/transformers (python)

Issue:
[Bug] transformers `TPU` support broken on `v4.45.0`
### System Info

transformers: v4.45.0 and up (any of v4.45.0 / v4.45.1 / v4.45.2)
accelerate: v1.0.1 (same result on v0.34.2)

### Who can help?

trainer experts: @muellerzr @SunMarc
accelerate expert: @muellerzr
text models expert: @ArthurZucker 
Thank you guys!

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

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

### Reproduction

Minimal working code is [Here](https://gist.github.com/steveepreston/acd125a08214c631ba8389eb61a13798). Code follows [GoogleCloudPlatform example](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/tpuv5e_llama2_pytorch_finetuning_and_serving.ipynb)

on TPU VM, train done like a charm on transformers from v4.43.1 to v4.44.2, but when upgrading to any of v4.45.0 / v4.45.1 / v4.45.2 it throws this Error: `RuntimeError: There are currently no available devices found, must be one of 'XPU', 'CUDA', or 'NPU'.`

&nbsp;

**Error Traceback:**

General traceback is: callling `SFTTrainer()` >  `self.accelerator = Accelerator(**args)` (transformers/trainer.py)

<details>
  <summary>Click here to Show Full Error Traceback</summary>

  ```python
  ---------------------------------------------------------------------------
  RuntimeError                              Traceback (most recent call last)
  Cell In[48], line 4
        1 from trl import SFTTrainer
        2 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
  ----> 4 trainer = SFTTrainer(
        5     model=base_model,
        6     train_dataset=data,
        7     args=TrainingArguments(
        8         per_device_train_batch_size=BATCH_SIZE,  # This is actually the global batch size for SPMD.
        9         num_train_epochs=1,
       10         max_steps=-1,
       11         output_dir="/output_dir",
       12         optim="adafactor",
       13         logging_steps=1,
       14         dataloader_drop_last = True,  # Required for SPMD.
       15         fsdp="full_shard",
       16         fsdp_config=fsdp_config,
       17     ),
       18     peft_config=lora_config,
       19     dataset_text_field="quote",
       20     max_seq_length=max_seq_length,
       21     packing=True,
       22 )
  
  File /usr/local/lib/python3.10/site-packages/huggingface_hub/utils/_deprecation.py:101, in _deprecate_arguments.<locals>._inner_deprecate_positional_args.<locals>.inner_f(*args, **kwargs)
       99         message += "\n\n" + custom_message
      100     warnings.warn(message, FutureWarning)
  --> 101 return f(*args, **kwargs)
  
  File /usr/local/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:401, in SFTTrainer.__init__(self, model, args, data_collator, train_dataset, eval_dataset, tokenizer, model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics, peft_config, dataset_text_field, packing, formatting_func, max_seq_length, infinite, num_of_sequences, chars_per_token, dataset_num_proc, dataset_batch_size, neftune_noise_alpha, model_init_kwargs, dataset_kwargs, eval_packing)
      395 if tokenizer.padding_side is not None and tokenizer.padding_side != "right":
      396     warnings.warn(
      397         "You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to "
      398         "overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code."
      399     )
  --> 401 super().__init__(
      402     model=model,
      403     args=args,
      404     data_collator=data_collator,
      405     train_dataset=train_dataset,
      406     eval_dataset=eval_dataset,
      407     tokenizer=tokenizer,
      408     model_init=model_init,
      409     compute_metrics=compute_metrics,
      410     callbacks=callbacks,
      411     optimizers=optimizers,
      412     preprocess_logits_for_metrics=preprocess_logits_for_metrics,
      413 )
      415 # Add tags for models that have been loaded with the correct transformers version
      416 if hasattr(self.model, "add_model_tags"):
  
  File /usr/local/lib/python3.10/site-packages/transformers/trainer.py:411, in Trainer.__init__(self, model, args, data_collator, train_dataset, eval_dataset, tokenizer, model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)
      408 self.deepspeed = None
      409 self.is_in_train = False
  --> 411 self.create_accelerator_and_postprocess()
      413 # memory metrics - must set up as early as possible
      414 self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
  
  File /usr/local/lib/python3.10/site-packages/transformers/trainer.py:4858, in Trainer.create_accelerator_and_postprocess(self)
     4855     args.update(accelerator_config)
     4857 # create accelerator object
  -> 4858 self.accelerator = Accelerator(**args)
     4859 # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag
     4860 self.gather_function = self.accelerator.gather_for_metrics
  
  File /usr/local/lib/python3.10/site-packages/accelerate/accelerator.py:349, in Accelerator.__init__(self, device_placement, split_batches, mixed_precision, gradient_accumulation_steps, cpu, dataloader_config, deepspeed_plugin, fsdp_plugin, megatron_lm_plugin, rng_types, log_with, project_dir, project_config, gradient_accumulation_plugin, step_scheduler_with_optimizer, kwargs_handlers, dynamo_backend, deepspeed_plugins)
      345         raise ValueError(f"FSDP requires PyTorch >= {FSDP_PYTORCH_VERSION}")
      347 if fsdp_plugin is None:  # init from env variables
      348     fsdp_plugin = (
  --> 349         FullyShardedDataParallelPlugin() if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" else None
      350     )
      351 else:
      352     if not isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
  
  File <string>:21, in __init__(self, sharding_strategy, backward_prefetch, mixed_precision_policy, auto_wrap_policy, cpu_offload, ignored_modules, state_dict_type, state_dict_config, optim_state_dict_config, limit_all_gathers, use_orig_params, param_init_fn, sync_module_states, forward_prefetch, activation_checkpointing, cpu_ram_efficient_loading, transformer_cls_names_to_wrap, min_num_params)
  
  File /usr/local/lib/python3.10/site-packages/accelerate/utils/dataclasses.py:1684, in FullyShardedDataParallelPlugin.__post_init__(self)
     1682     device = torch.xpu.current_device()
     1683 else:
  -> 1684     raise RuntimeError(
     1685         "There are currently no available devices found, must be one of 'XPU', 'CUDA', or 'NPU'."
     1686     )
     1687 # Create a function that will be used to initialize the parameters of the model
     1688 # when using `sync_module_states`
     1689 self.param_init_fn = lambda x: x.to_empty(device=device, recurse=False)
  
  RuntimeError: There are currently no available devices found, must be one of 'XPU', 'CUDA', or 'NPU'.
  ```
</details>

&nbsp;

**My observation and guess**

I tested multiple times, and can confirm that this error is Directly Caused by only changing version of `transformers`. Therefore `accelerate` version was fixed during all runs, my guess is something changed on `v4.45.0` (maybe on `trainer.py`) that affects `args` in the `self.accelerator = Accelerator(**args)`, so that error will raised by `accelerate` .

### Expected behavior

my guess: `args` corrected and `self.accelerator = Accelerator(**args)` called correctly. so `accelerate` can work on `TPU`.

Pull request description:
Revert `accelerate` error caused by `46d09af`
# What does this PR do?

Fixes #34176 

Starting from `v4.45.0` Trainer was not able to train on TPU VMs.
so i tested commits between `v4.44.2` and `v4.45.0` one by one, then found this 1 line change commit [46d09af](https://github.com/huggingface/transformers/commit/46d09af4fc62db59cb836099d17a2c6102b39835) causes this error.
You can see full test result on #34176.


## Before submitting
- [x] This PR fixes a bug on code, not a typo
- [x] Read Pull Request Guideline
- [x] Discussed before on github issue.
- [x] No Document update is required. 1 line change.
- [x] Did you write any new necessary tests?


## Reviewers
Trainer: @muellerzr and @SunMarc
Deepspeed: HF Trainer/Accelerate: @muellerzr

Patch:
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index 6d71b754d6..86ea702dd9 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -1419,7 +1419,7 @@ class GenerationMixin:
         input_ids_length,
         inputs_tensor,
     ):
-        """Prepared max and min length in generaion configs to avoid clashes between similar attributes"""
+        """Prepared max and min length in generation configs to avoid clashes between similar attributes"""
 
         if generation_config.max_new_tokens is not None:
             if not has_default_max_length and generation_config.max_length is not None:
@@ -1662,7 +1662,7 @@ class GenerationMixin:
         device: torch.device,
     ) -> bool:
         """
-        Prepares the cache for generation (if applicable), given `generate`'s paramaterization. If a cache is
+        Prepares the cache for generation (if applicable), given `generate`'s parameterization. If a cache is
         instantiated, writes it to `model_kwargs`, under the name expected by the model.
         """
 
@@ -1925,7 +1925,7 @@ class GenerationMixin:
                 deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`.
             assistant_model (`PreTrainedModel`, *optional*):
                 An assistant model that can be used to accelerate generation. The assistant model must have the exact
-                same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistent model
+                same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistant model
                 is much faster than running generation with the model you're calling generate from. As such, the
                 assistant model should be much smaller.
             streamer (`BaseStreamer`, *optional*):
@@ -2442,7 +2442,15 @@ class GenerationMixin:
         # replace bos with pad to not condition healing on it
         input_ids = torch.where(input_ids == bos_token_id, pad_token_id, input_ids)
 
+        """
+        the latter code assumes the input_ids is not empty,
+        input_id has to be checked if contains elements
+		"""
+        if input_ids.numel() == 0:
+            return input_ids
+
         tail_ids = input_ids[:, -1].tolist()
+
         space_tok = tokenizer.convert_ids_to_tokens(tokenizer.convert_tokens_to_ids(" "))[0]
         # tail tokens are used for a prefix search, thus, whitespaces are replaced with
         # their tokenization (e.g. 'Ġ') to enable search for tokens prefixed with a whitespace
@@ -2454,7 +2462,14 @@ class GenerationMixin:
                 continue  # skip empty sequences (all pad ids)
 
             # apply bias for alternatives (extensions) to the tail token
-            seq_bias = {(alt_tok,): 10.0 for alt_tok in vocab_trie.values(prefix=tail_tok)}
+            """
+            seq_bias key has to be tuple with int so have to use
+            tokenizer function to convert str to int
+			"""
+            seq_bias = {
+                (tokenizer.convert_tokens_to_ids(alt_tok),): 10.0 for alt_tok in vocab_trie.extensions(prefix=tail_tok)
+            }
+
             if len(seq_bias) == 1:
                 continue  # skip if there are no token alternatives to heal with
 
@@ -2463,6 +2478,14 @@ class GenerationMixin:
             generation_config.update(sequence_bias=seq_bias)
 
             trimmed_ids = batch_ids[:-1]
+
+            """
+            the latter code assumes trimmed_ids is not empty
+            so have to check the its element count
+			"""
+            if trimmed_ids.numel() == 0:
+                continue
+
             # if the prompt is a single (non-pad) token, regenerate from bos
             if len(batch_ids[batch_ids != pad_token_id]) == 1:
                 trimmed_ids[-1] = bos_token_id
@@ -2915,7 +2938,7 @@ class GenerationMixin:
                     output_attentions=output_attentions,
                 )
 
-            # This is essential to avoid having a last reference to the big past K-V and double the necesary memory
+            # This is essential to avoid having a last reference to the big past K-V and double the necessary memory
             # in the next loop
             del next_model_inputs
 
@@ -3658,7 +3681,7 @@ class GenerationMixin:
             )
 
         # initialise score of first beam of each group with 0 and the rest with -1e9. This ensures that the beams in
-        # the same group don't produce same tokens everytime.
+        # the same group don't produce same tokens every time.
         beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
         beam_scores[:, ::num_sub_beams] = 0
         beam_scores = beam_scores.view((batch_size * num_beams,))
diff --git a/src/transformers/tokenization_utils.py b/src/transformers/tokenization_utils.py
index df13a029a6..d2433868cf 100644
--- a/src/transformers/tokenization_utils.py
+++ b/src/transformers/tokenization_utils.py
@@ -316,6 +316,9 @@ class ExtensionsTrie(Trie):
         """
         node = self.data
         for char in token:
+            if char not in node:
+                break
+
             node = node[char]
         return node
 
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
index a2d83b2915..485610dd9b 100644
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -1936,7 +1936,7 @@ class TrainingArguments:
                 warnings.warn("`--xla_fsdp_grad_ckpt` is useful only when `--xla` is set to true.")
 
         # accelerate integration for FSDP
-        if len(self.fsdp) > 0 and is_accelerate_available("0.28.0"):
+        if len(self.fsdp) > 0 and not self.fsdp_config["xla"]:
             os.environ["ACCELERATE_USE_FSDP"] = "true"
             from accelerate.utils.constants import (
                 FSDP_AUTO_WRAP_POLICY,


Test patch:
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
index 5165e43c09..6766fa22b9 100644
--- a/tests/generation/test_utils.py
+++ b/tests/generation/test_utils.py
@@ -28,7 +28,6 @@ from transformers import AutoConfig, is_torch_available, pipeline, set_seed
 from transformers.testing_utils import (
     is_flaky,
     require_accelerate,
-    require_auto_gptq,
     require_optimum_quanto,
     require_torch,
     require_torch_gpu,
@@ -3912,11 +3911,6 @@ class GenerationIntegrationTests(unittest.TestCase, GenerationIntegrationTestsMi
 class TokenHealingTestCase(unittest.TestCase):
     @parameterized.expand(
         [
-            (
-                "square_bracket",
-                'An example ["like this"] and another example [',
-                'An example ["like this"] and another example ["',
-            ),
             ("url", 'The link is <a href="http:', 'The link is <a href="http://'),
             # aggressive_healing: "http" shouldn't be replaced with "https"
             ("aggressive_healing", 'The link is <a href="http', 'The link is <a href="http'),
@@ -3926,9 +3920,8 @@ class TokenHealingTestCase(unittest.TestCase):
             ("empty_prompt", "", ""),
         ]
     )
-    @require_auto_gptq
     def test_prompts(self, name, input, expected):
-        model_name_or_path = "TheBloke/deepseek-llm-7B-base-GPTQ"
+        model_name_or_path = "distilbert/distilgpt2"
         tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
         completion_model = AutoModelForCausalLM.from_pretrained(
             model_name_or_path,
@@ -3937,9 +3930,16 @@ class TokenHealingTestCase(unittest.TestCase):
             revision="main",
             use_cache=True,
         )
+
+        """
+        tokenizer.pad_token value can be empty but it is required in the latter codes
+        so assigned it here with eos_token
+		"""
+        tokenizer.pad_token = tokenizer.eos_token
+
         input_ids = tokenizer(input, return_tensors="pt").input_ids.to(completion_model.device)
 
-        healed_ids = completion_model.heal_tokens(input_ids)
+        healed_ids = completion_model.heal_tokens(input_ids, tokenizer=tokenizer)
         predicted = tokenizer.decode(healed_ids[0], skip_special_tokens=True)
 
         self.assertEqual(predicted, expected)
diff --git a/tests/utils/test_tokenization_utils.py b/tests/utils/test_tokenization_utils.py
index edd580d887..3600be91a7 100644
--- a/tests/utils/test_tokenization_utils.py
+++ b/tests/utils/test_tokenization_utils.py
@@ -325,8 +325,9 @@ class ExtensionsTrieTest(unittest.TestCase):
     def test_no_extension_match(self):
         trie = ExtensionsTrie()
         # Test searching for a prefix that doesn't match any key
-        with self.assertRaises(KeyError):
-            trie.extensions("unknown")
+        values = trie.extensions("unknown")
+
+        self.assertEqual(len(values), 0)
 
     def test_update_value(self):
         trie = ExtensionsTrie()