Can't find the matched layer names when restore_from nemo file

[wangzhongruo]

Hi,

I use the following code to restore the parameters from a nemo file:

conf = OmegaConf.load(cfg_path)

with open_dict(conf):
    conf.trainer.accelerator='cpu'
    # conf.trainer.devices=1

nemo_path = 'nemo_gpt1.3B_fp16.nemo'
connector = NLPSaveRestoreConnector()
trainer = hydra.utils.instantiate(conf.trainer)

print('trainer', trainer)

connector.model_extracted_dir= nemo_path

gpt_cfg = MegatronGPTModel.restore_from(
    restore_path=nemo_path,
    trainer=trainer,
    return_config=True,
    save_restore_connector=connector,
)
print('cfg modified before', gpt_cfg)

with open_dict(gpt_cfg):
    gpt_cfg['global_batch_size'] = 96
    gpt_cfg['megatron_amp_O2'] = True
    gpt_cfg['use_cpu_initialization'] = True

print('cfg modified after', gpt_cfg)

save_restore_connector = NLPSaveRestoreConnector()
save_restore_connector.model_extracted_dir = nemo_path

model = MegatronGPTModel.restore_from(
    restore_path=nemo_path,
    trainer=trainer,
    override_config_path=gpt_cfg,
    save_restore_connector=save_restore_connector,
    map_location="cpu",
)

But it provides the error like the following:

RuntimeError: Error(s) in loading state_dict for MegatronGPTModel:
	Unexpected key(s) in state_dict: "model.module.language_model.embedding.word_embeddings.weight", "model.module.language_model.embedding.position_embeddings.weight", "model.module.language_model.encoder.layers.0.input_layernorm.weight", "model.module.language_model.encoder.layers.0.input_layernorm.bias", "model.module.language_model.encoder.layers.0.self_attention.query_key_value.weight", "model.module.language_model.encoder.layers.0.self_attention.query_key_value.bias", "model.module.language_model.encoder.layers.0.self_attention.dense.weight", "model.module.language_model.encoder.layers.0.self_attention.dense.bias", "model.module.language_model.encoder.layers.0.post_attention_layernorm.weight",

My guess is the layer name is not matched up so it can't be copied to the initialized MegatronGPTModel.

May I ask is there anyway we can get the layer names from an initialized MegatronGPTModel?

I tried methods like model.named_parameters() or 'model.modules()', they all doesn't work.