def broadcast_state(self, rank, src_rank): data = None if rank == src_rank: save_stream = io.BytesIO() self.save(save_stream) # Note: save_stream.getbuffer() will return a memoryview, which # cannot be convert to a tensor, need convert it to np array first data = numpy.asarray(save_stream.getbuffer()) data = dist.broadcast_binary(data, src_rank) load_stream = io.BytesIO(data) self.load(load_stream)
def broadcast_run(rank, world_size, input): data = edist.broadcast_binary( np.asarray(input[rank]) if input[rank] else None, 1) return compute_checksum(data)