"TypeError: Object of type device is not JSON serializable" when saving the model on TPU
Author: ayukh
Pull request: no

### System Info

Running on v4-64 Google TPU Pod:
- torch~=2.3.1
- torch_xla[tpu]~=2.3.0
- accelerate=0.32.1
- transformers=4.42.4
CLI env output:
- `transformers` version: 4.42.4
- Platform: Linux-5.19.0-1022-gcp-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.32.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.1+cu121 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed

### Who can help?

@muellerz @SunMarc 

### Information

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

### Tasks

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

### Reproduction

I am trying to finetune Gemma/Llama models on Google TPU Pod (v4-64). When I save the model at the end of trainer or separately using trainer.save_model() I encounter the following error which happens when saving tokenizer:

```
Traceback (most recent call last):
  File "/home/user_name/train_llama.py", line 254, in <module>
    main()
  File "/home/user_name/train_llama.py", line 246, in main
    trainer.train()
  File "/home/user_name/.local/lib/python3.10/site-packages/trl/trainer/sft_trainer.py", line 440, in train
    output = super().train(*args, **kwargs)
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/trainer.py", line 1932, in train
    return inner_training_loop(
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2345, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2796, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2875, in _save_checkpoint
    self.save_model(output_dir, _internal_call=True)
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3395, in save_model
    self._save_tpu(output_dir)
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3472, in _save_tpu
    self.tokenizer.save_pretrained(output_dir)
  File "/home/user_name/.local/lib/python3.10/site-packages/transformers/tokenization_utils_base.py", line 2572, in save_pretrained
    out_str = json.dumps(tokenizer_config, indent=2, sort_keys=True, ensure_ascii=False) + "\n"
  File "/usr/lib/python3.10/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
  File "/usr/lib/python3.10/json/encoder.py", line 201, in encode
    chunks = list(chunks)
  File "/usr/lib/python3.10/json/encoder.py", line 431, in _iterencode
    yield from _iterencode_dict(o, _current_indent_level)
  File "/usr/lib/python3.10/json/encoder.py", line 405, in _iterencode_dict
    yield from chunks
  File "/usr/lib/python3.10/json/encoder.py", line 438, in _iterencode
    o = _default(o)
  File "/usr/lib/python3.10/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type device is not JSON serializable
wandb: \ 0.044 MB of 0.044 MB uploaded
```
### Code reproduction
Similar lunch process as in the example [here](https://github.com/huggingface/transformers/blob/350aed70765cac5acfa7a4cfac4d1e6e2dc0b19b/examples/legacy/seq2seq/finetune_tpu.sh) either with spawning or directly launching the training script.
```python
def main():
    device = xm.xla_device()
    print(device)
    max_seq_length = 1024 
    packing = True
    base_model = "google/gemma-2-9b"
    output_dir = f"{base_model}-lora" 
    fsdp_config = {"fsdp_transformer_layer_cls_to_wrap": [
            "GemmaDecoderLayer",
        ],
        "xla": True,
        "xla_fsdp_v2": True,
        "xla_fsdp_grad_ckpt": True} 
    
    model = AutoModelForCausalLM.from_pretrained(base_model, device_map=device, token=token)
    tokenizer = AutoTokenizer.from_pretrained(base_model, token=token, return_tensors="pt", device_map=device)
    tokenizer.pad_token = "[PAD]"
    tokenizer.padding_side = "left"
    lora_config = LoraConfig(
            r=32,
            lora_alpha=64,
            target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],#"all-linear",
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
        )

    model.enable_input_require_grads()
    prepare_model_for_kbit_training(model)
    peft_model = get_peft_model(model, lora_config)
    peft_model.print_trainable_parameters()

    val_set_size = 0.1

    dataset = load_dataset("Abirate/english_quotes")
    len(dataset["train"])
    train_val = dataset["train"].train_test_split(
        test_size=val_set_size, shuffle=True, seed=2024
    )
    train_data = train_val["train"].shuffle(seed=2024)
    val_data = train_val["test"].shuffle(seed=2024)

    
    # Set the training arguments
    args = TrainingArguments(
        output_dir = output_dir,
        # auto_find_batch_size=True,
        per_device_train_batch_size=1, # Batch size = per_device_train_batch_size * gradient_accumulation_steps * num_gpus (if using DataParallel)
        per_device_eval_batch_size=1, # Batch size = per_device_train_batch_size * gradient_accumulation_steps * num_gpus (if using DataParallel)
        eval_accumulation_steps=5,
        gradient_checkpointing=False,
        tpu_num_cores = num_cores,
        warmup_steps=100, 
        max_steps=3, #3
        learning_rate=1e-5,
        bf16=True,
        deepspeed=deepspeed_config,
        ddp_backend="mpi",
        logging_strategy="steps",
        logging_steps=2,
        optim="adamw_torch",
        logging_dir="/data/training_output",
        lr_scheduler_type="cosine",
        evaluation_strategy="steps",
        fsdp="full_shard",
        fsdp_config=fsdp_config,
        save_strategy="steps",
        report_to="wandb",
        eval_steps=5,
        dataloader_drop_last = True,
        save_steps=50,
        save_total_limit=4,
        load_best_model_at_end=True,
        save_on_each_node=True,
    )

    trainer = SFTTrainer(
            model=peft_model,
            args=args,
            train_dataset=train_data,
            eval_dataset=val_data,
            tokenizer=tokenizer,
            packing=packing,
            max_seq_length=max_seq_length,
            dataset_text_field="quote",
        )

    # Train the model
    trainer.train()
    trainer.save_model(f"{output_dir}/trainer")
    tokenizer.save_pretrained(output_dir)

def _mp_fn(index):
    # For xla_spawn (multi-host TPUs)
    args = parse_arguments()
    main(deepspeed_config=args.deepspeed_config, num_cores=args.tpu_num_cores)

if __name__ == "__main__":
    args = parse_arguments()
    main(deepspeed_config=args.deepspeed_config, num_cores=args.tpu_num_cores)
```

### Expected behavior

The error arises [here](https://github.com/huggingface/transformers/blob/350aed70765cac5acfa7a4cfac4d1e6e2dc0b19b/src/transformers/tokenization_utils_base.py#L2572). I tried somehow changing the device for model and tokenizer, but as it is configured internally inside Trainer I do not think I can fix it. I checked the source code and did not see device mentioned anywhere during config saving, not sure why this issue arises on TPU specifically.

### Expected behavior

- Both model and tokenizer trained on TPU are saved correctly without issues
- TPU device in configuration should be ignored when saving? In other cases I assume the code should be modified to handle TPU device

Comments:
Hi @ayukh, thanks for reporting ! Could you try to create a colab with the error since we don't have access to local TPU ? Could you also check what is inside of tokenizer_config with printing it ? (in `tokenizer.save_pretrained(output_dir)`)

> Hi @ayukh, thanks for reporting ! Could you try to create a colab with the error since we don't have access to local TPU ? Could you also check what is inside of tokenizer_config with printing it ? (in `tokenizer.save_pretrained(output_dir)`)

Hi @SunMarc, thanks for the quick reply! I am not able to create or connect colab for those TPU Pods in particular and Google Colab itself offers only v2-8 TPU (different type), which cannot fit models like Gemma2-9B. However, I can try get more details about tokenizer config in my case.

For some reason, the tokenizer is being saved when I call `trainer.save_model()`. I can get config from tokenizer itself using AutoConfig for gemma2-9b tokenizer:
```
Gemma2Config {
  "_name_or_path": "google/gemma-2-9b",
  "architectures": [
    "Gemma2ForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "attn_logit_softcapping": 50.0,
  "bos_token_id": 2,
  "cache_implementation": "hybrid",
  "eos_token_id": 1,
  "final_logit_softcapping": 30.0,
  "head_dim": 256,
  "hidden_act": "gelu_pytorch_tanh",
  "hidden_activation": "gelu_pytorch_tanh",
  "hidden_size": 3584,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "gemma2",
  "num_attention_heads": 16,
  "num_hidden_layers": 42,
  "num_key_value_heads": 8,
  "pad_token_id": 0,
  "query_pre_attn_scalar": 256,
  "rms_norm_eps": 1e-06,
  "rope_theta": 10000.0,
  "sliding_window": 4096,
  "sliding_window_size": 4096,
  "torch_dtype": "float32",
  "transformers_version": "4.42.4",
  "use_cache": true,
  "vocab_size": 256000
}
```
After training I can print tokenizer by itself as config is not added there:
```
GemmaTokenizerFast(name_or_path='google/gemma-2-9b', vocab_size=256000, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<bos>', 'eos_token': '<eos>', 'unk_token': '<unk>', 'pad_token': '[PAD]', 'additional_special_tokens': ['<start_of_turn>', '<end_of_turn>']}, clean_up_tokenization_spaces=False),  added_tokens_decoder={tokens dict}
```

> Could you also check what is inside of tokenizer_config with printing it ? (in tokenizer.save_pretrained(output_dir))

Sorry, I don't quite understand what do you mean by this, can you give me an example how can I retrieve config of tokenizer (passed into Trainer) after training?

From your traceback, i can see that the error happens in save_pretrained method of the tokenizer: 
```python
        with open(tokenizer_config_file, "w", encoding="utf-8") as f:
            out_str = json.dumps(tokenizer_config, indent=2, sort_keys=True, ensure_ascii=False) + "\n"
```
Could you print `tokenizer_config` to see what is stored inside and compare with a non tpu setup. LMK if this is clear now ! 

> Could you print tokenizer_config to see what is stored inside and compare with a non tpu setup. 

Now it is clear :) 

Ok, so the output is quite big, I have uploaded output from TPU [here](https://github.com/ayukh/hf_tpu/blob/main/tpu_log.txt) and from GPU run [here](https://github.com/ayukh/hf_tpu/blob/main/gpu_log.txt).

The only difference between them is the presence of `device_map` in TPU config: `'device_map': device(type='xla', index=0)`. I can create a pull request to handle this type of device, maybe even just keep it is string format as 'xla'? I am not sure whether this would be used anywhere when loading tokenizer as we can't pass device_map like when loading model - maybe it is better even to remove device_map from config knowing it is not saved in non-tpu case.

I have linked a PR which helped solve this issue in my case, let me know if it looks good.
On less related note, I would also like to clarify a minor thing since I haven't found the confirmation for this anywhere - when I use HuggingFace trainer for multi-host TPU training (imagine many hosts with 8 TPU each) I see several runs created on wandb (no. of runs = no. of hosts):
* does the Trainer save in the end the same file for the best model accross all hosts?
* does the Trainer require to replicate the whole model on each worker separately? From the training it is not really clear whether the hosts train separately by themselves or they host model file locally but training is managed on global level. I have checked source code and it looks like like option 2 hold, but I did not see `xmp.spawn()` used anywhere, which from my understanding handles this.

Thanks! 😃 

> I have linked a PR which helped solve this issue in my case, let me know if it looks good.
Nice ! 

cc @muellerzr for the questions

A similar issue occured when I used `accelerator` to save a tokenizer, i.e., there is a "device" key inside the tokenizer config that cannot be json serialized:

```python
{
    "do_lower_case": True,
    "unk_token": "[UNK]",
    "sep_token": "[SEP]",
    "pad_token": "[PAD]",
    "cls_token": "[CLS]",
    "mask_token": "[MASK]",
    "tokenize_chinese_chars": True,
    "strip_accents": None,
    "clean_up_tokenization_spaces": True,
    "device": device(type="cuda", index=0),
    "do_basic_tokenize": True,
    "model_max_length": 512,
    "never_split": None,
    "max_length": 64,
    "truncation_side": "right",
    "stride": 0,
    "truncation_strategy": "longest_first",
    "pad_token_type_id": 0,
    "padding_side": "right",
    "pad_to_multiple_of": None,
    "extra_special_tokens": {},
    "added_tokens_decoder": {
        0: {
            "content": "[PAD]",
            "single_word": False,
            "lstrip": False,
            "rstrip": False,
            "normalized": False,
            "special": True,
        },
        100: {
            "content": "[UNK]",
            "single_word": False,
            "lstrip": False,
            "rstrip": False,
            "normalized": False,
            "special": True,
        },
        101: {
            "content": "[CLS]",
            "single_word": False,
            "lstrip": False,
            "rstrip": False,
            "normalized": False,
            "special": True,
        },
        102: {
            "content": "[SEP]",
            "single_word": False,
            "lstrip": False,
            "rstrip": False,
            "normalized": False,
            "special": True,
        },
        103: {
            "content": "[MASK]",
            "single_word": False,
            "lstrip": False,
            "rstrip": False,
            "normalized": False,
            "special": True,
        },
    },
    "tokenizer_class": "BertTokenizer",
}
```