mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-27 19:36:13 +00:00
ccabcf6485
* [fp8] support fp8 amp for hybrid parallel plugin * [test] add fp8 hook test * [fp8] fix fp8 linear compatibility
24 lines
521 B
Python
24 lines
521 B
Python
import torch.nn.functional as F
|
|
|
|
from colossalai.quantization.fp8 import linear_fp8
|
|
from colossalai.tensor.param_op_hook import ColoParamOpHook
|
|
|
|
|
|
class FP8Hook(ColoParamOpHook):
|
|
def pre_forward(self, params) -> None:
|
|
pass
|
|
|
|
def post_forward(self, params) -> None:
|
|
pass
|
|
|
|
def pre_backward(self, params) -> None:
|
|
pass
|
|
|
|
def post_backward(self, params) -> None:
|
|
pass
|
|
|
|
def rewrite_op(self, func):
|
|
if func is F.linear:
|
|
return linear_fp8
|
|
return func
|