Skip to content

ZeRO stage 1 refresh - #1042

Merged
jeffra merged 12 commits into
masterfrom
jeffra/z1-refresh-4
May 19, 2021
Merged

ZeRO stage 1 refresh#1042
jeffra merged 12 commits into
masterfrom
jeffra/z1-refresh-4

Conversation

@jeffra

@jeffra jeffra commented May 5, 2021

Copy link
Copy Markdown
Collaborator

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

Comment thread deepspeed/runtime/engine.py
@jeffra
jeffra merged commit cfa63f5 into master May 19, 2021
@mrwyattii
mrwyattii deleted the jeffra/z1-refresh-4 branch July 7, 2023 02:40
vineethsaivs added a commit to vineethsaivs/DeepSpeed that referenced this pull request Aug 12, 2026
config-json.md documents grad_hooks as a ZeRO option with a default of True, but
DeepSpeedZeroConfig has no such field and DeepSpeedConfigModel sets
extra="forbid", so following the docs is a hard failure:

    DeepSpeedZeroConfig(**{"stage": 1, "grad_hooks": False})
    pydantic_core._pydantic_core.ValidationError: 1 validation error
    grad_hooks
      Extra inputs are not permitted

The option was never wired up rather than removed later. cfa63f5 ("ZeRO stage 1
refresh", deepspeedai#1042) added the doc entry and DeepSpeedEngine.zero_grad_hooks() in the
same commit without ever adding the field to the config or to the constants
module, and it is absent from deepspeed/runtime/zero/config.py in every release
back to v0.3.0. So there is nothing to restore, and zero_grad_hooks() reads
zero_config.grad_hooks, which can only raise AttributeError. Nothing in the
package or the tests calls it.

Remove the doc entry and the dead accessor.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants