BERTLMModel not serializable?
Hello, I tried to run a BERT pre-training in a distributed fashion with [Ray](https://github.com/ray-project), more specifically using [ray_lightning](https://github.com/ray-project/ray_lightning) as a plugin. Apparently, the BERTLMModel cannot be serialized (stack trace below)
```
Traceback (most recent call last):
  File "models/bert-nvidia/bert_pretraining.py", line 34, in main
    trainer.fit(bert_model)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 499, in fit
    self.dispatch()
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 546, in dispatch
    self.accelerator.start_training(self)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 73, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray_lightning/ray_ddp.py", line 213, in start_training
    results = self.execution_loop(trainer, tune_enabled=False)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray_lightning/ray_ddp.py", line 180, in execution_loop
    model_ref = ray.put(model)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/_private/client_mode_hook.py", line 62, in wrapper
    return func(*args, **kwargs)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/worker.py", line 1527, in put
    object_ref = worker.put_object(value)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/worker.py", line 280, in put_object
    serialized_value = self.get_serialization_context().serialize(value)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/serialization.py", line 326, in serialize
    return self._serialize_to_msgpack(value)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/serialization.py", line 306, in _serialize_to_msgpack
    self._serialize_to_pickle5(metadata, python_objects)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/serialization.py", line 266, in _serialize_to_pickle5
    raise e
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/serialization.py", line 262, in _serialize_to_pickle5
    inband = pickle.dumps(
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/cloudpickle/cloudpickle_fast.py", line 73, in dumps
    cp.dump(obj)
  File "/home/mallonus/tunbert-pip/lib/python3.8/site-packages/ray/cloudpickle/cloudpickle_fast.py", line 580, in dump
    return Pickler.dump(self, obj)
NotImplementedError: object proxy must define __reduce_ex__()
Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
```
I was wondering if there was a quick work-around about this and if it's possible for me to fix it and still use the [BERT pre-training script](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/bert_pretraining.py). Any help will be much appreciated, thanks in advance!
