InternLM/internlm/model/__init__.py

32 lines
944 B
Python

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from .embedding import Embedding1D, RotaryEmbedding
from .linear import FeedForward, RewardModelLinear, ScaleColumnParallelLinear
from .metrics import AccPerplex
from .modeling_internlm import build_model_with_cfg
from .modeling_llama import build_model_with_cfg as build_model_with_llama_cfg
from .modeling_llama_moe import (
build_model_with_moe_cfg as build_model_with_llama_moe_cfg,
)
from .modeling_moe import build_model_with_moe_cfg
from .moe import MoE
from .multi_head_attention import MHA
from .utils import gather_forward_split_backward
__all__ = [
"Embedding1D",
"FeedForward",
"MoE",
"RotaryEmbedding",
"RewardModelLinear",
"ScaleColumnParallelLinear",
"AccPerplex",
"MHA",
"gather_forward_split_backward",
"build_model_with_cfg",
"build_model_with_moe_cfg",
"build_model_with_llama_cfg",
"build_model_with_llama_moe_cfg",
]