Changing `on_step` in `self.log` causes `batch_to_device` to change

[davidgill97]

Bug description

There seem to be a bottleneck to the training pipeline, so i tried profiling and found that setting on_step=True in self.log causes progress bar to take very long. This makes sense, as logging to progress bar requires calling .item() on GPU tensor, causing cudaStreamSynchronize.

FIT Profiler Report
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|  Action                                                                                                                                                                     	|  Mean duration (s)	|  Num calls      	|  Total time (s) 	|  Percentage %   	|
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|  Total                                                                                                                                                                      	|  -              	|  32523          	|  164.08         	|  100 %          	|
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|  run_training_epoch                                                                                                                                                         	|  8.0065         	|  12             	|  96.078         	|  58.556         	|
|  [Callback]ExportCallback.on_fit_end                                                                                                                                        	|  62.609         	|  1              	|  62.609         	|  38.158         	|
|  [Callback]TQDMProgressBar.on_train_batch_end                                                                                                                               	|  0.24358        	|  144            	|  35.075         	|  21.377         	|
|  [Callback]ModelCheckpoint{'monitor': 'valid/macro_f1', 'mode': 'max', 'every_n_train_steps': 0, 'every_n_epochs': 1, 'train_time_interval': None}.on_validation_end        	|  1.3881         	|  12             	|  16.657         	|  10.152         	|
|  [Strategy]SingleDeviceStrategy.batch_to_device                                                                                                                             	|  0.0040199      	|  1056           	|  4.245          	|  2.5872         	|

But when i change on_step=False, though TQDMProgressBar.on_train_batch_end takes much less time, transfer_batch_to_device takes almost 10 times more. This behavior is completely baffling, and i don't know what to suspect.

Is this expected and what's causing this behavior?

FIT Profiler Report
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|  Action                                                                                                                                                                     	|  Mean duration (s)	|  Num calls      	|  Total time (s) 	|  Percentage %   	|
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|  Total                                                                                                                                                                      	|  -              	|  32523          	|  163.98         	|  100 %          	|
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|  run_training_epoch                                                                                                                                                         	|  7.9402         	|  12             	|  95.282         	|  58.104         	|
|  [Strategy]SingleDeviceStrategy.batch_to_device                                                                                                                             	|  0.036941       	|  1056           	|  39.01          	|  23.789         	|
|  [LightningModule]SingleLabelClassification.transfer_batch_to_device                                                                                                        	|  0.036852       	|  1056           	|  38.916         	|  23.732         	|
|  [Callback]ModelCheckpoint{'monitor': 'valid/macro_f1', 'mode': 'max', 'every_n_train_steps': 0, 'every_n_epochs': 1, 'train_time_interval': None}.on_validation_end        	|  1.3867         	|  12             	|  16.64          	|  10.147         	|

[SkafteNicki]

Hi @davidgill97, I think I can explain this:

First off, it's important to note that the overall training time remains constant regardless of on_step=True/False, indicating that the same operations are occurring but in different parts of the pipeline. This behavior is likely due to CUDA stream synchronization.

Explanation:

• on_step=True: Logging at every step calls .item() on GPU tensors, forcing synchronization between the CPU and GPU. This causes the GPU to wait for all queued operations to complete, leading to the observed slowdown in TQDMProgressBar.on_train_batch_end.
• on_step=False: Logging happens less frequently, deferring synchronization. This shifts the synchronization overhead to other operations, such as transfer_batch_to_device, which now appears slower.

Thus, your code need to synchronize at some point and depending on the on_step argument this happens in different parts of the code.

[davidgill97]

Hi @davidgill97, I think I can explain this:

First off, it's important to note that the overall training time remains constant regardless of on_step=True/False, indicating that the same operations are occurring but in different parts of the pipeline. This behavior is likely due to CUDA stream synchronization.

Explanation:

• on_step=True: Logging at every step calls .item() on GPU tensors, forcing synchronization between the CPU and GPU. This causes the GPU to wait for all queued operations to complete, leading to the observed slowdown in TQDMProgressBar.on_train_batch_end.
• on_step=False: Logging happens less frequently, deferring synchronization. This shifts the synchronization overhead to other operations, such as transfer_batch_to_device, which now appears slower.

Thus, your code need to synchronize at some point and depending on the on_step argument this happens in different parts of the code.

Thank you for the detailed explanation, now it makes a lot more sense! Is there a recipe for minimizing CPU-GPU synchronization during training? It is quite difficult to decide what to control for maximum training throughput, compared to vanilla pytorch where i can control exactly when to synchronize..

[SkafteNicki]

You are welcome. I am not really sure if we have a specific recipe/doc for how to minimize CPU-GPU sync. You can look at these pages for general advice to improve training time:
https://lightning.ai/docs/pytorch/stable/levels/intermediate_level_13.html

One advice is that every time you want to print, log, save anything, that involves a CPU-GPU sync. You can try running the trainer in barebone mode:

Trainer(barebones=True)

which will disable a lot of lightning stuff and see what kind of performance you can get for a comparison of how your current setup is performing (without having to rewrite everything to vanilla pytorch).

[davidgill97]

You are welcome. I am not really sure if we have a specific recipe/doc for how to minimize CPU-GPU sync. You can look at these pages for general advice to improve training time: https://lightning.ai/docs/pytorch/stable/levels/intermediate_level_13.html

One advice is that every time you want to print, log, save anything, that involves a CPU-GPU sync. You can try running the trainer in barebone mode:

Trainer(barebones=True)
which will disable a lot of lightning stuff and see what kind of performance you can get for a comparison of how your current setup is performing (without having to rewrite everything to vanilla pytorch).

I tested barebones, and found out that training is twice as fast... so i now need to figure out which callback/logger is causing the slowdown. Thank you for your help!