tune save checkpoint throwing error due to float32
Author: BramVanroy
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.21.0.dev0
- Platform: Linux-5.13.0-27-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.0+cu113 (True)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no, although Ray tune runs in parallel

### Who can help?

@richardliaw @amogkam

### Information

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

### Tasks

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

### Reproduction

I run a hyperparameter search with ray[tune] which consists of these parts:

```python
hp_space = {
    "num_train_epochs": tune.choice([1, 2, 3, 4]),
}

scheduler = PB2(
    metric="eval_f1",
    mode="max",
    hyperparam_bounds={
        "weight_decay": [0.0, 0.3],  # default (in transformers): 0.
        "learning_rate": [1e-4, 1e-5],
        "gradient_accumulation_steps": [4, 8],
        "adam_epsilon": [1e-07, 1e-9],  # default: 1e-8
        "adam_beta1": [0.85, 0.9999],  # default: 0.9
        "adam_beta2": [0.95, 0.9999],  # default: 0.999
    },
)

resources_per_trial = {"cpu": min(4, (os.cpu_count() - 1) // device_count), "gpu": 1}

best_params = trainer.hyperparameter_search(hp_space=lambda _: hpspace,
                                            backend="ray",
                                            n_trials=8,
                                            resources_per_trial=resources_per_trial,
                                            keep_checkpoints_num=1,
                                            scheduler=scheduler,
                                            compute_objective=compute_objective)
```

However, after the processes have been running for a long time (async 4x V100), I get the following error trace:

```
ray::ImplicitFunc.train()ESC[39m (pid=3354338, ip=157.193.228.18, repr=_objective)
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/trainable.py", line 360, in train
    result = self.step()
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/function_runner.py", line 404, in step
    self._report_thread_runner_error(block=True)
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/function_runner.py", line 574, in _report_thread_runner_error
    raise e
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/function_runner.py", line 277, in run
    self._entrypoint()
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/function_runner.py", line 349, in entrypoint
    return self._trainable_func(
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/function_runner.py", line 645, in _trainable_func
    output = fn()
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/integrations.py", line 288, in dynamic_modules_import_trainable
    return trainable(*args, **kwargs)
  File "/home/bram/.local/share/virtualenvs/transformers-finetuner-rzmJjOSV/lib/python3.8/site-packages/ray/tune/utils/trainable.py", line 410, in inner
    trainable(config, **fn_kwargs)
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/integrations.py", line 189, in _objective
    local_trainer.train(resume_from_checkpoint=checkpoint, trial=trial)
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/trainer.py", line 1410, in train
    return inner_training_loop(
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/trainer.py", line 1729, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/trainer.py", line 1914, in _maybe_log_save_evaluate
    self._report_to_hp_search(trial, epoch, metrics)
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/trainer.py", line 1153, in _report_to_hp_search
    self._tune_save_checkpoint()
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/trainer.py", line 1165, in _tune_save_checkpoint
    self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
  File "/home/bram/Python/projects/transformers-finetuner/transformers/src/transformers/trainer_callback.py", line 97, in save_to_json
    json_string = json.dumps(dataclasses.asdict(self), indent=2, sort_keys=True) + "\n"
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/__init__.py", line 234, in dumps
    return cls(
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/encoder.py", line 201, in encode
    chunks = list(chunks)
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/encoder.py", line 431, in _iterencode
    yield from _iterencode_dict(o, _current_indent_level)
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/encoder.py", line 405, in _iterencode_dict
    yield from chunks
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/encoder.py", line 405, in _iterencode_dict
    yield from chunks
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/encoder.py", line 438, in _iterencode
    o = _default(o)
  File "/home/bram/.pyenv/versions/3.8.10/lib/python3.8/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type float32 is not JSON serializable
```

This occurs at around step 1840/2500. I do not know if it is relevant, but I am also running in `fp16`. If I had to guess, I'd think that during the duping of the [TrainerState](https://github.com/huggingface/transformers/blob/692e61e91a0b83f5b847902ed619b7c74c0a5dda/src/transformers/trainer_callback.py#L97), one of the [trial_params](https://github.com/huggingface/transformers/blob/692e61e91a0b83f5b847902ed619b7c74c0a5dda/src/transformers/trainer_callback.py#L89) was a np/torch float32 rather than a Python primitive, which could not be serialized. It is unclear to me why this would only happen already far into the training, though. Maybe it's a nan, or another kind of overflow of some kind?

### Expected behavior

No errors. It would also be nice if the error message could tell us which key is causing this issue, but I am not sure how feasible that is.

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Bump

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I am getting the same error almost a year later. It seems that noone is using PB2 with transformers...

> I am getting the same error almost a year later. It seems that noone is using PB2 with transformers...

me too