deepspeed config for bf16 + flash

stanford-crfm · Dec 29, 2022 · a0ba902 · a0ba902
1 parent 1659a9e
commit a0ba902
Showing 1 changed file with 38 additions and 0 deletions.
diff --git a/conf/deepspeed/z2-small-bf16-conf.json b/conf/deepspeed/z2-small-bf16-conf.json
@@ -0,0 +1,38 @@
+{
+  "optimizer": {
+    "type": "AdamW",
+    "params": {
+      "lr": 0.0006,
+      "betas": "auto",
+      "eps": 1e-8,
+      "weight_decay": 0.1
+    }
+  },
+
+  "scheduler": {
+    "type": "WarmupDecayLR",
+    "params": {
+      "total_num_steps": "auto",
+      "warmup_max_lr": 0.0006,
+      "warmup_num_steps": 10000
+    }
+  },
+
+  "zero_optimization": {
+    "stage": 2,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 2e8,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 2e8,
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "cpu_offload": false
+  },
+
+  "bf16": {
+   "enabled": true
+  },
+
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto"
+}