Is it possible to call dist.all_reduce manually in train_step?

[sandylaker]

In my code, I would like to synchronize a tensor across all the gpus in train_step, which is a temporary variable. Is it allowed to call torch.distributed.all_reduce in this case? Or there is a specific function in pytorch_lightning that does the job?

[SeanNaren]

Hey @sandylaker! You can use torch.distributed.all_reduce. There is also within the LightningModule this function, however it may be better to expose this within lightning to make it easier to access:

x = self.trainer.accelerator.training_type_plugin.reduce(x)