revert

Signed-off-by: ericharper <[email protected]>
NVIDIA · Aug 10, 2023 · 3213cac · 3213cac
1 parent d4a7576
commit 3213cac
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -643,7 +643,7 @@ def _get_total_params_across_model_parallel_groups_gpt_bert(self, model):
                 word_embeddings_weight = (
                     model[-1].module.shared_embedding_or_output_weight()
                     if getattr(self, 'mcore_gpt', False)
-                    else model[-1].model_embeddings_weight()
+                    else model[-1].word_embeddings_weight()
                 )
                 # substract the embedding weights on the last virtual stage
                 num_word_embedding_parameters = sum([p.nelement() for p in word_embeddings_weight])
@@ -658,7 +658,7 @@ def _get_total_params_across_model_parallel_groups_gpt_bert(self, model):
                 word_embeddings_weight = (
                     model.module.shared_embedding_or_output_weight()
                     if getattr(self, 'mcore_gpt', False)
-                    else model.model_embeddings_weight()
+                    else model.word_embeddings_weight()
                 )
                 # substract the embedding weights on the last stage
                 num_word_embedding_parameters = sum([p.nelement() for p in word_embeddings_weight])

diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -724,7 +724,7 @@ def allreduce_first_last_embeddings(self):
             )
             if share_embeddings:
                 word_embeddings_weight = (
-                    module.shared_embedding_or_output_weight() if self.mcore_gpt else module.model_embeddings_weight()
+                    module.shared_embedding_or_output_weight() if self.mcore_gpt else module.word_embeddings_weight()
                 )
                 # (@adithyare) adapter training now extends MegatronGPTModel so we have to add this check here to ensure we do not perform all_reduce when grad is None.
                 # grad can be None when performing PeFT training.