Skip to content

Commit

Permalink
[fp8] fix linear hook (#6046)
Browse files Browse the repository at this point in the history
  • Loading branch information
ver217 authored Sep 3, 2024
1 parent c3b5caf commit 26e5539
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion colossalai/booster/plugin/hybrid_parallel_plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 26e5539

Please sign in to comment.