turn off offload_params since ZeRO-2 doesn't have that option.
The performance will likely improve significantly with just offload_params turned off, even if you don't change
stage3_param_persistence_threshold. Of course, these changes will impact the size of the model you can train. So
these help you to trade scalability for speed depending on your needs.
ZeRO-2 Example
Here is a full ZeRO-2 auto-configuration file ds_config_zero2.json:
```json
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
"optimizer": {
    "type": "AdamW",
    "params": {
        "lr": "auto",
        "betas": "auto",
        "eps": "auto",
        "weight_decay": "auto"
