PR types

Performance optimization

PR changes

Others

Description

Pcard-76459
Support parallel cross entropy in static semi-auto training.

Introduction of parallel cross entropy
Parallel cross entropy is a performance optimization strategy when the tensor is sharded on the softmax normalize axis. Parallel cross entropy will first perform some local computation and then do communication. Compared to the original pipline that performs communication before computation, parallel cross entropy reduces the communication elements size from (b,s,v) to (b,s), and also reduces the computation elements number on each process. The following figure shows the pipline difference between parallel cross entropy and the original cross entropy.

What this pr does
This pr supports parallel cross entropy in static semi-auto training. When there is a softmax_with_cross_entropy operator in the model, we will first use spmd rule (#58913) to infer the sharding status of its input and output tensors. If the input tensor is sharded on the softmax normalize axis, we will select c_softmax_with_cross_entropy kernel.