Unable to integrate ModelPruning callback in ASR training
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

When adding ModelPruning callback to the Pytorch Lightning trainer in order to train ASR models (in particular, I tried Quartznet15x5) I get the following exception:

TypeError                                 Traceback (most recent call last)
<ipython-input-16-2116e952a23e> in <module>
      1 # Start training!!!
----> 2 trainer.fit(first_asr_model)

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    497 
    498         # dispath `start_training` or `start_testing` or `start_predicting`
--> 499         self.dispatch()
    500 
    501         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py in dispatch(self)
    544 
    545         else:
--> 546             self.accelerator.start_training(self)
    547 
    548     def train_or_test_or_predict(self):

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/accelerators/accelerator.py in start_training(self, trainer)
     71 
     72     def start_training(self, trainer):
---> 73         self.training_type_plugin.start_training(trainer)
     74 
     75     def start_testing(self, trainer):

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in start_training(self, trainer)
    112     def start_training(self, trainer: 'Trainer') -> None:
    113         # double dispatch to initiate the training loop
--> 114         self._results = trainer.run_train()
    115 
    116     def start_testing(self, trainer: 'Trainer') -> None:

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
    635                 with self.profiler.profile("run_training_epoch"):
    636                     # run train epoch
--> 637                     self.train_loop.run_training_epoch()
    638 
    639                 if self.max_steps and self.max_steps <= self.global_step:

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
    575 
    576         if should_check_val:
--> 577             self.trainer.run_evaluation(on_epoch=True)
    578 
    579             # reset stage to train

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py in run_evaluation(self, max_batches, on_epoch)
    750 
    751         # hook
--> 752         self.evaluation_loop.on_evaluation_end()
    753 
    754         # log epoch metrics

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/evaluation_loop.py in on_evaluation_end(self, *args, **kwargs)
     98             self.trainer.call_hook('on_test_end', *args, **kwargs)
     99         else:
--> 100             self.trainer.call_hook('on_validation_end', *args, **kwargs)
    101 
    102     def reload_evaluation_dataloaders(self):

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py in call_hook(self, hook_name, *args, **kwargs)
   1093             if hasattr(self, hook_name):
   1094                 trainer_hook = getattr(self, hook_name)
-> 1095                 trainer_hook(*args, **kwargs)
   1096 
   1097             # next call hook in lightningModule

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/callback_hook.py in on_validation_end(self)
    183         """Called when the validation loop ends."""
    184         for callback in self.callbacks:
--> 185             callback.on_validation_end(self, self.lightning_module)
    186 
    187     def on_test_start(self):

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in on_validation_end(self, trainer, pl_module)
    210         checkpoints can be saved at the end of the val loop
    211         """
--> 212         self.save_checkpoint(trainer, pl_module)
    213 
    214     def on_save_checkpoint(self, trainer, pl_module, checkpoint: Dict[str, Any]) -> Dict[str, Any]:

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in save_checkpoint(self, trainer, pl_module)
    257         # Mode 1: save all checkpoints OR only the top k
    258         if self.save_top_k:
--> 259             self._save_top_k_checkpoints(trainer, pl_module, monitor_candidates)
    260 
    261         # Mode 2: save the last checkpoint

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_top_k_checkpoints(self, trainer, pl_module, metrics)
    561 
    562         if self.check_monitor_top_k(trainer, current):
--> 563             self._update_best_and_save(current, epoch, step, trainer, pl_module, metrics)
    564         elif self.verbose:
    565             rank_zero_info(f"Epoch {epoch:d}, step {step:d}: {self.monitor} was not in top {self.save_top_k}")

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _update_best_and_save(self, current, epoch, step, trainer, pl_module, ckpt_name_metrics)
    603                 f' (best {self.best_model_score:0.5f}), saving model to "{filepath}" as top {k}'
    604             )
--> 605         self._save_model(filepath, trainer, pl_module)
    606 
    607         if del_filepath is not None and filepath != del_filepath:

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_model(self, filepath, trainer, pl_module)
    333         # delegate the saving to the trainer
    334         if self.save_function is not None:
--> 335             self.save_function(filepath, self.save_weights_only)
    336         else:
    337             raise ValueError(".save_function() not set")

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/properties.py in save_checkpoint(self, filepath, weights_only)
    325 
    326     def save_checkpoint(self, filepath, weights_only: bool = False) -> None:
--> 327         self.checkpoint_connector.save_checkpoint(filepath, weights_only)
    328 
    329     @property

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in save_checkpoint(self, filepath, weights_only)
    395         """
    396         # dump states as a checkpoint dictionary object
--> 397         checkpoint = self.dump_checkpoint(weights_only)
    398         if self.trainer.is_global_zero:
    399             # write the checkpoint dictionary on the file

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in dump_checkpoint(self, weights_only)
    282         if not weights_only:
    283             # dump callbacks
--> 284             checkpoint['callbacks'] = self.trainer.on_save_checkpoint(checkpoint)
    285 
    286             optimizer_states = []

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/callback_hook.py in on_save_checkpoint(self, checkpoint)
    220                 state = callback.on_save_checkpoint(self, self.lightning_module)  # noqa: parameter-unfilled
    221             else:
--> 222                 state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
    223             if state:
    224                 callback_states[type(callback)] = state

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/callbacks/pruning.py in on_save_checkpoint(self, trainer, pl_module, checkpoint)
    398             prev_device = pl_module.device
    399             # prune a copy so training can continue with the same buffers
--> 400             copy = deepcopy(pl_module.to("cpu"))
    401             self.make_pruning_permanent(copy)
    402             checkpoint["state_dict"] = copy.state_dict()

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    178                     y = x
    179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
    181 
    182     # If is its own copy, don't memoize.

/opt/conda/lib/python3.6/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    278     if state is not None:
    279         if deep:
--> 280             state = deepcopy(state, memo)
    281         if hasattr(y, '__setstate__'):
    282             y.__setstate__(state)

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    148     copier = _deepcopy_dispatch.get(cls)
    149     if copier:
--> 150         y = copier(x, memo)
    151     else:
    152         try:

/opt/conda/lib/python3.6/copy.py in _deepcopy_dict(x, memo, deepcopy)
    238     memo[id(x)] = y
    239     for key, value in x.items():
--> 240         y[deepcopy(key, memo)] = deepcopy(value, memo)
    241     return y
    242 d[dict] = _deepcopy_dict

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    178                     y = x
    179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
    181 
    182     # If is its own copy, don't memoize.

/opt/conda/lib/python3.6/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    278     if state is not None:
    279         if deep:
--> 280             state = deepcopy(state, memo)
    281         if hasattr(y, '__setstate__'):
    282             y.__setstate__(state)

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    148     copier = _deepcopy_dispatch.get(cls)
    149     if copier:
--> 150         y = copier(x, memo)
    151     else:
    152         try:

/opt/conda/lib/python3.6/copy.py in _deepcopy_dict(x, memo, deepcopy)
    238     memo[id(x)] = y
    239     for key, value in x.items():
--> 240         y[deepcopy(key, memo)] = deepcopy(value, memo)
    241     return y
    242 d[dict] = _deepcopy_dict

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    178                     y = x
    179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
    181 
    182     # If is its own copy, don't memoize.

/opt/conda/lib/python3.6/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    278     if state is not None:
    279         if deep:
--> 280             state = deepcopy(state, memo)
    281         if hasattr(y, '__setstate__'):
    282             y.__setstate__(state)

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    148     copier = _deepcopy_dispatch.get(cls)
    149     if copier:
--> 150         y = copier(x, memo)
    151     else:
    152         try:

/opt/conda/lib/python3.6/copy.py in _deepcopy_dict(x, memo, deepcopy)
    238     memo[id(x)] = y
    239     for key, value in x.items():
--> 240         y[deepcopy(key, memo)] = deepcopy(value, memo)
    241     return y
    242 d[dict] = _deepcopy_dict

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    178                     y = x
    179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
    181 
    182     # If is its own copy, don't memoize.

/opt/conda/lib/python3.6/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    278     if state is not None:
    279         if deep:
--> 280             state = deepcopy(state, memo)
    281         if hasattr(y, '__setstate__'):
    282             y.__setstate__(state)

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    148     copier = _deepcopy_dispatch.get(cls)
    149     if copier:
--> 150         y = copier(x, memo)
    151     else:
    152         try:

/opt/conda/lib/python3.6/copy.py in _deepcopy_dict(x, memo, deepcopy)
    238     memo[id(x)] = y
    239     for key, value in x.items():
--> 240         y[deepcopy(key, memo)] = deepcopy(value, memo)
    241     return y
    242 d[dict] = _deepcopy_dict

/opt/conda/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    167                     reductor = getattr(x, "__reduce_ex__", None)
    168                     if reductor:
--> 169                         rv = reductor(4)
    170                     else:
    171                         reductor = getattr(x, "__reduce__", None)

TypeError: can't pickle BoundFunctionWrapper objects


**Steps/Code to reproduce bug**
The easiest way I could think of to replicate this problem is to take the ASR with NeMo tutorial (No. 1), and replace 

```python
trainer = pl.Trainer(gpus=1, max_epochs=50)
```
with
```python
mp = pl.callbacks.ModelPruning("l1_unstructured", amount=0.5)
trainer = pl.Trainer(gpus=1, callbacks=[mp], max_epochs=50)
```

and run until the first fit. It will work with any simple training process, since I met this bug when using NeMo/examples/asr/speech_to_text.py on my own data.

**Expected behavior**

I expected the training to proceed as usual, and the model to be pruned during training

**Environment overview (please complete the following information)**

 - Environment location: Docker (clean installation of Nvidia's Pytorch 20.11)
 - Method of NeMo install: I installed from source using 
```shell
git clone https://github.com/NVIDIA/NeMo
cd NeMo
./reinstall.sh 
pip install ipywidgets #otherwise the Jupyter Notebook crashes when trying to show progress bar
```

 - If method of install is [Docker], provide `docker pull` & `docker run` commands used
```shell
sudo docker pull nvcr.io/nvidia/pytorch:20.11-py3
sudo docker run --gpus all -it --rm -v ./nemo:/NeMo --shm-size=8g \
-p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit \
stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:21.03-py3
```

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version

**Additional context**

I don't think additional context is needed. My GPU is GTX 1080-ti, but that's probably irrelevant 


@titu1994  any thoughts? 

Just to text, can you add accelerator="ddp" (even for 1 gpu, just a test). 

As far as I can parse the error, it's due to a deep copy of a model which can't be pickled, so it is raising a bound function wrapper error. 

You'll probably have to come up with a toy example (which is unpickleable) and try the same to see if that's indeed the reason. 

That actually did something. I got this other exception instead:

---------------------------------------------------------------------------
ModuleAttributeError                      Traceback (most recent call last)
/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
    636                     # run train epoch
--> 637                     self.train_loop.run_training_epoch()
    638 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
    576         if should_check_val:
--> 577             self.trainer.run_evaluation(on_epoch=True)
    578 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in run_evaluation(self, max_batches, on_epoch)
    751         # hook
--> 752         self.evaluation_loop.on_evaluation_end()
    753 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py in on_evaluation_end(self, *args, **kwargs)
     99         else:
--> 100             self.trainer.call_hook('on_validation_end', *args, **kwargs)
    101 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in call_hook(self, hook_name, *args, **kwargs)
   1094                 trainer_hook = getattr(self, hook_name)
-> 1095                 trainer_hook(*args, **kwargs)
   1096 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/callback_hook.py in on_validation_end(self)
    184         for callback in self.callbacks:
--> 185             callback.on_validation_end(self, self.lightning_module)
    186 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in on_validation_end(self, trainer, pl_module)
    211         """
--> 212         self.save_checkpoint(trainer, pl_module)
    213 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in save_checkpoint(self, trainer, pl_module)
    258         if self.save_top_k:
--> 259             self._save_top_k_checkpoints(trainer, pl_module, monitor_candidates)
    260 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_top_k_checkpoints(self, trainer, pl_module, metrics)
    562         if self.check_monitor_top_k(trainer, current):
--> 563             self._update_best_and_save(current, epoch, step, trainer, pl_module, metrics)
    564         elif self.verbose:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _update_best_and_save(self, current, epoch, step, trainer, pl_module, ckpt_name_metrics)
    604             )
--> 605         self._save_model(filepath, trainer, pl_module)
    606 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_model(self, filepath, trainer, pl_module)
    334         if self.save_function is not None:
--> 335             self.save_function(filepath, self.save_weights_only)
    336         else:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/properties.py in save_checkpoint(self, filepath, weights_only)
    326     def save_checkpoint(self, filepath, weights_only: bool = False) -> None:
--> 327         self.checkpoint_connector.save_checkpoint(filepath, weights_only)
    328 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in save_checkpoint(self, filepath, weights_only)
    396         # dump states as a checkpoint dictionary object
--> 397         checkpoint = self.dump_checkpoint(weights_only)
    398         if self.trainer.is_global_zero:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in dump_checkpoint(self, weights_only)
    283             # dump callbacks
--> 284             checkpoint['callbacks'] = self.trainer.on_save_checkpoint(checkpoint)
    285 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/callback_hook.py in on_save_checkpoint(self, checkpoint)
    221             else:
--> 222                 state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
    223             if state:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/pruning.py in on_save_checkpoint(self, trainer, pl_module, checkpoint)
    399             # prune a copy so training can continue with the same buffers
--> 400             copy = deepcopy(pl_module.to("cpu"))
    401             self.make_pruning_permanent(copy)

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    171                 else:
--> 172                     y = _reconstruct(x, memo, *rv)
    173 

/opt/conda/lib/python3.8/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269         if deep:
--> 270             state = deepcopy(state, memo)
    271         if hasattr(y, '__setstate__'):

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    145     if copier is not None:
--> 146         y = copier(x, memo)
    147     else:

/opt/conda/lib/python3.8/copy.py in _deepcopy_dict(x, memo, deepcopy)
    229     for key, value in x.items():
--> 230         y[deepcopy(key, memo)] = deepcopy(value, memo)
    231     return y

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    171                 else:
--> 172                     y = _reconstruct(x, memo, *rv)
    173 

/opt/conda/lib/python3.8/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269         if deep:
--> 270             state = deepcopy(state, memo)
    271         if hasattr(y, '__setstate__'):

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    145     if copier is not None:
--> 146         y = copier(x, memo)
    147     else:

/opt/conda/lib/python3.8/copy.py in _deepcopy_dict(x, memo, deepcopy)
    229     for key, value in x.items():
--> 230         y[deepcopy(key, memo)] = deepcopy(value, memo)
    231     return y

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    171                 else:
--> 172                     y = _reconstruct(x, memo, *rv)
    173 

/opt/conda/lib/python3.8/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269         if deep:
--> 270             state = deepcopy(state, memo)
    271         if hasattr(y, '__setstate__'):

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    145     if copier is not None:
--> 146         y = copier(x, memo)
    147     else:

/opt/conda/lib/python3.8/copy.py in _deepcopy_dict(x, memo, deepcopy)
    229     for key, value in x.items():
--> 230         y[deepcopy(key, memo)] = deepcopy(value, memo)
    231     return y

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    171                 else:
--> 172                     y = _reconstruct(x, memo, *rv)
    173 

/opt/conda/lib/python3.8/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269         if deep:
--> 270             state = deepcopy(state, memo)
    271         if hasattr(y, '__setstate__'):

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    145     if copier is not None:
--> 146         y = copier(x, memo)
    147     else:

/opt/conda/lib/python3.8/copy.py in _deepcopy_dict(x, memo, deepcopy)
    229     for key, value in x.items():
--> 230         y[deepcopy(key, memo)] = deepcopy(value, memo)
    231     return y

/opt/conda/lib/python3.8/copy.py in deepcopy(x, memo, _nil)
    171                 else:
--> 172                     y = _reconstruct(x, memo, *rv)
    173 

/opt/conda/lib/python3.8/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    271         if hasattr(y, '__setstate__'):
--> 272             y.__setstate__(state)
    273         else:

/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py in __setstate__(self, state)
    549         self.__dict__.setdefault('require_backward_grad_sync', True)
--> 550         self._ddp_init_helper()
    551 

/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py in _ddp_init_helper(self)
    479 
--> 480         self.modules_params = [list(parameters(m)) for m in self._module_copies]
    481         self.modules_buffers = [list(m.buffers()) for m in self._module_copies]

/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py in <listcomp>(.0)
    479 
--> 480         self.modules_params = [list(parameters(m)) for m in self._module_copies]
    481         self.modules_buffers = [list(m.buffers()) for m in self._module_copies]

/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py in parameters(m, recurse)
    441             for m in m.modules() if recurse else [m]:
--> 442                 for p in model_parameters(m):
    443                     yield p

/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py in model_parameters(m)
    437                     else m.parameters(recurse=False)
--> 438                 for p in ps:
    439                     yield p

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in parameters(self, recurse)
   1088         """
-> 1089         for name, param in self.named_parameters(recurse=recurse):
   1090             yield param

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in named_parameters(self, prefix, recurse)
   1114             prefix=prefix, recurse=recurse)
-> 1115         for elem in gen:
   1116             yield elem

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in _named_members(self, get_members_fn, prefix, recurse)
   1059         for module_prefix, module in modules:
-> 1060             members = get_members_fn(module)
   1061             for k, v in members:

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in <lambda>(module)
   1112         gen = self._named_members(
-> 1113             lambda module: module._parameters.items(),
   1114             prefix=prefix, recurse=recurse)

/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py in __getattr__(self, name)
    777                 return modules[name]
--> 778         raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
    779             type(self).__name__, name))

ModuleAttributeError: 'EncDecCTCModel' object has no attribute '_parameters'

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
<ipython-input-17-2116e952a23e> in <module>
      1 # Start training!!!
----> 2 trainer.fit(first_asr_model)

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    497 
    498         # dispath `start_training` or `start_testing` or `start_predicting`
--> 499         self.dispatch()
    500 
    501         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in dispatch(self)
    544 
    545         else:
--> 546             self.accelerator.start_training(self)
    547 
    548     def train_or_test_or_predict(self):

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py in start_training(self, trainer)
     71 
     72     def start_training(self, trainer):
---> 73         self.training_type_plugin.start_training(trainer)
     74 
     75     def start_testing(self, trainer):

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in start_training(self, trainer)
    112     def start_training(self, trainer: 'Trainer') -> None:
    113         # double dispatch to initiate the training loop
--> 114         self._results = trainer.run_train()
    115 
    116     def start_testing(self, trainer: 'Trainer') -> None:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
    668         finally:
    669             # hook
--> 670             self.train_loop.on_train_end()
    671 
    672     def run_evaluation(self, max_batches=None, on_epoch=False):

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in on_train_end(self)
    132         # when a checkpoint was saved at the last step
    133         self.trainer.global_step -= 1
--> 134         self.check_checkpoint_callback(should_update=True, is_last=True)
    135         self.trainer.global_step += 1
    136 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in check_checkpoint_callback(self, should_update, is_last)
    162 
    163             for cb in callbacks:
--> 164                 cb.on_validation_end(self.trainer, model)
    165 
    166     def check_early_stopping_callback(self, should_update):

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in on_validation_end(self, trainer, pl_module)
    210         checkpoints can be saved at the end of the val loop
    211         """
--> 212         self.save_checkpoint(trainer, pl_module)
    213 
    214     def on_save_checkpoint(self, trainer, pl_module, checkpoint: Dict[str, Any]) -> Dict[str, Any]:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in save_checkpoint(self, trainer, pl_module)
    257         # Mode 1: save all checkpoints OR only the top k
    258         if self.save_top_k:
--> 259             self._save_top_k_checkpoints(trainer, pl_module, monitor_candidates)
    260 
    261         # Mode 2: save the last checkpoint

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_top_k_checkpoints(self, trainer, pl_module, metrics)
    560         step = metrics.get("step")
    561 
--> 562         if self.check_monitor_top_k(trainer, current):
    563             self._update_best_and_save(current, epoch, step, trainer, pl_module, metrics)
    564         elif self.verbose:

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py in check_monitor_top_k(self, trainer, current)
    360 
    361         # If using multiple devices, make sure all processes are unanimous on the decision.
--> 362         should_update_best_and_save = trainer.training_type_plugin.reduce_boolean_decision(should_update_best_and_save)
    363 
    364         return should_update_best_and_save

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/parallel.py in reduce_boolean_decision(self, decision)
     90     def reduce_boolean_decision(self, decision: bool) -> bool:
     91         decision = torch.tensor(int(decision), device=self.lightning_module.device)
---> 92         decision = self.reduce(decision, reduce_op=ReduceOp.SUM)
     93         decision = bool(decision == self.world_size)
     94         return decision

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py in reduce(self, output, group, reduce_op)
    308     def reduce(self, output, group: Optional[Any] = None, reduce_op: Optional[Union[ReduceOp, str]] = None):
    309         if isinstance(output, torch.Tensor):
--> 310             output = sync_ddp_if_available(output, group, reduce_op)
    311         return output
    312 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py in sync_ddp_if_available(result, group, reduce_op)
    130     """
    131     if torch.distributed.is_available() and torch.distributed.is_initialized():
--> 132         return sync_ddp(result, group=group, reduce_op=reduce_op)
    133     return result
    134 

/opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py in sync_ddp(result, group, reduce_op)
    163     # sync all processes before reduction
    164     torch.distributed.barrier(group=group)
--> 165     torch.distributed.all_reduce(result, op=op, group=group, async_op=False)
    166 
    167     if divide_by_world_size:

/opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py in all_reduce(tensor, op, group, async_op)
    952     if group == GroupMember.WORLD:
    953         _check_default_pg()
--> 954         work = _default_pg.allreduce([tensor], opts)
    955     else:
    956         work = group.allreduce([tensor], opts)

RuntimeError: Tensors must be CUDA and dense

BTW I changed docker settings to check if the the issue also happens in other python versions. I downloaded the latest pytorch docker, and added the following lines to the requirements so the pytorch versions will match:
```bash
torch==1.7.1
torchtext==0.8.1
```
I'll try that again with the 20.11 docker later.

What version of cuda and cudnn is currently installed on your system? It's trying to all reduce the integer flag which acts as a boolean, but it fails cause ints are not dense and the step value is not on cuda. 

I think you can't use ddp + this pruning. We haven't tried and don't support it since it's not tested. I would suggest creating a minimal unpickleable PT lightning module and then trying with and without ddp if it works. If even that doesn't work, it should be brought up to Pytorch lightning team. 

I understand what's the problem, and I made a hacky solution.
Apparently, to make the pruning permanent in the saved model, during checkpoints the ModelPruning callback deep- copies the model, makes the pruning permanent in the copy and then replaces the checkpoint's state dict by the copy's dict:
```python
    def on_save_checkpoint(self, trainer, pl_module: LightningModule, checkpoint: Dict[str, Any]):
        if self._make_pruning_permanent:
            rank_zero_debug("`ModelPruning.on_save_checkpoint`. Pruning is made permanent for this checkpoint.")
            prev_device = pl_module.device
            # prune a copy so training can continue with the same buffers            
            copy = deepcopy(pl_module.to("cpu"))
            
            self.make_pruning_permanent(copy)
            checkpoint["state_dict"] = copy.state_dict()
            
            pl_module.to(prev_device)
```

The EncDecCTCModel module has several attributes that cannot be deep- copied for various reasons (I tried deep copying all of them by themselves and they caused exceptions, and removing all of them was required in my solution):
- `trainer` (has an attribute dev_debugger which is a weakref, which has a weakref back to itself)
- `_trainer` (same)
- `_train_dl` (Lightning Data Loader, I'm not sure why it's uncopiable)
- `_validation_dl` (same)
- `training_step` (a function wrapper, and so uncopiable)

My hacky solution is replacing all the offending attributes by `None`, letting the callback do its thing and then place them back in the module:
```python
    def on_save_checkpoint(self, trainer, pl_module: LightningModule, checkpoint: Dict[str, Any]):
        if self._make_pruning_permanent:
            rank_zero_debug("`ModelPruning.on_save_checkpoint`. Pruning is made permanent for this checkpoint.")
            prev_device = pl_module.device
            # prune a copy so training can continue with the same buffers
            
            trainer = pl_module.trainer
            _trainer = pl_module._trainer
            _train_dl = pl_module._train_dl
            _validation_dl = pl_module._validation_dl
            training_step = pl_module.training_step
            
            pl_module.trainer = None
            pl_module._trainer = None
            pl_module._train_dl = None
            pl_module._validation_dl = None
            pl_module.training_step = None
            
            copy = deepcopy(pl_module.to("cpu"))
            
            self.make_pruning_permanent(copy)
            checkpoint["state_dict"] = copy.state_dict()
            
            pl_module.trainer = trainer
            pl_module._trainer = _trainer
            pl_module._train_dl = _train_dl
            pl_module._validation_dl = _validation_dl
            pl_module.training_step = training_step
            
            pl_module.to(prev_device)
```
A less hacky solution would be to replace the `__deepcolpy__` method of EncDecCTCModel with something that does not try to copy uncopyable attributes. An even less hacky solution would be to have a ModelPruning implementation which doesn't deepcopies the whole model just to operate on its weights, but I'm not sure how that would work.
Thanks for your help!