aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

test broadcast and all reduce

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-24 18:31:59 -03:00
parent 55db13836b
commit 94b3582a46
5 changed files with 17 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

6
norch/csrc/distributed.cpp
View file

@ -46,17 +46,13 @@ void broadcast_tensor(Tensor* tensor) {
cudaStreamDestroy(stream);
}
void allreduce_mean_tensor(Tensor* tensor) {
void allreduce_sum_tensor(Tensor* tensor) {
cudaStream_t stream;
cudaStreamCreate(&stream);
// Perform NCCL AllReduce operation to calculate the sum of all tensors across all processes
NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
// Divide the sum by the world size to get the mean
float scale = 1.0f / world_size;
scalar_mul_tensor_cuda(tensor, scale, tensor->data);
cudaStreamSynchronize(stream);
cudaStreamDestroy(stream);
}

2
norch/csrc/distributed.h
View file

@ -21,6 +21,8 @@
extern "C" {
void init_process_group(int rank, int world_size);
void broadcast_tensor(Tensor* tensor);
void allreduce_sum_tensor(Tensor* tensor);
}
#endif /* DISTRIBUTED_H */

BIN
norch/csrc/distributed.o
View file

Binary file not shown.

2
norch/distributed/distributed.py
View file

@ -16,7 +16,7 @@ def broadcast_tensor(tensor):
Tensor._C.broadcast_tensor(tensor)
def allreduce_mean_tensor(tensor):
def allreduce_sum_tensor(tensor):
Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.allreduce_mean_tensor.restype = None

21
train.py
View file

@ -10,17 +10,22 @@ def main():
dist.init_process_group(rank, world_size)
tensor = norch.Tensor([1,2,3])
tensor = rank * tensor
print(f"BEFORE on rank {rank}: ", tensor, '\n\n')
tensor = norch.Tensor([1,1,1])
tensor = (rank + 1) * tensor
print(f"BEFORE on rank {rank}: {tensor} \n\n")
dist.allreduce_sum_tensor(tensor)
print(f"AFTER ALLREDUCE on rank {rank}: {tensor} \n\n")
print("###############\n\n\n")
tensor = tensor * 10
print(f"BEFORE BROADCAST on rank {rank}: {tensor} \n\n")
dist.broadcast_tensor(tensor)
print(f"AFTER BROADCAST on rank {rank}: ", tensor, '\n\n')
dist.allreduce_mean_tensor(tensor)
print(f"AFTER ALLREDUCE MEAN on rank {rank}: ", tensor, '\n\n')
print(f"AFTER BROADCAST on rank {rank}: {tensor} \n\n")
if __name__ == "__main__":
main()