Alit/nano v2 (#14464)

* add nano v2 support

* add convertor fix

* minor change to test script

* remove use_mamba_mem_eff_path

* Apply isort and black reformatting

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

---------

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>
Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos02.eos.clusters.nvidia.com>
Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com>
Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com>

Author: 3 people

nemo/collections/llm/__init__.py

@@ -153,6 +153,7 @@
     NemotronHConfig47B,
     NemotronHConfig56B,
     NemotronModel,
+    NemotronNano9Bv2,
     NVIDIAMambaConfig8B,
     NVIDIAMambaHybridConfig8B,
     Phi3Config,
@@ -279,6 +280,7 @@
     "NemotronHConfig8B",
     "NemotronHConfig47B",
     "NemotronHConfig56B",
+    "NemotronNano9Bv2",
     "MambaModel",
     "LlamaConfig",
     "Llama2Config7B",

nemo/collections/llm/gpt/model/__init__.py

@@ -166,6 +166,7 @@
     NemotronHConfig8B,
     NemotronHConfig47B,
     NemotronHConfig56B,
+    NemotronNano9Bv2,
     NVIDIAMambaConfig8B,
     NVIDIAMambaHybridConfig8B,
     SSMConfig,
@@ -304,6 +305,7 @@
     "NemotronHConfig8B",
     "NemotronHConfig47B",
     "NemotronHConfig56B",
+    "NemotronNano9Bv2",
     "MambaModel",
     "DeepSeekModel",
     "DeepSeekV2Config",

nemo/collections/llm/gpt/model/ssm.py

@@ -575,6 +575,8 @@ def make_vocab_size_divisible_by(vocab_size):
             nemotron_h_config = NemotronHConfig47B()
         elif "56B" in source._name_or_path:
             nemotron_h_config = NemotronHConfig56B()
+        elif "Nano-9B-v2" in source._name_or_path:
+            nemotron_h_config = NemotronNano9Bv2()
         else:
             raise ValueError(f"Unsupported model size: {source._name_or_path}")
 
@@ -702,6 +704,9 @@ def config(self):
         elif type(source) == NemotronHConfig56B:
             model_path = local_model_path if local_model_path else "nvidia/Nemotron-H-56B-Base-8K"
             hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        elif type(source) == NemotronNano9Bv2:
+            model_path = local_model_path if local_model_path else "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+            hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
         else:
             raise ValueError(f"Unsupported model size: {source}")
 
@@ -1039,6 +1044,21 @@ class NemotronHConfig56B(NemotronHConfigBase):
     num_attention_heads: int = 64
 
 
+@dataclass
+class NemotronNano9Bv2(NemotronHConfigBase):
+    """NemotronNano9Bv2"""
+
+    hybrid_override_pattern: str = "M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M-"
+    num_layers: int = 56
+    hidden_size: int = 4480
+    mamba_num_heads: int = 128
+    kv_channels: int = 128
+    mamba_state_dim: int = 128
+    ffn_hidden_size: int = 15680
+    num_attention_heads: int = 40
+    mamba_head_dim: int = 80
+
+
 __all__ = [
     "SSMConfig",
     "BaseMambaConfig130M",
@@ -1053,4 +1073,5 @@ class NemotronHConfig56B(NemotronHConfigBase):
     "NemotronHConfig8B",
     "NemotronHConfig47B",
     "NemotronHConfig56B",
+    "NemotronNano9Bv2",
 ]

nemo/collections/llm/recipes/__init__.py

@@ -81,6 +81,7 @@
     nemotron4_15b_16k,
     nemotron4_15b_64k,
     nemotron4_340b,
+    nemotron_nano_9b_v2,
     nemotronh_4b,
     nemotronh_8b,
     nemotronh_47b,
@@ -163,6 +164,7 @@
     "nemotronh_8b",
     "nemotronh_47b",
     "nemotronh_56b",
+    "nemotron_nano_9b_v2",
     "mistral_7b",
     "mistral_nemo_12b",
     "hyena_base",