Skip to content

Commit

Permalink
Enable torch compile on _allgather_params (#6769)
Browse files Browse the repository at this point in the history
* Previosuly ZerO3 was crashing when trying to compile _allgather_params
* Disabling grad solves the issue
  • Loading branch information
deepcharm authored Nov 21, 2024
1 parent 83e4364 commit b5709cc
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions deepspeed/runtime/zero/partition_parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -1882,6 +1882,7 @@ def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False):

return None

@torch.no_grad()
def _allgather_params(self, param_list, hierarchy=0):
if len(param_list) == 0:
return
Expand Down

0 comments on commit b5709cc

Please sign in to comment.