ds_config_stage2.json

{
    "micro_batch_per_gpu": 5,
    "train_batch_size": 5,
    "steps_per_print": 2000,

    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 12,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

   "zero_optimization": {
       "stage": 2,
       "allgather_partitions": true,
       "allgather_bucket_size": 200000000,
       "overlap_comm": true,
       "reduce_scatter": true,
       "reduce_bucket_size": 200000000,
       "contiguous_gradients": true,
       "cpu_offload": false
   },

   "wall_clock_breakdown": false
}