github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/engine/gradient_handler/_moe_gradient_handler.py (#3260)

pull/3313/head
LuGY 2023-03-27 18:47:44 +08:00 committed by binmakeswell
parent 204ca2f09a
commit 1ff7d5bfa5
1 changed files with 46 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/engine/gradient_handler/_moe_gradient_handler.py
View File

@ -1,10 +1,11 @@
from colossalai.context.moe_context import MOE_CONTEXT
from colossalai.core import global_context as gpc from colossalai.core import global_context as gpc
from colossalai.registry import GRADIENT_HANDLER from colossalai.registry import GRADIENT_HANDLER
from colossalai.utils.moe import get_moe_epsize_param_dict from colossalai.utils.moe import get_moe_epsize_param_dict
from ._base_gradient_handler import BaseGradientHandler
from ...context.parallel_mode import ParallelMode from ...context.parallel_mode import ParallelMode
from ._base_gradient_handler import BaseGradientHandler
from .utils import bucket_allreduce from .utils import bucket_allreduce
from colossalai.context.moe_context import MOE_CONTEXT
@GRADIENT_HANDLER.register_module @GRADIENT_HANDLER.register_module