Updated gemma_2b_it CI (huggingface#1561)

Co-authored-by: regisss <15324346+regisss@users.noreply.github.com>
HabanaAI · Jan 20, 2025 · 142fbb5 · 142fbb5
1 parent 0f2767b
commit 142fbb5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/tests/baselines/gemma_2b_it.json b/tests/baselines/gemma_2b_it.json
@@ -25,6 +25,18 @@
                         "--dataset_config_name wikitext-2-raw-v1",
                         "--use_hpu_graphs_for_inference"
                     ]
+                },
+                "deepspeed": {
+                    "learning_rate": 8e-4,
+                    "train_batch_size": 4,
+                    "perplexity": 924.062,
+                    "train_runtime": 75.518,
+                    "train_samples_per_second": 81.097,
+                    "extra_arguments": [
+                        "--dataset_config_name wikitext-2-raw-v1",
+                        "--use_hpu_graphs_for_inference",
+                        "--deepspeed tests/configs/deepspeed_zero_2.json"
+                    ]
                 }
             }
         }

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -288,6 +288,8 @@ def to_test(
             return False
         elif eager_mode and model_name not in models_measured_on_eager_mode:
             return False
+        elif "gemma" in model_name and not IS_GAUDI2:
+            return False
         elif model_name not in models_with_specific_rules and not deepspeed:
             return True
         elif model_name == "gpt2-xl" and deepspeed: