.

ds_config_zero3.json

@@ -9,20 +9,19 @@
     "reduce_bucket_size": 200000000,
     "stage3_prefetch_bucket_size": 100000000,
     "stage3_param_persistence_threshold": 1000000,
+
     "offload_optimizer": { "device": "cpu", "pin_memory": true },
-    "offload_param": { "device": "none" }
+    "offload_param":     { "device": "none" }
+  },
+
+  "optimizer": {
+    "type": "AdamW",
+    "params": { "lr": 2e-5, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.1 }
   },
 
   "bf16": { "enabled": true },
   "fp16": { "enabled": false },
 
-  "optimizer": {
-    "type": "CPUAdam",
-    "params": { "lr": 2e-5, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.1 }
-  },
-
   "gradient_clipping": 1.0,
-  "zero_force_ds_cpu_optimizer": true,
-
   "wall_clock_breakdown": false
 }