Large Performance Regression with FusedAdam

I'm running some reinforcement learning experiments and noticed a large performance regression when swapping out PyTorch modules for APEX ones. I've narrowed it down to FusedAdam being used to replace AdamW, as can be seen in the plots below.

Optimiser	Train	Test
optim.AdamW
optimizers.FusedAdam

Relevant code snippets below are a switch at the beginning of my code, setting the optimisers for all my component models, and the training loop:

 if importlib.util.find_spec('apex') and torch.cuda.is_available():  # Use FusedAdam if NVIDIA Apex available           
   from apex import optimizers                                                                                          
   Adam = optimizers.FusedAdam  # Implements AdamW weight decay by default                                              
   Adam.zero_grad = lambda *args, **kwargs: None  # Patch the APEX optimiser to match the standard PyTorch API (set grad to None is done on init by default in APEX)
 else:                                                                                                                  
   Adam = optim.AdamW

encoder_optimiser = Adam(encoder.parameters(), lr=training_cfg.learning_rate, weight_decay=training_cfg.weight_decay)
critic_optimiser = Adam(critic.parameters(), lr=training_cfg.learning_rate, weight_decay=training_cfg.weight_decay)

 encoder_optimiser.zero_grad(set_to_none=True)                                                                   
 critic_optimiser.zero_grad(set_to_none=True)                                                                    
 value_loss.backward()                                                                                                 
 encoder_optimiser.step()                                                                                        
 critic_optimiser.step()

I would expect FusedAdam to act as a drop-in faster version of AdamW, especially with lr and weight_decay set. For reference the experiments above use a learning rate of 0.0003 and a weight decay of 0.

Environment information below (APEX was built on commit 1d77111):

PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.2.0-19ubuntu1) 11.2.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-47-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 515.65.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] botorch==0.6.4
[pip3] gpytorch==1.8.0
[pip3] numpy==1.21.5
[pip3] torch==1.12.1
[pip3] torchaudio==0.12.1
[pip3] torchvision==0.13.1
[conda] blas                      1.0                         mkl  
[conda] botorch                   0.6.4                         0    pytorch
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] gpytorch                  1.8.0              pyhd8ed1ab_0    conda-forge
[conda] libblas                   3.9.0            12_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            12_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            12_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            12_linux64_mkl    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.21.5           py39h6c91a56_3  
[conda] numpy-base                1.21.5           py39ha15fc14_3  
[conda] pytorch                   1.12.1          py3.9_cuda11.3_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.12.1               py39_cu113    pytorch
[conda] torchvision               0.13.1               py39_cu113    pytorch

Start with optimizers.FusedAdam and reproduce the comparison against optim.AdamW using the optimizer setup and training loop shown in the report. Check whether the regression is reproducible under the listed PyTorch, CUDA, Apex, and GPU versions; done means the cause is identified and the reported performance difference is addressed or clearly characterized.