add check if pos embed

Signed-off-by: jiemingz <jiemingz@nvidia.com>
JimmyZhang12 · Apr 9, 2024 · 72f7fd2 · 72f7fd2
1 parent 35e400f
commit 72f7fd2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_base_model.py b/nemo/collections/nlp/models/language_modeling/megatron_base_model.py
@@ -775,7 +775,8 @@ def get_config_arg(key: str, default_value: Optional[Any] = None) -> Any:
  if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
  if self.mcore_gpt:
  fp32_params.append(modules[0].shared_embedding_or_output_weight())
- fp32_params.append(modules[0].embedding.position_embeddings.weight)
+ if modules[0].embedding.add_position_embedding:
+ fp32_params.append(modules[0].embedding.position_embeddings.weight)
  else:
  fp32_params.append(modules[0].word_embeddings_weight())
  fp32_params.append(modules[0].position_embeddings_weight())