/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/zero/low_level/zero_hook.py
33 строки
800 B
Hongxin Liu
[zero] support all-gather overlap (#5898)
11 июл 2024, 13:59
Не верифицирован
11 июл 2024, 13:59
c068ef0
Код
Авторство
О чём код?
from typing import List from torch._tensor import Tensor from colossalai.tensor.param_op_hook import ColoParamOpHook _ALL_GATHER_HANDLE = "_all_gather_handle" def wait_all_gather_handle(p): if hasattr(p, _ALL_GATHER_HANDLE): handle = getattr(p, _ALL_GATHER_HANDLE) handle.wait() delattr(p, _ALL_GATHER_HANDLE) def set_all_gather_handle(p, handle): setattr(p, _ALL_GATHER_HANDLE, handle) class ZeroOpHook(ColoParamOpHook): def pre_forward(self, params: List[Tensor]) -> None: for p in params: wait_all_gather_handle(p) def post_forward(self, params: List[Tensor]) -> None: pass def pre_backward(self, params: List[Tensor]) -> None: pass def post_backward(self, params: List[Tensor]) -> None: pass