github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-03 14:08:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

[plugin] hotfix zero plugin ()

Browse Source 
* [plugin] hotfix zero plugin

* [plugin] hotfix zero plugin
This commit is contained in:
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

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)