A machine learning compiler for GPUs, CPUs, and ML accelerators
2.55k
stars
394
forks
source link
[XLA:SPMD] Optimize the partitioning for element-wise operations when all operands share the same sharding. #14873
Open
copybara-service[bot] opened 1 month ago
[XLA:SPMD] Optimize the partitioning for element-wise operations when all operands share the same sharding.
Let us take
C with S3 = add(A with S1, B with S2)
as an example, whereA, B, C
are tensors,S1, S2, S3
are their shardings.Before this change, we always have
With this cl, if S1 and S2 are the same, we will have
The new partitioning method can reduce the number of reshards.