While it could be that PyTorch’s kernels are less optimized, I do think that apex uses three kernels if you want all gradients for input and parameters (cuComputePartGradGammaBeta, cuComputeGradGammaBeta, cuComputeGradInput) so the difference is not quite as stark.
That said, if you have benchmarks that show a PyTorch being slower than Apex, I’m quite certain people will look into it more (PyToch did adopt the LayerNorm forward at some point).
Best regards
Thomas