Skip to content

Commit 26e5539

Browse files
authored
[fp8] fix linear hook (#6046)
1 parent c3b5caf commit 26e5539

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

colossalai/booster/plugin/hybrid_parallel_plugin.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,8 @@ def __init__(
119119
if use_fp8:
120120
self.op_hooks.append(FP8Hook())
121121
if overlap_allgather:
122-
self.op_hook = ZeroOpHook()
122+
self.op_hooks.append(ZeroOpHook())
123+
if use_fp8 or overlap_allgather:
123124
for p in module.parameters():
124125
if p.requires_grad and type(p) is not ColoParameter:
125126
p.__class__ = ColoParameter

0 commit comments

Comments
 (0)