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

import internlm.moe

pull/567/head
Wenwen Qu 2024-01-12 14:01:23 +08:00
parent 8acf8455fe
commit ebf2942746
2 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
internlm/model/moe.py
View File

@ -1,5 +1,6 @@
import torch import torch
import internlm.moe # noqa # pylint: disable=W0611
from internlm.core.context import ParallelMode from internlm.core.context import ParallelMode
from internlm.core.context import global_context as gpc from internlm.core.context import global_context as gpc
from internlm.model.linear import FeedForward from internlm.model.linear import FeedForward

2
internlm/moe/__init__.py
View File

@ -1,3 +1,3 @@
from internlm.moe.sharded_moe import GShardMOELayer from .sharded_moe import GShardMOELayer
__all__ = ["GShardMOELayer"] __all__ = ["GShardMOELayer"]