Skip to content

fix errors when setting zero3 leaf modules with torch.compile #6564

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
loadams merged 4 commits into from
Sep 26, 2024
Merged

fix errors when setting zero3 leaf modules with torch.compile #6564

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
53155e8
fix errors when set zero3 leaf modules
NirSonnenschein Sep 24, 2024
ae82739
Merge branch 'master' into Fix_errors_when_set_zero3_leaf_modules
tjruwase Sep 24, 2024
8e3c331
Merge branch 'master' into Fix_errors_when_set_zero3_leaf_modules
tjruwase Sep 26, 2024
2ffcad7
Merge branch 'master' into Fix_errors_when_set_zero3_leaf_modules
loadams Sep 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions deepspeed/runtime/zero/parameter_offload.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ def _apply_forward_and_backward_to_tensors_only(module, forward_function, backwa

class ZeROOrderedDict(OrderedDict):

def __init__(self, parent_module, *args, **kwargs):
def __init__(self, parent_module=None, *args, **kwargs):
"""A replacement for ``collections.OrderedDict`` to detect external ZeRO params.

Args:
Expand All @@ -56,7 +56,7 @@ def __getitem__(self, key):
if param is None:
return param

if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
if hasattr(param, "ds_status") and param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
if self._parent_module._parameters._in_forward:
register_external_parameter(FWD_MODULE_STACK[-1], param)
param.all_gather()
Expand Down