huggingface / huggingface/accelerate

DeepSpeedEngineWrapper.backward() does a bit too much

Open
#2,951 19 comments 13 reactions 0 assignees View on GitHub
enhancement feature request
Dominant language
Python
Stars
9.9k
Forks
1.5k
Avg merge
5d 2h
Merged PRs (30d)
27

Description

The source code of DeepSpeedEngineWrapper:
```
class DeepSpeedEngineWrapper:
"""
Internal wrapper for deepspeed.runtime.engine.DeepSpeedEngine. This is used to follow conventional training loop.

Args:
engine (deepspeed.runtime.engine.DeepSpeedEngine): deepspeed engine to wrap
"""

def __init__(self, engine):
self.engine = engine

def backward(self, loss, **kwargs):
# runs backpropagation and handles mixed precision
self.engine.backward(loss, **kwargs)

# Deepspeed's `engine.step` performs the following operations:
# - gradient accumulation check
# - gradient clipping
# - optimizer step
# - zero grad
# - checking overflow
# - lr_scheduler step (only if engine.lr_scheduler is not None)
self.engine.step()
# and this plugin overrides the above calls with no-ops when Accelerate runs under
# Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
# training loop that works transparently under many training regimes.
```

My question is: Why do we need to do `self.engine.step()` here immediately? This behavior zeros grad and change the parameter without noticing the user. It might be out of expectation. Since backward step is internally binded with zeroing grad and changing parameter, this blocks users from checking the gradient or parameter manually before stepping.

I know deepspeed-wrapped models can't be seen as normal models, but this behavior still elimiates a lot of flexibility.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.