[fp8] fix linear hook (#6046)

2024-09-03 16:37:16 +08:00 · 2024-09-03 16:37:16 +08:00 · 26e553937b
parent c3b5caff0e
commit 26e553937b
1 changed files with 2 additions and 1 deletions
--- a/colossalai/booster/plugin/hybrid_parallel_plugin.py
+++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@ -119,7 +119,8 @@ class HybridParallelModule(ModelWrapper, AMPModelMixin):
        if use_fp8:
            self.op_hooks.append(FP8Hook())
        if overlap_allgather:
-            self.op_hook = ZeroOpHook()
+            self.op_hooks.append(ZeroOpHook())
+        if use_fp8 or overlap_allgather:
            for p in module.parameters():
                if p.requires_grad and type(p) is not ColoParameter:
                    p.__class__ = ColoParameter