[feat] fp8 training (#8960)

Co-authored-by: Benjamin Feuer <penfever@gmail.com> Co-authored-by: Yaowei Zheng <hiyouga@buaa.edu.cn>
2025-09-30 23:32:53 -07:00
parent e2b1594d31
commit 1c44b60e3e
8 changed files with 322 additions and 3 deletions
--- a/setup.py
+++ b/setup.py
@@ -70,6 +70,9 @@ extra_require = {
    ],
    "openmind": ["openmind"],
    "swanlab": ["swanlab"],
+    "fp8": ["torchao>=0.8.0", "accelerate>=1.10.0"],
+    "fp8-te": ["transformer_engine[pytorch]>=2.0.0", "accelerate>=1.10.0"],
+    "fp8-all": ["torchao>=0.8.0", "transformer_engine[pytorch]>=2.0.0", "accelerate>=1.10.0"],
    "dev": ["pre-commit", "ruff", "pytest", "build"],
 }