Simulate four distributed collectives locally over a rank-ordered list of one-dimensional PyTorch tensors. A rank is one participating worker. Let R be the number of ranks, L the input length at each rank, xr rank r's input, and yr its output; s indexes source ranks.
For all_gather, every rank receives the rank-ordered concatenation:
yr=concat(x0,…,xR−1).For all_reduce, every rank receives the elementwise sum:
yr=s=0∑R−1xs.For reduce_scatter, rank r receives chunk r of the sum, using R equal contiguous chunks. For all_to_all, rank r receives chunk r from each source, concatenated in source-rank order. Return a Python list of R PyTorch tensors in rank order, preserving input dtype and device. Each output has length R times L for all_gather, L divided by R for reduce_scatter, and L otherwise.
Input: rank_tensors = [[1,2],[3,4]], collective = "all_gather"
Output: [[1,2,3,4],[1,2,3,4]]
Explanation: Rank-order concatenation produces one gathered tensor, and every rank receives a separate copy.
Input: rank_tensors = [[1,2,3,4],[10,20,30,40]], collective = "reduce_scatter"
Output: [[11,22],[33,44]]
Input: rank_tensors = [[1,2,3,4],[5,6,7,8]], collective = "all_to_all"
Output: [[1,2,5,6],[3,4,7,8]]
torch.cat joins tensors along an existing axis.
torch.chunk divides a tensor into ordered chunks.
Tensor.clone creates independent tensor storage.
Accepts: array
Accepts: string
Simulate four distributed collectives locally over a rank-ordered list of one-dimensional PyTorch tensors. A rank is one participating worker. Let R be the number of ranks, L the input length at each rank, xr rank r's input, and yr its output; s indexes source ranks.
For all_gather, every rank receives the rank-ordered concatenation:
yr=concat(x0,…,xR−1).For all_reduce, every rank receives the elementwise sum:
yr=s=0∑R−1xs.For reduce_scatter, rank r receives chunk r of the sum, using R equal contiguous chunks. For all_to_all, rank r receives chunk r from each source, concatenated in source-rank order. Return a Python list of R PyTorch tensors in rank order, preserving input dtype and device. Each output has length R times L for all_gather, L divided by R for reduce_scatter, and L otherwise.
Input: rank_tensors = [[1,2],[3,4]], collective = "all_gather"
Output: [[1,2,3,4],[1,2,3,4]]
Explanation: Rank-order concatenation produces one gathered tensor, and every rank receives a separate copy.
Input: rank_tensors = [[1,2,3,4],[10,20,30,40]], collective = "reduce_scatter"
Output: [[11,22],[33,44]]
Input: rank_tensors = [[1,2,3,4],[5,6,7,8]], collective = "all_to_all"
Output: [[1,2,5,6],[3,4,7,8]]
torch.cat joins tensors along an existing axis.
torch.chunk divides a tensor into ordered chunks.
Tensor.clone creates independent tensor storage.
Accepts: array
Accepts: string
