From 165e5f99f382586895ceaa9da08064fb31ba8837 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Wed, 23 Oct 2024 18:06:57 -0400 Subject: [PATCH] add new args from mcore Signed-off-by: Chen Cui --- .../models/language_modeling/megatron/gemma2/gemma2_modules.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nemo/collections/nlp/models/language_modeling/megatron/gemma2/gemma2_modules.py b/nemo/collections/nlp/models/language_modeling/megatron/gemma2/gemma2_modules.py index 11123d03b4e8..10524cc2f61a 100644 --- a/nemo/collections/nlp/models/language_modeling/megatron/gemma2/gemma2_modules.py +++ b/nemo/collections/nlp/models/language_modeling/megatron/gemma2/gemma2_modules.py @@ -66,6 +66,8 @@ def __init__( attn_mask_type: AttnMaskType, attention_type: str, attention_dropout: float = None, + softmax_scale: float = None, + cp_comm_type: str = None, ): super().__init__(config=config)