fix: accept fractional warmup rates

finetune_moss.py

@@ -287,7 +287,7 @@ def train(args):
     parser.add_argument('--eval_bsz_per_gpu', default=4, type=int)
     parser.add_argument('--weight_decay', default=0.1, type=float)
     parser.add_argument('--learning_rate', default=9e-6, type=float)
-    parser.add_argument('--warmup_rates', default=0.05, type=int)
+    parser.add_argument('--warmup_rates', default=0.05, type=float)
     parser.add_argument('--n_epochs', default=2, type=int)
 
     # Other Args

tests/test_finetune_cli.py

@@ -0,0 +1,29 @@
+import subprocess
+import sys
+import unittest
+from pathlib import Path
+
+
+REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
+
+
+class FinetuneCliTest(unittest.TestCase):
+    def test_accepts_fractional_warmup_rate(self):
+        result = subprocess.run(
+            [
+                sys.executable,
+                str(REPOSITORY_ROOT / "finetune_moss.py"),
+                "--warmup_rates",
+                "0.1",
+                "--help",
+            ],
+            cwd=REPOSITORY_ROOT,
+            capture_output=True,
+            text=True,
+        )
+
+        self.assertEqual(result.returncode, 0, result.stderr)
+
+
+if __name__ == "__main__":
+    unittest.main()