def test_dist_broadcast_coalesced(self):
# Set up process group.
store = c10d.TCPStore('localhost', self.port, self.is_master)
options = c10d.ProcessGroupGloo.Options()
options.devices = [c10d.ProcessGroupGloo.create_tcp_device(interface="lo")]
process_group = c10d.ProcessGroupGloo(store, self.rank, self.world_size, options)
device = torch.device('cuda')
target = torch.arange(10, dtype=torch.float64, device=device).chunk(5)
if self.is_master:
# All processes should have these tensors in the end.
tensors = target
else:
# Non-master processes start with empty tensors and should be
# filled with the tensors from the master.
tensors = torch.zeros(10, device=device).chunk(5)
c10d._dist_broadcast_coalesced(
tensors,
buffer_size=10,
process_group=process_group)
if not self.is_master:
self.assertEqual(tensors, target)
def _dist_broadcast_coalesced(self, tensors, buffer_size):
c10d._dist_broadcast_coalesced(tensors, buffer_size,
self.process_group)