/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/quantization/fp8_hook.py
23 строки
521 B
Hongxin Liu
[fp8] support gemini plugin (#5978)
09 авг 2024, 09:09
Не верифицирован
09 авг 2024, 09:09
8241c0c
Код
Авторство
О чём код?
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