2022-06-06 07:34:41 +00:00
|
|
|
from .colo_module import ColoModule
|
|
|
|
|
from .linear import ColoLinear
|
|
|
|
|
from .embedding import ColoEmbedding
|
|
|
|
|
from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module
|
|
|
|
|
|
2022-09-01 09:55:41 +00:00
|
|
|
from .cache_embedding import FreqAwareEmbeddingBag, ParallelFreqAwareEmbeddingBag, CachedParamMgr, LimitBuffIndexCopyer, EvictionStrategy, \
|
2022-09-05 07:12:53 +00:00
|
|
|
ParallelFreqAwareEmbeddingBagTablewise, TablewiseEmbeddingBagConfig, ParallelFreqAwareEmbeddingBagTablewiseSpiltCache
|
2022-08-11 05:43:24 +00:00
|
|
|
|
2022-06-06 07:34:41 +00:00
|
|
|
__all__ = [
|
2022-08-11 05:43:24 +00:00
|
|
|
'ColoModule', 'register_colo_module', 'is_colo_module', 'get_colo_module', 'init_colo_module', 'check_colo_module',
|
|
|
|
|
'ColoLinear', 'ColoEmbedding', 'FreqAwareEmbeddingBag', 'ParallelFreqAwareEmbeddingBag', 'CachedParamMgr',
|
2022-09-05 07:12:53 +00:00
|
|
|
'LimitBuffIndexCopyer', 'EvictionStrategy', 'ParallelFreqAwareEmbeddingBagTablewise', 'TablewiseEmbeddingBagConfig',
|
|
|
|
|
'ParallelFreqAwareEmbeddingBagTablewiseSpiltCache'
|
2022-06-06 07:34:41 +00:00
|
|
|
]
|
