Merge pull request #192 from srowen/bf16flag

Set bf16 flags correctly for a10/a100
databrickslabs · Jun 7, 2023 · a357706 · a357706
2 parents fd1a733 + 6609c85
commit a357706
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 3 deletions.
diff --git a/config/a100_config.json b/config/a100_config.json
@@ -1,6 +1,9 @@
 {
+    "fp16": {
+      "enabled": false
+    },
     "bf16": {
-      "enabled": "auto"
+      "enabled": true
     },
     "optimizer": {
       "type": "AdamW",

diff --git a/config/a10_config.json b/config/a10_config.json
@@ -1,6 +1,9 @@
 {
+    "fp16": {
+      "enabled": false
+    },
     "bf16": {
-      "enabled": "auto"
+      "enabled": true
     },
     "optimizer": {
       "type": "AdamW",

diff --git a/config/v100_config.json b/config/v100_config.json
@@ -2,6 +2,9 @@
     "fp16": {
       "enabled": true
     },
+    "bf16": {
+      "enabled": false
+    },
     "optimizer": {
       "type": "AdamW",
       "params": {

diff --git a/train_dolly.py b/train_dolly.py
@@ -160,6 +160,11 @@
     num_gpus = int(num_gpus)
     num_gpus_flag = f"--num_gpus={num_gpus}"
 
+if gpu_family == "v100":
+    bf16_flag = "--bf16 false"
+else:
+    bf16_flag = "--bf16 true"
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # COMMAND ----------
@@ -184,7 +189,8 @@
     --eval-steps 50 \
     --warmup-steps 50 \
     --test-size 200 \
-    --lr 5e-6
+    --lr 5e-6 \
+    {bf16_flag}
 
 # COMMAND ----------