github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[plugin] hotfix zero plugin (#6036) 

* [plugin] hotfix zero plugin

* [plugin] hotfix zero plugin
pull/6038/head
Hongxin Liu 2024-08-28 10:16:48 +08:00 committed by GitHub
parent d383449fc4
commit cc1b0efc17
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/booster/plugin/low_level_zero_plugin.py
View File

@ -100,7 +100,7 @@ class LowLevelZeroModel(ModelWrapper, AMPModelMixin):
if self.convert_fn is not None:
args = tree_map(self.convert_fn, args)
kwargs = tree_map(self.convert_fn, kwargs)
ctx = ColoParamOpHookManager.use_hooks(self.op_hook) if self.overlap_allgather else nullcontext()
ctx = ColoParamOpHookManager.use_hooks(*self.op_hooks) if self.overlap_allgather else nullcontext()
with ctx:
return super().forward(*args, **kwargs)