add new args from mcore

Signed-off-by: Chen Cui <[email protected]>
NVIDIA · Oct 23, 2024 · 721f910 · 721f910
1 parent ffca434
commit 721f910
Showing 1 changed file with 2 additions and 0 deletions.
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
@@ -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)