[BUG] AttributeError: 'FusedAdam' object has no attribute 'set_grad_none'

Describe the bug
When I use FusedAdam from deepspeed as an optimizer, without any other DeepSpeed features, it seems to crash due to a missing attribute. This is clear from the error trace. This is odd to me, because the init clearly has a set_grad_none attribute.

https://github.com/microsoft/DeepSpeed/blob/b4e5826a60b275600673a77478ecb749519caaf0/deepspeed/ops/adam/fused_adam.py#L70

ds_report output

DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [YES] ...... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [YES] ...... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/bram/.local/share/virtualenvs/tolkien-transformers-aNvOnYab/lib/python3.9/site-packages/torch']
torch version .................... 1.9.1+cu102
torch cuda version ............... 10.2
nvcc version ..................... 10.2
deepspeed install path ........... ['/home/bram/.local/share/virtualenvs/tolkien-transformers-aNvOnYab/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.5.4, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.9, cuda 10.2

Error trace
Traceback (most recent call last):
File "python3.9/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
fn(i, *args)
File "python3.9/site-packages/pytorch_lightning/plugins/training_type/ddp_spawn.py", line 201, in new_process
results = trainer.run_stage()
File "python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1000, in run_stage
return self._run_train()
File "python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1049, in _run_train
self.fit_loop.run()
File "python3.9/site-packages/pytorch_lightning/loops/base.py", line 111, in run
self.advance(*args, **kwargs)
File "python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 200, in advance
epoch_output = self.epoch_loop.run(train_dataloader)
File "python3.9/site-packages/pytorch_lightning/loops/base.py", line 111, in run
self.advance(*args, **kwargs)
File "python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 130, in advance
batch_output = self.batch_loop.run(batch, self.iteration_count, self._dataloader_idx)
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 100, in run
super().run(batch, batch_idx, dataloader_idx)
File "python3.9/site-packages/pytorch_lightning/loops/base.py", line 111, in run
self.advance(*args, **kwargs)
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 147, in advance
result = self._run_optimization(batch_idx, split_batch, opt_idx, optimizer)
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 201, in _run_optimization
self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 395, in _optimizer_step
model_ref.optimizer_step(
File "python3.9/site-packages/pytorch_lightning/core/lightning.py", line 1616, in optimizer_step
optimizer.step(closure=optimizer_closure)
File "python3.9/site-packages/pytorch_lightning/core/optimizer.py", line 206, in step
self.__optimizer_step(closure=closure, profiler_name=profiler_name, **kwargs)
File "python3.9/site-packages/pytorch_lightning/core/optimizer.py", line 128, in __optimizer_step
trainer.accelerator.optimizer_step(self._optimizer, self._optimizer_idx, lambda_closure=closure, **kwargs)
File "python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py", line 292, in optimizer_step
make_optimizer_step = self.precision_plugin.pre_optimizer_step(
File "python3.9/site-packages/pytorch_lightning/plugins/precision/native_amp.py", line 57, in pre_optimizer_step
result = lambda_closure() # native amp does not support closures
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 235, in _training_step_and_backward_closure
result = self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, hiddens)
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 543, in training_step_and_backward
self._optimizer_zero_grad(batch_idx, optimizer, opt_idx)
File "python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 426, in _optimizer_zero_grad
self.trainer.accelerator.optimizer_zero_grad(self.trainer.current_epoch, batch_idx, optimizer, opt_idx)
File "python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py", line 308, in optimizer_zero_grad
model_ref.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx)
File "python3.9/site-packages/pytorch_lightning/core/lightning.py", line 1639, in optimizer_zero_grad
optimizer.zero_grad()
File "python3.9/site-packages/deepspeed/ops/adam/fused_adam.py", line 78, in zero_grad
if self.set_grad_none:
AttributeError: 'FusedAdam' object has no attribute 'set_grad_none'

System info (please complete the following information):

• OS: Ubuntu 20.04
• GPU count and types: 4x Tesla V100
• Python version: 3.9

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else? Through PyTorch Lightning