use batch_size of 4 for A10 GPUs

Author: tnixon

train_dolly.py

@@ -34,14 +34,14 @@
 
 # COMMAND ----------
 
-# MAGIC !wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcusparse-dev-11-3_11.5.0.58-1_amd64.deb -O /tmp/libcusparse-dev-11-3_11.5.0.58-1_amd64.deb && \
-# MAGIC   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcublas-dev-11-3_11.5.1.109-1_amd64.deb -O /tmp/libcublas-dev-11-3_11.5.1.109-1_amd64.deb && \
-# MAGIC   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcusolver-dev-11-3_11.1.2.109-1_amd64.deb -O /tmp/libcusolver-dev-11-3_11.1.2.109-1_amd64.deb && \
-# MAGIC   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcurand-dev-11-3_10.2.4.109-1_amd64.deb -O /tmp/libcurand-dev-11-3_10.2.4.109-1_amd64.deb && \
-# MAGIC   dpkg -i /tmp/libcusparse-dev-11-3_11.5.0.58-1_amd64.deb && \
-# MAGIC   dpkg -i /tmp/libcublas-dev-11-3_11.5.1.109-1_amd64.deb && \
-# MAGIC   dpkg -i /tmp/libcusolver-dev-11-3_11.1.2.109-1_amd64.deb && \
-# MAGIC   dpkg -i /tmp/libcurand-dev-11-3_10.2.4.109-1_amd64.deb
+!wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcusparse-dev-11-3_11.5.0.58-1_amd64.deb -O /tmp/libcusparse-dev-11-3_11.5.0.58-1_amd64.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcublas-dev-11-3_11.5.1.109-1_amd64.deb -O /tmp/libcublas-dev-11-3_11.5.1.109-1_amd64.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcusolver-dev-11-3_11.1.2.109-1_amd64.deb -O /tmp/libcusolver-dev-11-3_11.1.2.109-1_amd64.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/libcurand-dev-11-3_10.2.4.109-1_amd64.deb -O /tmp/libcurand-dev-11-3_10.2.4.109-1_amd64.deb && \
+  dpkg -i /tmp/libcusparse-dev-11-3_11.5.0.58-1_amd64.deb && \
+  dpkg -i /tmp/libcublas-dev-11-3_11.5.1.109-1_amd64.deb && \
+  dpkg -i /tmp/libcusolver-dev-11-3_11.1.2.109-1_amd64.deb && \
+  dpkg -i /tmp/libcurand-dev-11-3_10.2.4.109-1_amd64.deb
 
 # COMMAND ----------
 
@@ -148,7 +148,9 @@
 
 # configure the batch_size
 batch_size = 3
-if gpu_family == "a100":
+if gpu_family == "a10":
+    batch_size = 4
+elif gpu_family == "a100":
     batch_size = 6
 
 # configure num_gpus, if specified
@@ -167,22 +169,22 @@
 
 # COMMAND ----------
 
-# MAGIC !deepspeed {num_gpus_flag} \
-# MAGIC     --module training.trainer \
-# MAGIC     --input-model {input_model} \
-# MAGIC     --deepspeed {deepspeed_config} \
-# MAGIC     --epochs 2 \
-# MAGIC     --local-output-dir {local_output_dir} \
-# MAGIC     --dbfs-output-dir {dbfs_output_dir} \
-# MAGIC     --per-device-train-batch-size {batch_size} \
-# MAGIC     --per-device-eval-batch-size {batch_size} \
-# MAGIC     --logging-steps 10 \
-# MAGIC     --save-steps 200 \
-# MAGIC     --save-total-limit 20 \
-# MAGIC     --eval-steps 50 \
-# MAGIC     --warmup-steps 50 \
-# MAGIC     --test-size 200 \
-# MAGIC     --lr 5e-6
+!deepspeed {num_gpus_flag} \
+    --module training.trainer \
+    --input-model {input_model} \
+    --deepspeed {deepspeed_config} \
+    --epochs 2 \
+    --local-output-dir {local_output_dir} \
+    --dbfs-output-dir {dbfs_output_dir} \
+    --per-device-train-batch-size {batch_size} \
+    --per-device-eval-batch-size {batch_size} \
+    --logging-steps 10 \
+    --save-steps 200 \
+    --save-total-limit 20 \
+    --eval-steps 50 \
+    --warmup-steps 50 \
+    --test-size 200 \
+    --lr 5e-6
 
 # COMMAND ----------