Using BaseFinetuning callback with DeepSpeed

[aakaashjois]

I am trying to use the MilestoneFinetuning as given in this example. I am using DeepSpeed optimization to distribute the job on 8 GPUs. When the .fit() starts, an error is raised by BaseFinetuning.__apply_mapping_to_param_groups():

  File "/fast/cqa-fine-tune/learning_box_2/cli/lb_lightning_cli.py", line 221, in fit
    self.trainer.fit(self.model, datamodule=self.data_module)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 552, in fit
    self._run(model)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 917, in _run
    self._dispatch()
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 985, in _dispatch
    self.accelerator.start_training(self)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 92, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 161, in start_training
    self._results = trainer.run_stage()
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 995, in run_stage
    return self._run_train()
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1044, in _run_train
    self.fit_loop.run()
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 111, in run
    self.advance(*args, **kwargs)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/loops/fit_loop.py", line 200, in advance
    epoch_output = self.epoch_loop.run(train_dataloader)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 106, in run
    self.on_run_start(*args, **kwargs)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 106, in on_run_start
    self.trainer.call_hook("on_train_epoch_start")
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1219, in call_hook
    trainer_hook(*args, **kwargs)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/trainer/callback_hook.py", line 92, in on_train_epoch_start
    callback.on_train_epoch_start(self, self.lightning_module)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/callbacks/finetuning.py", line 298, in on_train_epoch_start
    self._store(pl_module, opt_idx, num_param_groups, current_param_groups)
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/callbacks/finetuning.py", line 284, in _store
    current_param_groups, mapping
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/callbacks/finetuning.py", line 270, in __apply_mapping_to_param_groups
    group_state["params"] = [mapping[p] for p in g["params"]]
  File "/home/aj337a/.pyenv/versions/miniconda3-3.7-4.9.2/envs/py371/lib/python3.7/site-packages/pytorch_lightning/callbacks/finetuning.py", line 270, in <listcomp>
    group_state["params"] = [mapping[p] for p in g["params"]]
KeyError: tensor([-0.0161, -0.0145,  0.0069,  ..., -0.0033,  0.0071,  0.0056],
       device='cuda:0', requires_grad=True)

My understanding is, DeepSpeed changes the way Optimizer is distributed across the GPUs, and this might be causing a problem here. Anyone experienced a similar issue, and know of a solution? I had to revert back to using DDP for now to get past this error.