⚗️ Tiny MoE for test

scripts/generate_tiny_models.py

@@ -64,6 +64,8 @@
     Qwen3Config,
     Qwen3ForCausalLM,
     Qwen3ForSequenceClassification,
+    Qwen3MoeConfig,
+    Qwen3MoeForCausalLM,
     SiglipVisionConfig,
     T5Config,
     T5ForConditionalGeneration,
@@ -143,6 +145,24 @@ def push_to_hub(model, tokenizer, prefix=None, suffix=None):
     model = model_class(config)
     push_to_hub(model, tokenizer, "tiny", suffix)
 
+# MoE models
+for model_id, config_class, model_class, suffix in [
+    ("Qwen/Qwen3-30B-A3B", Qwen3MoeConfig, Qwen3MoeForCausalLM, None),
+]:
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    config = config_class(
+        vocab_size=tokenizer.vocab_size + len(tokenizer.added_tokens_encoder.keys()),
+        hidden_size=8,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        num_hidden_layers=2,
+        intermediate_size=32,
+        num_experts=4,
+        num_experts_per_tok=2,
+    )
+    model = model_class(config)
+    push_to_hub(model, tokenizer, "tiny", suffix)
+
 
 # Two slightly bigger models, required for vLLM testing
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-32B-Instruct")

tests/test_sft_trainer.py

@@ -1129,16 +1129,20 @@ def test_torch_dtype(self):
 
 # This new tester aims to replace the first one at some point
 class SFTTrainerTester2(unittest.TestCase):
-    def test_train(self):
+    @parameterized.expand(
+        [
+            ("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",),
+            ("trl-internal-testing/tiny-Qwen3MoeForCausalLM",),
+        ]
+    )
+    def test_train(self, model_id):
         # Get the dataset
         dataset = load_dataset("trl-internal-testing/zen", "standard_language_modeling", split="train")
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             # Initialize the trainer
             training_args = SFTConfig(output_dir=tmp_dir, report_to="none")
-            trainer = SFTTrainer(
-                model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=training_args, train_dataset=dataset
-            )
+            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()}