Update ZeRO config for fp16 compatibility

haotian-liu · haotian-liu · commit 6b3079ea5da5 · 2023-08-04T12:33:45.000-05:00
diff --git a/scripts/zero2.json b/scripts/zero2.json
@@ -1,6 +1,14 @@
 {
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
     "bf16": {
-        "enabled": true
+        "enabled": "auto"
     },
     "train_micro_batch_size_per_gpu": "auto",
     "train_batch_size": "auto",
diff --git a/scripts/zero3.json b/scripts/zero3.json
@@ -1,6 +1,14 @@
 {
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
     "bf16": {
-        "enabled": true
+        "enabled": "auto"
     },
     "train_micro_batch_size_per_gpu": "auto",
     "train_batch_size": "auto",
diff --git a/scripts/zero3_offload.json b/scripts/zero3_offload.json
@@ -8,7 +8,7 @@
     "min_loss_scale": 1
   },
   "bf16": {
-    "enabled": true
+    "enabled": "auto"
   },
   "optimizer": {
     "type": "AdamW",
