Add Nemotron 3 to tests via tiny model

scripts/generate_tiny_models.py

@@ -59,6 +59,8 @@
     LlavaNextForConditionalGeneration,
     MistralConfig,
     MistralForCausalLM,
+    NemotronHConfig,
+    NemotronHForCausalLM,
     OPTConfig,
     OPTForCausalLM,
     PaliGemmaForConditionalGeneration,
@@ -227,6 +229,39 @@ def init_weights_tiny_model(model):
     init_weights_tiny_model(model)
     push_to_hub(model, tokenizer, generation_config, "tiny", suffix)
 
+# Hybrid Mamba-Attention models
+tokenizer = AutoTokenizer.from_pretrained("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16")
+generation_config = GenerationConfig.from_pretrained("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16")
+config = NemotronHConfig(
+    vocab_size=len(tokenizer.vocab),
+    hidden_size=16,
+    num_attention_heads=4,
+    num_key_value_heads=2,
+    intermediate_size=32,
+    layers_block_type=["mamba", "attention"],  # 2 layers: one Mamba + one Attention
+    mamba_num_heads=8,
+    mamba_head_dim=4,
+    mamba_n_groups=1,
+    ssm_state_size=16,
+    mamba_d_conv=4,
+    mamba_expand=2,
+    n_routed_experts=4,
+    num_experts_per_tok=2,
+    moe_intermediate_size=32,
+    moe_shared_expert_intermediate_size=32,
+    use_mamba_kernels=False,  # CPU-friendly for testing
+)
+model = NemotronHForCausalLM(config).to(dtype=torch.bfloat16)
+init_weights_tiny_model(model)
+# NemotronH keeps mixer.D and mixer.A_log in float32 in the reference model; mirror that here.
+for layer in model.model.layers:
+    if hasattr(layer, "mixer"):
+        if hasattr(layer.mixer, "D"):
+            layer.mixer.D.data = layer.mixer.D.data.float()
+        if hasattr(layer.mixer, "A_log"):
+            layer.mixer.A_log.data = layer.mixer.A_log.data.float()
+push_to_hub(model, tokenizer, generation_config, "tiny")
+
 # Two slightly bigger models, required for vLLM testing
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-32B-Instruct")
 generation_config = GenerationConfig.from_pretrained("Qwen/Qwen2.5-32B-Instruct")

tests/test_dpo_trainer.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from contextlib import nullcontext
+from unittest.mock import patch
+
 import pytest
 import torch
 import transformers
@@ -171,19 +174,36 @@ class TestDPOTrainer(TrlTestCase):
             "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
             "trl-internal-testing/tiny-Qwen3MoeForCausalLM",
             "trl-internal-testing/tiny-GptOssForCausalLM",
+            pytest.param(
+                "trl-internal-testing/tiny-NemotronHForCausalLM",
+                marks=pytest.mark.skipif(
+                    Version(transformers.__version__) < Version("5.3.0"),
+                    reason="NemotronH models were introduced in transformers-5.3.0",
+                ),
+            ),
         ],
     )
     def test_train(self, model_id):
         # Get the dataset
         dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")
 
+        # NemotronH (hybrid Mamba-Attention) does not support gradient checkpointing.
+        # Workaround: hide kernels package so transformers doesn't unconditionally load mamba CUDA kernels.
+        # See: https://github.com/huggingface/transformers/pull/44853
+        kwargs = {}
+        ctx = patch("transformers.integrations.hub_kernels.lazy_load_kernel", return_value=None) if "NemotronH" in model_id else nullcontext()
+        if "NemotronH" in model_id:
+            kwargs["gradient_checkpointing"] = False
+
         # Initialize the trainer
-        training_args = DPOConfig(
-            output_dir=self.tmp_dir,
-            learning_rate=0.1,  # use higher lr because gradients are tiny and default lr can stall updates
-            report_to="none",
-        )
-        trainer = DPOTrainer(model=model_id, args=training_args, train_dataset=dataset)
+        with ctx:
+            training_args = DPOConfig(
+                output_dir=self.tmp_dir,
+                learning_rate=0.1,  # use higher lr because gradients are tiny and default lr can stall updates
+                report_to="none",
+                **kwargs,
+            )
+            trainer = DPOTrainer(model=model_id, args=training_args, train_dataset=dataset)
 
         # Save the initial parameters to compare them later
         previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}

tests/test_sft_trainer.py

@@ -15,6 +15,8 @@
 import gc
 import json
 import pathlib
+from contextlib import nullcontext
+from unittest.mock import patch
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -284,15 +286,31 @@ def test_init_with_training_arguments(self):
             "trl-internal-testing/tiny-GptOssForCausalLM",
             "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
             "trl-internal-testing/tiny-Qwen3MoeForCausalLM",
+            pytest.param(
+                "trl-internal-testing/tiny-NemotronHForCausalLM",
+                marks=pytest.mark.skipif(
+                    Version(transformers.__version__) < Version("5.3.0"),
+                    reason="NemotronH models were introduced in transformers-5.3.0",
+                ),
+            ),
         ],
     )
     def test_train(self, model_id):
         # Get the dataset
         dataset = load_dataset("trl-internal-testing/zen", "standard_language_modeling", split="train")
 
+        # NemotronH (hybrid Mamba-Attention) does not support gradient checkpointing.
+        # Workaround: hide kernels package so transformers doesn't unconditionally load mamba CUDA kernels.
+        # See: https://github.com/huggingface/transformers/pull/44853
+        kwargs = {}
+        ctx = patch("transformers.integrations.hub_kernels.lazy_load_kernel", return_value=None) if "NemotronH" in model_id else nullcontext()
+        if "NemotronH" in model_id:
+            kwargs["gradient_checkpointing"] = False
+
         # Initialize the trainer
-        training_args = SFTConfig(output_dir=self.tmp_dir, report_to="none")
-        trainer = SFTTrainer(model=model_id, args=training_args, train_dataset=dataset)
+        with ctx:
+            training_args = SFTConfig(output_dir=self.tmp_dir, report_to="none", **kwargs)
+            trainer = SFTTrainer(model=model_id, args=training_args, train_dataset=dataset)
 
         # Save the initial parameters to compare them later
         previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}