Set non_blocking=True When moving data from the CPU to the GPU

[Hukongtao]

System Info

No Need

Who can help?

No response

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

I used the transformer's Trainer to train the model, but used my own Dataloader.
Then I used pytorch profile to check my training performance and found that the CPU execution time accounted for a high proportion

After a period of investigation, it was found that the non_blocking was not set when the data was transferred from the CPU to the GPU.
https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3625-L3631
The modified code is:

kwargs = {"device": self.args.device，“non_blocking”:  True}

Then I re-profiled my code and the results were as follows：

You can see that the performance has been greatly improved.

I'm not sure if this is a bug in the code or a problem with the way I'm using it.

But there is no doubt that setting non_blocking=True has brought a great performance improvement to my training.

Looking forward to your reply

Expected behavior

No Need

[Hukongtao]

If this is a bug, do you have plans to fix it?

[Rocketknight1]

cc @muellerzr @SunMarc

[Hukongtao]

I submitted a #36408, can you check it for me? @muellerzr @SunMarc