mirror of https://github.com/hpcaitech/ColossalAI
34 lines
962 B
Python
34 lines
962 B
Python
from .cache_embedding import (
|
|
CachedEmbeddingBag,
|
|
CachedParamMgr,
|
|
EvictionStrategy,
|
|
LimitBuffIndexCopyer,
|
|
ParallelCachedEmbeddingBag,
|
|
ParallelCachedEmbeddingBagTablewise,
|
|
ParallelCachedEmbeddingBagTablewiseSpiltCache,
|
|
TablewiseEmbeddingBagConfig,
|
|
)
|
|
from .colo_module import ColoModule
|
|
from .embedding import ColoEmbedding
|
|
from .linear import ColoLinear
|
|
from .module_utils import check_colo_module, get_colo_module, init_colo_module, is_colo_module, register_colo_module
|
|
|
|
__all__ = [
|
|
"ColoModule",
|
|
"register_colo_module",
|
|
"is_colo_module",
|
|
"get_colo_module",
|
|
"init_colo_module",
|
|
"check_colo_module",
|
|
"ColoLinear",
|
|
"ColoEmbedding",
|
|
"CachedEmbeddingBag",
|
|
"ParallelCachedEmbeddingBag",
|
|
"CachedParamMgr",
|
|
"LimitBuffIndexCopyer",
|
|
"EvictionStrategy",
|
|
"ParallelCachedEmbeddingBagTablewise",
|
|
"TablewiseEmbeddingBagConfig",
|
|
"ParallelCachedEmbeddingBagTablewiseSpiltCache",
|
|
]
|