Print input / output / grad / loss at every step/epoch when training Transformers HuggingFace model
I'm working on HuggingFace Transformers and using toy example from here:
https://huggingface.co/transformers/custom_datasets.html#fine-tuning-with-trainer
What I actually need: ability to print input, output, grad and loss at every step.
It is trivial using Pytorch training loop, but it is not obvious using HuggingFace Trainer.
At the current moment I have next idea: create a CustomCallback like this:
class MyCallback(TrainerCallback):
    &quot;A callback that prints a grad at every step&quot;

    def on_step_begin(self, args, state, control, **kwargs):
        print(&quot;next step&quot;)
        print(kwargs['model'].classifier.out_proj.weight.grad.norm())

args = TrainingArguments(
    output_dir='test_dir',
    overwrite_output_dir=True,
    num_train_epochs=1,
    logging_steps=100,
    report_to=&quot;none&quot;,
    fp16=True,
    disable_tqdm=True,
)


trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    callbacks=[MyCallback],
)

trainer.train()

This way I can print grad and weights for any model layer.
But I still can't figure out how to print input/output (for example, I want to check them on nan) and loss?
P.S. I also read something about forward_hook but still can't find good code examples for it.

While using hooks and custom callbacks is the right way to solve the problem I find better solution - use built-in utility for finding nan/Inf in losses / weights / inputs / outputs:
https://huggingface.co/transformers/internal/trainer_utils.html#transformers.debug_utils.DebugUnderflowOverflow
since the 4.6.0 transformers has such option.
You can use it manually in forward function or just use additional option for TrainingArguments like this:
args = TrainingArguments(
    output_dir='test_dir',
    overwrite_output_dir=True,
    num_train_epochs=1,
    logging_steps=100,
    report_to=&quot;none&quot;,
    fp16=True,
    disable_tqdm=True,
    debug=&quot;debug underflow_overflow&quot;
)

