load gatortron checkpoint into biomegatron model, hit Key 'optim' is not in struct
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**
not able to complete the training process using pretrained model [biomegatron345m_biovocab_50k_cased]

**Steps/Code to reproduce bug**
NVIDIA NeMo Introduction Notebook
Setup
using the NGC nvcr.io/nvidia/pytorch:22.04-py3 PyTorch container.
Install [NeMo](https://github.com/NVIDIA/NeMo) using the following steps:
apt-get update && apt-get install -y libsndfile1 ffmpeg
pip install Cython
pip install nemo_toolkit['all']

![image](https://user-images.githubusercontent.com/53548997/208076338-0ac6587d-c79f-4e6d-a1f2-021dd58b4e75.png)

# ERROR at this step
# load gatortron checkpoint into biomegatron model
# (the architecture is the same, just the weights are different)

config.model.language_model.lm_checkpoint = "gatortron_og/MegatronBERT.nemo"
config.model.language_model.pretrained_model_name = 'biomegatron345m_biovocab_50k_cased'
config.model.tokenizer.tokenizer_name=None

model = nemo_nlp.models.TextClassificationModel(cfg=config.model, trainer=trainer)

############ ERROR ############
![image](https://user-images.githubusercontent.com/53548997/208077032-a2a2fd16-0722-42d0-91cb-71e2bfa7b494.png)

[NeMo W 2022-12-16 07:59:06 lm_utils:80] biomegatron345m_biovocab_50k_cased is not in get_pretrained_lm_models_list(include_external=False), will be using AutoModel from HuggingFace.

---------------------------------------------------------------------------
ConfigAttributeError                      Traceback (most recent call last)
Input In [18], in <cell line: 8>()
      5 config.model.language_model.pretrained_model_name = 'biomegatron345m_biovocab_50k_cased'
      6 config.model.tokenizer.tokenizer_name=None
----> 8 model = nemo_nlp.models.TextClassificationModel(cfg=config.model, trainer=trainer)

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/text_classification/text_classification_model.py:43, in TextClassificationModel.__init__(self, cfg, trainer)
     40 self.dataset_cfg = cfg.dataset
     41 self.class_weights = None
---> 43 super().__init__(cfg=cfg, trainer=trainer)
     45 self.classifier = SequenceClassifier(
     46     hidden_size=self.hidden_size,
     47     num_classes=cfg.dataset.num_classes,
   (...)
     53     idx_conditioned_on=0,
     54 )
     56 self.create_loss_module()

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/nlp_model.py:81, in NLPModel.__init__(self, cfg, trainer, no_lm_init)
     79 elif pretrain_model_name in all_pretrained_megatron_bert_models:
     80     copy_cfg = copy.deepcopy(cfg)
---> 81     bert_model = get_lm_model(
     82         config_file=config_file,
     83         config_dict=config_dict,
     84         vocab_file=vocab_file,
     85         trainer=trainer,
     86         cfg=copy_cfg,
     87     )
     88     # set the tokenizer if it is not initialized explicitly
     89     if (
     90         (hasattr(self, 'tokenizer') and self.tokenizer is None) or not hasattr(self, 'tokenizer')
     91     ) and hasattr(bert_model, 'tokenizer'):

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/lm_utils.py:112, in get_lm_model(config_dict, config_file, vocab_file, trainer, cfg)
    109         return x
    111 if cfg.language_model.get("lm_checkpoint"):
--> 112     model = MegatronBertModel.restore_from(restore_path=cfg.language_model.lm_checkpoint, trainer=trainer)
    113 else:
    114     model = MegatronBertModel.from_pretrained(cfg.language_model.get('pretrained_model_name'), trainer=trainer)

File /opt/conda/lib/python3.8/site-packages/nemo/core/classes/modelPT.py:316, in ModelPT.restore_from(cls, restore_path, override_config_path, map_location, strict, return_config, save_restore_connector, trainer)
    313 app_state.model_restore_path = restore_path
    315 cls.update_save_restore_connector(save_restore_connector)
--> 316 instance = cls._save_restore_connector.restore_from(
    317     cls, restore_path, override_config_path, map_location, strict, return_config, trainer
    318 )
    319 if isinstance(instance, ModelPT):
    320     instance._save_restore_connector = save_restore_connector

File /opt/conda/lib/python3.8/site-packages/nemo/core/connectors/save_restore_connector.py:235, in SaveRestoreConnector.restore_from(self, calling_cls, restore_path, override_config_path, map_location, strict, return_config, trainer)
    210 """
    211 Restores model instance (weights and configuration) into .nemo file
    212 
   (...)
    231     An instance of type cls or its underlying config (if return_config is set).
    232 """
    233 # Get path where the command is executed - the artifacts will be "retrieved" there
    234 # (original .nemo behavior)
--> 235 loaded_params = self.load_config_and_state_dict(
    236     calling_cls, restore_path, override_config_path, map_location, strict, return_config, trainer,
    237 )
    238 if not isinstance(loaded_params, tuple):
    239     return loaded_params

File /opt/conda/lib/python3.8/site-packages/nemo/core/connectors/save_restore_connector.py:158, in SaveRestoreConnector.load_config_and_state_dict(self, calling_cls, restore_path, override_config_path, map_location, strict, return_config, trainer)
    156 # get the class
    157 calling_cls._set_model_restore_state(is_being_restored=True, folder=tmpdir)
--> 158 instance = calling_cls.from_config_dict(config=conf, trainer=trainer)
    159 instance = instance.to(map_location)
    160 # add load_state_dict override

File /opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py:505, in Serialization.from_config_dict(cls, config, trainer)
    503             if prev_error:
    504                 logging.error(prev_error)
--> 505             raise e
    507 if not hasattr(instance, '_cfg'):
    508     instance._cfg = config

File /opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py:497, in Serialization.from_config_dict(cls, config, trainer)
    495 accepts_trainer = Serialization._inspect_signature_for_trainer(cls)
    496 if accepts_trainer:
--> 497     instance = cls(cfg=config, trainer=trainer)
    498 else:
    499     instance = cls(cfg=config)

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron_bert_model.py:56, in MegatronBertModel.__init__(self, cfg, trainer)
     52 if not HAVE_APEX:
     53     raise ImportError(
     54         "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
     55     )
---> 56 super().__init__(cfg, trainer=trainer, no_lm_init=False)
     57 self.cfg = cfg
     59 # used in NVIDIA NGC PyTorch containers
     60 # buffer used during train_step for logging average loss over gradient accumulation steps

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron_base_model.py:78, in MegatronBaseModel.__init__(self, cfg, trainer, no_lm_init)
     74 self.tokenizer = None
     76 super().__init__(cfg, trainer=trainer, no_lm_init=no_lm_init)
---> 78 self.with_distributed_adam = cfg.optim.get('name') == 'distributed_fused_adam'
     80 # used in NVIDIA NGC PyTorch containers
     81 self._enable_nvidia_optimizations()

File /opt/conda/lib/python3.8/site-packages/omegaconf/dictconfig.py:357, in DictConfig.__getattr__(self, key)
    353     self._format_and_raise(
    354         key=key, value=None, cause=e, type_override=ConfigAttributeError
    355     )
    356 except Exception as e:
--> 357     self._format_and_raise(key=key, value=None, cause=e)

File /opt/conda/lib/python3.8/site-packages/omegaconf/base.py:190, in Node._format_and_raise(self, key, value, cause, type_override)
    187 def _format_and_raise(
    188     self, key: Any, value: Any, cause: Exception, type_override: Any = None
    189 ) -> None:
--> 190     format_and_raise(
    191         node=self,
    192         key=key,
    193         value=value,
    194         msg=str(cause),
    195         cause=cause,
    196         type_override=type_override,
    197     )
    198     assert False

File /opt/conda/lib/python3.8/site-packages/omegaconf/_utils.py:738, in format_and_raise(node, key, value, msg, cause, type_override)
    736         ex = type_override(str(cause))
    737         ex.__dict__ = copy.deepcopy(cause.__dict__)
--> 738     _raise(ex, cause)
    740 object_type: Optional[Type[Any]]
    741 object_type_str: Optional[str] = None

File /opt/conda/lib/python3.8/site-packages/omegaconf/_utils.py:716, in _raise(ex, cause)
    714 else:
    715     ex.__cause__ = None
--> 716 raise ex.with_traceback(sys.exc_info()[2])

File /opt/conda/lib/python3.8/site-packages/omegaconf/dictconfig.py:351, in DictConfig.__getattr__(self, key)
    348     raise AttributeError()
    350 try:
--> 351     return self._get_impl(key=key, default_value=_DEFAULT_MARKER_)
    352 except ConfigKeyError as e:
    353     self._format_and_raise(
    354         key=key, value=None, cause=e, type_override=ConfigAttributeError
    355     )

File /opt/conda/lib/python3.8/site-packages/omegaconf/dictconfig.py:438, in DictConfig._get_impl(self, key, default_value)
    436 def _get_impl(self, key: DictKeyType, default_value: Any) -> Any:
    437     try:
--> 438         node = self._get_node(key=key, throw_on_missing_key=True)
    439     except (ConfigAttributeError, ConfigKeyError):
    440         if default_value is not _DEFAULT_MARKER_:

File /opt/conda/lib/python3.8/site-packages/omegaconf/dictconfig.py:465, in DictConfig._get_node(self, key, validate_access, throw_on_missing_value, throw_on_missing_key)
    462         return None
    464 if validate_access:
--> 465     self._validate_get(key)
    467 value: Optional[Node] = self.__dict__["_content"].get(key)
    468 if value is None:

File /opt/conda/lib/python3.8/site-packages/omegaconf/dictconfig.py:166, in DictConfig._validate_get(self, key, value)
    164 else:
    165     msg = f"Key '{key}' is not in struct"
--> 166 self._format_and_raise(
    167     key=key, value=value, cause=ConfigAttributeError(msg)
    168 )

File /opt/conda/lib/python3.8/site-packages/omegaconf/base.py:190, in Node._format_and_raise(self, key, value, cause, type_override)
    187 def _format_and_raise(
    188     self, key: Any, value: Any, cause: Exception, type_override: Any = None
    189 ) -> None:
--> 190     format_and_raise(
    191         node=self,
    192         key=key,
    193         value=value,
    194         msg=str(cause),
    195         cause=cause,
    196         type_override=type_override,
    197     )
    198     assert False

File /opt/conda/lib/python3.8/site-packages/omegaconf/_utils.py:818, in format_and_raise(node, key, value, msg, cause, type_override)
    815     ex.ref_type = ref_type
    816     ex.ref_type_str = ref_type_str
--> 818 _raise(ex, cause)

File /opt/conda/lib/python3.8/site-packages/omegaconf/_utils.py:716, in _raise(ex, cause)
    714 else:
    715     ex.__cause__ = None
--> 716 raise ex.with_traceback(sys.exc_info()[2])

ConfigAttributeError: Key 'optim' is not in struct
    full_key: optim
    object_type=dict


**Expected behavior**
model training is successful

**Environment overview (please complete the following information)**
docker run --gpus all -it --rm --shm-size=8g -p 8888:8888 --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/pytorch:22.04-py3

**Environment details**
using the NGC nvcr.io/nvidia/pytorch:22.04-py3 PyTorch container.

**Additional context**

Add any other context about the problem here.
Model: RTX 3090 windows 11