[fp8] fix linear hook

hpcaitech · Sep 3, 2024 · fc76bea · fc76bea
1 parent e9032fb
commit fc76bea
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@@ -119,7 +119,8 @@ def __init__(
         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