Set bf16 flags corretly for a10/a100

Author: srowen

config/a100_config.json

@@ -1,6 +1,9 @@
 {
+    "fp16": {
+      "enabled": false
+    },
     "bf16": {
-      "enabled": "auto"
+      "enabled": true
     },
     "optimizer": {
       "type": "AdamW",

config/a10_config.json

@@ -1,6 +1,9 @@
 {
+    "fp16": {
+      "enabled": false
+    },
     "bf16": {
-      "enabled": "auto"
+      "enabled": true
     },
     "optimizer": {
       "type": "AdamW",

config/v100_config.json

@@ -2,6 +2,9 @@
     "fp16": {
       "enabled": true
     },
+    "bf16": {
+      "enabled": false
+    },
     "optimizer": {
       "type": "AdamW",
       "params": {

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 ----------