mirror of https://github.com/hpcaitech/ColossalAI
[plugin] hotfix zero plugin (#6036)
* [plugin] hotfix zero plugin * [plugin] hotfix zero pluginpull/6038/head
parent
d383449fc4
commit
cc1b0efc17
|
@ -100,7 +100,7 @@ class LowLevelZeroModel(ModelWrapper, AMPModelMixin):
|
||||||
if self.convert_fn is not None:
|
if self.convert_fn is not None:
|
||||||
args = tree_map(self.convert_fn, args)
|
args = tree_map(self.convert_fn, args)
|
||||||
kwargs = tree_map(self.convert_fn, kwargs)
|
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:
|
with ctx:
|
||||||
return super().forward(*args, **kwargs)
|
return super().forward(*args, **kwargs)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue