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

add moe module to `__init__.py`

pull/182/head
zhanglei 2023-09-25 15:36:35 +08:00
parent 85f4d4af58
commit 1986116527
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internlm/model/__init__.py
View File

@ -5,12 +5,14 @@ from .embedding import Embedding1D, RotaryEmbedding
from .linear import FeedForward, RewardModelLinear, ScaleColumnParallelLinear from .linear import FeedForward, RewardModelLinear, ScaleColumnParallelLinear
from .metrics import AccPerplex from .metrics import AccPerplex
from .modeling_internlm import build_model_with_cfg from .modeling_internlm import build_model_with_cfg
from .moe import MoE
from .multi_head_attention import MHA from .multi_head_attention import MHA
from .utils import gather_forward_split_backward from .utils import gather_forward_split_backward
__all__ = [ __all__ = [
"Embedding1D", "Embedding1D",
"FeedForward", "FeedForward",
"MoE",
"RotaryEmbedding", "RotaryEmbedding",
"RewardModelLinear", "RewardModelLinear",
"ScaleColumnParallelLinear", "ScaleColumnParallelLinear",