ZeRO stage 1 refresh

[jeffra]

This PR changes the default underlying implementation of ZeRO stage 1. We have had trouble keeping up with bug fixes and maintaining feature parity across ZeRO implementations (e.g., cpu offload, overlap, fp32 support). Instead we are deprecating the original stage 1 implementation in favor of a new mode of stage 2 that supports optimizer state partitioning only. A few implications to be aware of when this PR is merged.

ZeRO stage 1 now supports

• Optimizer state offload to CPU
• Communication overlap
  • This is done as in stage 2 through the use of backward hooks to reduce gradients when they become available. This can be turned on/off with "overlap_comm": [true|false]
• Pipeline parallelism + zero stage 1 now supports optimizer state offload to CPU
• FP32/TF32 support

Warnings

• ZeRO stage 1 checkpoints created before this PR will not be compatible with the new ZeRO stage 1
  • Until we release v0.4.0 we will support a new zero param JSON of "legacy_stage": true to use the old codebase.

Misc

• Memory overhead bug fix related to both zero stage 1 and 2. Fixes bug where if a single model parameter was larger than our reduce bucket size we would error our, which required arbitrarily inflating the bucket size (x4.5) in certain cases to accommodate the large param(s).
• Fixes a perf bug in pipeline parallelism where we were all-reducing all model gradients instead if relying on reduce(scatter) in zero for reduce comm volume.
• Adds DeepSpeed version info to checkpoints