mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
36 lines
1.2 KiB
36 lines
1.2 KiB
from .bf16 import BF16MixedPrecision
|
|
from .fp8 import FP8MixedPrecision
|
|
from .fp16_apex import FP16ApexMixedPrecision
|
|
from .fp16_naive import FP16NaiveMixedPrecision
|
|
from .fp16_torch import FP16TorchMixedPrecision
|
|
from .mixed_precision_base import MixedPrecision
|
|
|
|
__all__ = [
|
|
'MixedPrecision', 'mixed_precision_factory', 'FP16_Apex_MixedPrecision', 'FP16_Torch_MixedPrecision',
|
|
'FP32_MixedPrecision', 'BF16_MixedPrecision', 'FP8_MixedPrecision', 'FP16NaiveMixedPrecision'
|
|
]
|
|
|
|
_mixed_precision_mapping = {
|
|
'fp16': FP16TorchMixedPrecision,
|
|
'fp16_apex': FP16ApexMixedPrecision,
|
|
'fp16_naive': FP16NaiveMixedPrecision,
|
|
'bf16': BF16MixedPrecision,
|
|
'fp8': FP8MixedPrecision
|
|
}
|
|
|
|
|
|
def mixed_precision_factory(mixed_precision_type: str) -> MixedPrecision:
|
|
"""
|
|
Factory method to create mixed precision object
|
|
|
|
Args:
|
|
mixed_precision_type (str): mixed precision type, including None, 'fp16', 'fp16_apex', 'bf16', and 'fp8'.
|
|
"""
|
|
|
|
if mixed_precision_type in _mixed_precision_mapping:
|
|
return _mixed_precision_mapping[mixed_precision_type]()
|
|
else:
|
|
raise ValueError(
|
|
f'Mixed precision type {mixed_precision_type} is not supported, support types include {list(_mixed_precision_mapping.keys())}'
|
|
)
|