feat: use DeepSpeed ZeRO-3 for proper model sharding across GPUs

training/configs/ds_zero3.json

@@ -0,0 +1,22 @@
+{
+  "bf16": {
+    "enabled": true
+  },
+
+  "zero_optimization": {
+    "stage": 3,
+
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+
+    "reduce_bucket_size": 50000000,
+    "stage3_prefetch_bucket_size": 50000000,
+    "stage3_param_persistence_threshold": 100000
+  },
+
+  "gradient_clipping": 1.0,
+
+  "train_micro_batch_size_per_gpu": 1,
+
+  "steps_per_print": 100
+}