I did a high level mapping. The ComputePartGradGammaBeta and ComputeGradGammaBeta are together equivalent to the PyTorch GammaBetaBackwardKernel basically calculating the gradients for gamma and beta if those params are set to trainable. Despite PyTorch doing it through one kernel, the perf of this gradient computation is actually much worse than Apex since the GammaBetaBackward kernel does not effectively utilize the GPU. Basically low occupancy and memory throughput utilization. I can post some benchmark results but we have seen it in many cases. Other than that what Apex does through cuComputeGradInput maps to a set of three kernels in PyTorch(ComputeInternalGradients,ComputeGradientFusedParamsCUDAKernel and LayerNormBackwardKernel). Its not clear why this part is not fused though the GPU utilization is fairly good for these kernels.