question: Access a registered buffer is very slow
context: Hello,
  I implemented MoCo in Pytorch lightning. I was surprised to see that my lightning version was slower than Pytorch's and I ran the profiler to check which function is slow. I can't share all my code but here are the relevant parts:
  class MoCoModel(LightningModule):
      def __init__(
          ...
      ) -> None:
          ...
  
          self.register_buffer('queue', torch.randn(queue.feature_dim, queue.size))
          self.queue = nn.functional.normalize(self.queue, dim=0)
          self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))
      
      @torch.no_grad()
      def _update_queue(self, x: Tensor) -> None:
          x = self.concat_all_gather_without_backprop(x)
  
          #batch_size = x.shape[0]
          batch_size = self._get_batch_size(x)
  
          # for simplicity
          ptr = self._get_ptr()
          #ptr = int(self.queue_ptr)
  
          self._assert(batch_size)
          #assert self.queue_size % batch_size == 0
  
          # replace the keys at ptr (dequeue and enqueue)
          x = self._transpose(x)
          self._assign_in_queue(x, ptr, batch_size)
          #self.queue[:, ptr: ptr + batch_size] = x.T
  
          # move pointer
          ptr = self._compute_ptr(ptr, batch_size)
          self._assign_ptr(ptr)
          #ptr =  (ptr + batch_size) % self.queue_size
      
      def _get_batch_size(self, x):
          return x.shape[0]
  
      def _get_ptr(self):
          return int(self.queue_ptr)
  
      def _assert(self, batch_size):
          assert self.queue_size % batch_size == 0
      
      def _assign_ptr(self, ptr):
          self.queue_ptr[0] = ptr
      
      def _compute_ptr(self, batch_size, ptr):
          return (ptr + batch_size) % self.queue_size
  
      def _transpose(self, x):
          return x.T
      
      def _assign_in_queue(self, x, ptr, batch_size):
          self.queue[:, ptr: ptr + batch_size] = x
  
      def training_step(self, batch):
          ...
          self._update_queue(k)
  Here is the output of running simple profiler:
  Action                             	|  Mean duration (s)	|Num calls      	|  Total time (s) 	|  Percentage %   	|
  --------------------------------------------------------------------------------------------------------------------------------------
  Total                              	|  -              	|_              	|  53.595         	|  100 %          	|
  --------------------------------------------------------------------------------------------------------------------------------------
  run_training_epoch                 	|  45.224         	|1              	|  45.224         	|  84.381         	|
  run_training_batch                 	|  0.21673        	|195            	|  42.262         	|  78.854         	|
  optimizer_step_with_closure_0      	|  0.20378        	|195            	|  39.738         	|  74.145         	|
  training_step_and_backward         	|  0.19978        	|195            	|  38.957         	|  72.688         	|
  model_forward                      	|  0.1909         	|195            	|  37.225         	|  69.457         	|
  training_step                      	|  0.19077        	|195            	|  37.201         	|  69.411         	|
  backward                           	|  0.0083673      	|195            	|  1.6316         	|  3.0443         	|
  on_train_batch_end                 	|  0.0077772      	|195            	|  1.5166         	|  2.8296         	|
  get_train_batch                    	|  0.0034326      	|196            	|  0.6728         	|  1.2553         	|
  fetch_next_train_batch             	|  0.0034203      	|196            	|  0.67037        	|  1.2508         	|
  zero_grad                          	|  0.00049274     	|195            	|  0.096084       	|  0.17928        	|
  configure_optimizers               	|  0.093719       	|1              	|  0.093719       	|  0.17486        	|
  training_batch_to_device           	|  0.00028381     	|195            	|  0.055342       	|  0.10326        	|
  on_train_batch_start               	|  0.00018134     	|195            	|  0.03536        	|  0.065977       	|
  on_train_start                     	|  0.033906       	|1              	|  0.033906       	|  0.063264       	|
  on_pretrain_routine_start          	|  0.006531       	|1              	|  0.006531       	|  0.012186       	|
  on_batch_start                     	|  3.062e-05      	|195            	|  0.0059708      	|  0.011141       	|
  on_after_backward                  	|  3.0163e-05     	|195            	|  0.0058817      	|  0.010974       	|
  on_before_optimizer_step           	|  2.989e-05      	|195            	|  0.0058285      	|  0.010875       	|
  on_batch_end                       	|  2.9087e-05     	|195            	|  0.005672       	|  0.010583       	|
  on_before_zero_grad                	|  2.8804e-05     	|195            	|  0.0056167      	|  0.01048        	|
  on_before_backward                 	|  2.6982e-05     	|195            	|  0.0052616      	|  0.0098172      	|
  on_train_epoch_end                 	|  0.0014064      	|1              	|  0.0014064      	|  0.0026241      	|
  training_step_end                  	|  4.9198e-06     	|195            	|  0.00095937     	|  0.00179        	|
  on_train_epoch_start               	|  0.00025167     	|1              	|  0.00025167     	|  0.00046957     	|
  on_train_end                       	|  0.00017067     	|1              	|  0.00017067     	|  0.00031844     	|
  on_before_accelerator_backend_setup	|  6.968e-05      	|1              	|  6.968e-05      	|  0.00013001     	|
  setup                              	|  5.0209e-05     	|1              	|  5.0209e-05     	|  9.3682e-05     	|
  prepare_data                       	|  4.4779e-05     	|1              	|  4.4779e-05     	|  8.355e-05      	|
  on_fit_end                         	|  3.892e-05      	|1              	|  3.892e-05      	|  7.2618e-05     	|
  on_epoch_start                     	|  3.332e-05      	|1              	|  3.332e-05      	|  6.2169e-05     	|
  on_pretrain_routine_end            	|  3.009e-05      	|1              	|  3.009e-05      	|  5.6143e-05     	|
  on_epoch_end                       	|  2.741e-05      	|1              	|  2.741e-05      	|  5.1142e-05     	|
  on_configure_sharded_model         	|  2.556e-05      	|1              	|  2.556e-05      	|  4.7691e-05     	|
  on_fit_start                       	|  2.0869e-05     	|1              	|  2.0869e-05     	|  3.8938e-05     	|
  teardown                           	|  1.9379e-05     	|1              	|  1.9379e-05     	|  3.6158e-05     	|
  configure_sharded_model            	|  6.5197e-06     	|1              	|  6.5197e-06     	|  1.2165e-05     	|
  configure_callbacks                	|  5.16e-06       	|1              	|  5.16e-06       	|  9.6277e-06     	|
  on_train_dataloader                	|  4.2003e-06     	|1              	|  4.2003e-06     	|  7.837e-06      	|
  
  As we can see a large time is spent in training_step and here is the output of advanced profiler for this function:
  Profile stats for: training_step rank: 0
           1065072 function calls (862519 primitive calls) in 37.086 seconds
  
     Ordered by: cumulative time
  
     ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        195    0.001    0.000   37.082    0.190 accelerator.py:210(training_step)
        195    0.000    0.000   37.079    0.190 ddp.py:438(training_step)
  31980/195    0.053    0.000   37.079    0.190 module.py:1096(_call_impl)
        195    0.015    0.000   37.078    0.190 distributed.py:852(forward)
        195    0.002    0.000   36.629    0.188 base.py:76(forward)
        195    0.009    0.000   36.624    0.188 moco.py:201(training_step)
   1170/390    0.006    0.000   34.429    0.088 grad_mode.py:25(decorate_context)
        195    0.002    0.000   32.216    0.165 moco.py:83(_update_queue)
        195   32.171    0.165   32.171    0.165 moco.py:109(_get_ptr)
        390    0.000    0.000    3.942    0.010 resnet.py:268(forward)
       ...
        195    0.008    0.000    0.008    0.000 moco.py:124(_assign_in_queue)
        195    0.005    0.000    0.006    0.000 moco.py:115(_assign_ptr)
        195    0.002    0.000    0.002    0.000 moco.py:121(_transpose)
        195    0.001    0.000    0.001    0.000 gather.py:44(concat_all_gather_without_backprop)
        195    0.000    0.000    0.000    0.000 moco.py:106(_get_batch_size)
       ...
  
  The function _update_queue is very long and the function taking the most time is _get_ptr which should be really fast in comparison with forwards or computation of MoCo loss. I watched lightning bolts implementation that uses the same kind of operations so I don't really understand why it is this slow.
  I tested with DDP and SingleDevice strategy that resulted in the same kind of slow down on a SLURM cluster environment.
answer: Fixed it, lightning is now as fast as my previous implementation, the problem was elsewhere but I didn't detect it using the profiler because of the asynchronous computation from GPUs which were not synchronized during profiling.