diff --git a/ds_config_zero3.json b/ds_config_zero3.json index 48b3f7c..86d9d3b 100644 --- a/ds_config_zero3.json +++ b/ds_config_zero3.json @@ -6,12 +6,15 @@ "stage": 3, "overlap_comm": true, "contiguous_gradients": true, - "reduce_bucket_size": 200000000, - "stage3_prefetch_bucket_size": 100000000, + + "reduce_bucket_size": 150000000, + "stage3_prefetch_bucket_size": 75000000, "stage3_param_persistence_threshold": 1000000, - "offload_optimizer": { "device": "cpu", "pin_memory": true }, - "offload_param": { "device": "none" } + "offload_optimizer": { "device": "none" }, + "offload_param": { "device": "none" }, + + "stage3_gather_16bit_weights_on_model_save": false }, "optimizer": { diff --git a/ds_config_zero3_cpu_offload.json b/ds_config_zero3_cpu_offload.json new file mode 100644 index 0000000..48b3f7c --- /dev/null +++ b/ds_config_zero3_cpu_offload.json @@ -0,0 +1,27 @@ +{ + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + + "zero_optimization": { + "stage": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "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" } + }, + + "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 }, + + "gradient_clipping": 1.0, + "wall_clock_breakdown": false +}