7% of training time with `Trainer` is spent on a single line: `and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))`
Author: umarbutler
Pull request: no
Labels: name: trainer

name: bug

### System Info

- `transformers` version: 4.44.2
- Platform: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
- Python version: 3.12.5
- Huggingface_hub version: 0.25.1
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.1+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: yes
- GPU type: NVIDIA GeForce RTX 4090

**AND**

- `transformers` version: 4.44.2
- Platform: Windows-11-10.0.22631-SP0
- Python version: 3.12.4
- Huggingface_hub version: 0.24.6
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.1+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: yes
- GPU type: NVIDIA GeForce RTX 4090

### Who can help?

_No response_

### 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

1. Create a simple training script that trains a model (in my case, a BERT model being trained on the masked language modelling objective) for a few hundred steps.
2. Profile that script with a profiler such as Austin.
3. Observe that 6-7% of time is spent on a single line in `transformers.trainer.Trainer._inner_training_loop()` at line 2284: `and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))`, as shown below:
![image](https://github.com/user-attachments/assets/7e26a499-3df9-4ae3-8fc8-db5a70cd00aa)


### Expected behavior

I'm not entirely sure if this is expected behaviour or if that much time is not actually being spent in that line or if there is a problem with my training script, but this does strike me as something a bit weird.

Comments:
Commenting out the check didn't seem to help my training speed very much on my proper training script (which uses a much much larger number of training steps) which suggests to me that this may only be a bottleneck in short runs?

Pinging @muellerzr and @SunMarc 

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.

Hey @umarbutler, sorry for the delay. As the issue you saw only relevant for short runs, I guess it should not be an major issue and I don't really think we can actually replace this condition with something faster. Thanks for reporting and sharing your observation, really appreciate it.