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

2026-07-01 22:08:37 -04:00
parent 9d8c2b6cf0
commit a655040db2
3 changed files with 31 additions and 10 deletions
--- a/training/configs/ds_zero3.json
+++ b/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
+}