schedulefree optimizers (#30079)

* schedulefree optimizers * fix train instead of eval for optimizer * fixes and update docs * chore: lint * add tests and drop overly-verbose _32bit suffix * chore: lint * fix for docs * fix code review issues * use duck-typing to avoid per-optimizer patches * fixup style * fixup style * warn if incorrect accelerate version with schedule free Co-authored-by: Aman Gupta Karmani <aman@tmm1.net> --------- Co-authored-by: Aman Karmani <aman@tmm1.net>
2024-09-09 03:51:39 -04:00
parent 60226fdc1d
commit 62aecd85ff
9 changed files with 124 additions and 0 deletions
--- a/setup.py
+++ b/setup.py
@@ -163,6 +163,7 @@ _deps = [
    "sacremoses",
    "safetensors>=0.4.1",
    "sagemaker>=2.31.0",
+    "schedulefree>=1.2.6",
    "scikit-learn",
    "scipy<1.13.0",  # SciPy >= 1.13.0 is not supported with the current jax pin (`jax>=0.4.1,<=0.4.13`)
    "sentencepiece>=0.1.91,!=0.1.92",