It would be good if you could give more context. This is from fairseq?
The mixed precision parts have been merged (though PyTorch 1.8 is so last year…  ), but apex also has some unrelated optimizations, including kernels for operating on multiple tensors (I think).
The question always is whether these optimizations are worth having in PyTorch (which has >2000 functions last someone counted, which is quite a burden at times) or whether they’re only when you want to race implementations against each other.
You could install apex and compare performance before and after to find out. If it is a large difference, there might be an argument for getting these into PyTorch. If it is just to squeeze out the large percent of performance, maybe not as much.
Best regards
Thomas