25.05.19: add ds train code

2025-05-19 11:51:33 +08:00
parent c2d9b96acc
commit be180dfb6a
3 changed files with 160 additions and 0 deletions
@@ -0,0 +1,40 @@
+{
+    "fp16": {
+      "enabled": true,
+      "loss_scale": 0,
+      "loss_scale_window": 1000,
+      "initial_scale_power": 16,
+      "hysteresis": 2,
+      "min_loss_scale": 1
+    },
+    "optimizer": {
+      "type": "AdamW",
+      "params": {
+        "lr": 5e-5,
+        "weight_decay": 0.01
+      }
+    },
+    "scheduler": {
+      "type": "CosineAnnealing",
+      "params": {
+        "warmup_steps": 500
+      }
+    },
+    "zero_optimization": {
+      "stage": 2,
+      "offload_optimizer": {
+        "device": "cpu",
+        "pin_memory": true
+      },
+      "allgather_partitions": true,
+      "allgather_bucket_size": 2e8,
+      "overlap_comm": true,
+      "reduce_scatter": true,
+      "reduce_bucket_size": 2e8,
+      "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 2,
+    "train_batch_size": "auto",
+    "gradient_clipping": 1.0,
+    "steps_per_print": 10
+  }