add fp8 training with torchao

2026-02-03 20:51:26 +00:00
parent 8ebc14b348
commit 6079f78fc3
3 changed files with 180 additions and 180 deletions
@@ -19,7 +19,8 @@ dependencies = [
    "tabulate>=0.9.0",
    "tiktoken>=0.11.0",
    "tokenizers>=0.22.0",
-    "torch>=2.9.0",
+    "torch==2.9.1",
+    "torchao==0.15.0",
    "transformers>=4.57.3",
    "uvicorn>=0.36.0",
    "wandb>=0.21.3",
@@ -59,10 +60,10 @@ explicit = true

 [project.optional-dependencies]
 cpu = [
-    "torch>=2.9.1",
+    "torch==2.9.1",
 ]
 gpu = [
-    "torch>=2.9.1",
+    "torch==2.9.1",
 ]

 [tool.uv]