Learning rate is given as tensor, cannot serialize TrainerState in order to save checkpoint.

[kmckiern]

System Info

• transformers version: 4.21.1
• Platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.35
• Python version: 3.10.5
• Huggingface_hub version: 0.8.1
• PyTorch version (GPU?): 1.12.1.post200 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

Who can help?

@sgugger

Information

• The official example scripts
• 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

1. use Adafactor optimizer and schedule

optimizer = Adafactor(
    model.parameters(),
    relative_step=True,
    warmup_init=True,
)
lr_scheduler = AdafactorSchedule(optimizer)

2. save a checkpoint

Expected behavior

When using an AdafactorSchedule I can't use the Trainer class to save a checkpoint.

It breaks here since the learning rate attached to the TrainerState is given by a tensor and tensors are not JSON serializable.

I dropped a breakpoint at this line and took a look at my TrainerState:

In [4]: self.log_history[0]['learning_rate']
Out[4]: tensor(0.0001, device='cuda:0')

Expected behavior is that the learning rate attached to the log history would be given by a float and would therefore be JSON serializable.