From 7ad8cd7cae20a2503293142f73481c0b64e0116e Mon Sep 17 00:00:00 2001 From: ver217 Date: Tue, 27 Aug 2024 18:14:54 +0800 Subject: [PATCH 1/2] [plugin] hotfix zero plugin --- colossalai/booster/plugin/low_level_zero_plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py index 42bb49bc9ca1..2a2b86a03613 100644 --- a/colossalai/booster/plugin/low_level_zero_plugin.py +++ b/colossalai/booster/plugin/low_level_zero_plugin.py @@ -100,7 +100,7 @@ def forward(self, *args, **kwargs): 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) From 3cef5c982607ba179bf9bfe2c82050870334d907 Mon Sep 17 00:00:00 2001 From: ver217 Date: Tue, 27 Aug 2024 18:16:04 +0800 Subject: [PATCH 2/2] [plugin] hotfix zero plugin --- colossalai/booster/plugin/low_level_zero_plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py index 2a2b86a03613..8cc511a5610f 100644 --- a/colossalai/booster/plugin/low_level_zero_plugin.py +++ b/colossalai/booster/plugin/low_level_zero_plugin.py @@ -100,7 +100,7 @@ def forward(self, *args, **kwargs): 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_hooks) 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)