Check device map for saving tokenizer config on TPU (fix for issue #31971)
Author: ayukh
Pull request: yes

# What does this PR do?

Removes device_map from tokenizer_config when saving tokenizer on **TPU**. 

Fixes [#31971](https://github.com/huggingface/transformers/issues/31971)

**Why?**

As mentioned in the issue above, configs after training on TPU and non-TPU devices are different. After training on TPU, we get the following tokenizer config:
```
{'clean_up_tokenization_spaces': False, 'unk_token': '<unk>', 'bos_token': '<bos>', 'eos_token': '<eos>', 'pad_token': '[PAD]', 'add_bos_token': True, 'add_eos_token': False, 'additional_special_tokens': ['<start_of_turn>', '<end_of_turn>'], 'model_max_length': 1000000000000000019884624838656, 'sp_model_kwargs': {}, 'spaces_between_special_tokens': False, 'use_default_system_prompt': False, 'device_map': device(type='xla', index=0),...
```
While after training or non-TPU (e.g. GPU), we get the same config, but without `device_map`:
```
{'clean_up_tokenization_spaces': False, 'unk_token': '<unk>', 'bos_token': '<bos>', 'eos_token': '<eos>', 'pad_token': '<pad>', 'add_bos_token': True, 'add_eos_token': False, 'additional_special_tokens': ['<start_of_turn>', '<end_of_turn>'], 'model_max_length': 1000000000000000019884624838656, 'sp_model_kwargs': {}, 'spaces_between_special_tokens': False, 'use_default_system_prompt': False,...
```
`device_map` is not being used anywhere when loading tokenizer, therefore we can exclude it from config to support consistency across devices.

## 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?
- [x] 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. -> Discussed in the [issue](https://github.com/huggingface/transformers/issues/31971) listed above.
- [ ] 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? -> not sure if need to run big tests here as I am not adding any models; just a minor fix, which was tested on Google TPU Pods and fixes the error mentioned in the issue above.


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@SunMarc 


Comments:
> LGTM but how did device map end up here? AFAIK, device map should not go on the tokenizer, unless accelerate automatically dispatches the input ids?

To be honest, not sure why it appears there, maybe because it was a multi-TPU training and that code did not use accelerate, just SFTTrainer.
I should note that my experience with using HF tools on Colab TPU and big GCP TPU Pods is quite different as accelerate does not recognize TPU device in case with VMs even.

I mostly see this:
> tokenizer = AutoTokenizer.from_pretrained(base_model, token=token, return_tensors="pt", device_map=device)

passing device map to the tokenizer init

> passing device map to the tokenizer init

Reuploaded comment w/o edits. device_map appears if we pass `device_map` into tokenizer and call `tokenizer.save_pretrained()` on both TPU and non-TPU. When saving only tokenizer directly, device_map is used as 'cuda'/'cpu' in string format, but for TPU it is being saved a class. I haven't found anywhere in the source code where device_map is being used from tokenizer config, so do we need it anywhere? 

Also apologies for multiple edits before, I figured out which command saves device_map and which doesn't🤷
On non-TPU devices device_map using accelerate also looks ok, but I can't test it for TPU since it doesn't work properly on Google TPU Pod at the moment.

Not sure what's the best solution, we can raise and error if device map is passed as it seems recurring 

Potentially can just use a check for `device_map` and then throw an error if it is not in str format:
```
if "device_map" in tokenizer_config and not isinstance(tokenizer_config["device_map"], str):
    tokenizer_device = tokenizer_config["device_map"]
    raise TypeError(f"'device_map' is not in str type, got {type(tokenizer_device)}. Please set the value in str format (e.g. 'auto') or do not pass it into tokenizer (for TPU).")
```
This would throw an error for any non-str device, like xm.xla_device()/torch.device, which would be in torch.device class format.

We could, but overall tokenizer should not be related to any device! SO better just removing it and not serializing

> We could, but overall tokenizer should not be related to any device! SO better just removing it and not serializing

I agree, initially I wasn't sure if device is used anywhere for tokenizer after checking the source code. Then we can keep the before proposed code, where we just remove the device using:
```
if "device_map" in tokenizer_config:
    tokenizer_config.pop("device_map")
```
Does this look like a good solution to the issue?

yeah sure!