github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 20:40:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] solve hang

Browse Source
This commit is contained in:
hxwang
2024-07-05 07:19:37 +00:00
committed by Hongxin Liu
parent 0fad23c691
commit 46c069b0db
12 changed files with 113 additions and 390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/zero/low_level/bookkeeping/gradient_store.py
View File

@@ -91,7 +91,7 @@ class GradientStore(BaseStore):
return grad_list
def get_working_grad_by_param_id(self, param_id) -> Tensor:
def get_working_grad_by_param_id(self, param_id) -> Optional[Tensor]:
"""
Return the working gradient for the specified parameter.