Feat: Add support for gemma3_text and add e2e for gemma2 (#2406)

2025-03-23 07:33:21 +07:00
parent 86bac48d14
commit 9f00465a5c
12 changed files with 348 additions and 6 deletions
--- a/tests/e2e/test_schedulers.py
+++ b/tests/e2e/test_schedulers.py
@@ -54,7 +54,7 @@ class TestCustomSchedulers(unittest.TestCase):
                "gradient_accumulation_steps": 1,
                "output_dir": temp_dir,
                "learning_rate": 0.00001,
-                "optimizer": "adamw_hf",
+                "optimizer": "adamw_torch_fused",
                "max_steps": 20,
                "lr_scheduler": "rex",
                "warmup_steps": 5,