resolve gradient checkpointing issue.

Former-commit-id: 6df9135d063bb6102f0cbcdf0d702076f5febbae
2024-04-16 12:05:27 +08:00
parent d4d471450f
commit 6dd6b3e396
4 changed files with 8 additions and 14 deletions
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@ extra_require = {
    "metrics": ["nltk", "jieba", "rouge-chinese"],
    "unsloth": ["torch==2.2.0", "unsloth[cu121-ampere-torch220]"],
    "galore": ["galore-torch"],
+    "badam": ["torch>=2.1.0"],
    "vllm": ["vllm>=0.3.3"],
    "bitsandbytes": ["bitsandbytes>=0.39.0"],
    "gptq": ["optimum>=1.16.0", "auto-gptq>=0.5.0"],