ColossalAI/colossalai/inference/modeling/policy/__init__.py

10 lines
318 B
Python
Raw Normal View History

from .glide_llama import GlideLlamaModelPolicy
from .nopadding_llama import NoPaddingLlamaModelInferPolicy
model_policy_map = {
"nopadding_llama": NoPaddingLlamaModelInferPolicy,
"glide_llama": GlideLlamaModelPolicy,
}
__all__ = ["NoPaddingLlamaModelInferPolicy", "GlideLlamaModelPolicy", "model_polic_map"]