Skip to content

Add and Remove ZeRO 3 Hooks#5658

Merged
jomayeri merged 20 commits intomasterfrom
jomayeri/zero3-hooks
Aug 16, 2024
Merged

Add and Remove ZeRO 3 Hooks#5658
jomayeri merged 20 commits intomasterfrom
jomayeri/zero3-hooks

Conversation

@jomayeri
Copy link
Copy Markdown
Collaborator

Gives the ability to add and remove the forward hooks in ZeRO 3 by using a context manager. These code changes were taken from a Huggingface PR and integrated for direct support in DeepSpeed.

This is useful in the inference case and the speedup can be observed here.

@tjruwase
Copy link
Copy Markdown
Contributor

@jomayeri, please add some unit tests

@tjruwase tjruwase requested review from tohtana and removed request for mrwyattii June 18, 2024 00:25
Comment thread deepspeed/runtime/zero/utils.py Outdated
Comment thread deepspeed/runtime/zero/utils.py Outdated
Comment thread deepspeed/runtime/zero/utils.py Outdated
@jomayeri jomayeri requested review from arashb and awan-10 as code owners June 27, 2024 18:59
@jomayeri jomayeri requested a review from loadams as a code owner July 16, 2024 19:47
Comment thread csrc/aio/py_lib/deepspeed_py_copy.cpp
Comment thread tests/unit/runtime/zero/test_zero_context.py Outdated


@contextmanager
def unwrap_model_for_generation(model):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems my comment here was missed?

Suggested change
def unwrap_model_for_generation(model):
def unshard_and_remove_hooks(model):

@tjruwase
Copy link
Copy Markdown
Contributor

@jomayeri, can you please add documentation of this new API somewhere below?

https://deepspeed.readthedocs.io/en/latest/zero3.html

@jomayeri jomayeri added this pull request to the merge queue Aug 16, 2024
Merged via the queue into master with commit c2e3a70 Aug 16, 2024
mauryaavinash95 pushed a commit to DataStates/DeepSpeed that referenced this pull request Mar 20, 2025
Gives the ability to add and remove the forward hooks in ZeRO 3 by using
a context manager. These code changes were taken from a Huggingface
[PR](huggingface/trl#1617) and integrated for
direct support in DeepSpeed.

This is useful in the inference case and the speedup can be observed
[here](huggingface/trl#1483).

---------

Co-authored-by: root <root@deepspeed-c000004.2d1icxc5dsxehnpuwt3ifc34ph.gvxx.internal.cloudapp.net>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Heyang Qin <heyangqin@microsoft.com>
Co-authored-by: Logan Adams <loadams@microsoft.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.

4 participants