github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[amp] Add naive amp demo (#3774) 

* [mixed_precison] add naive amp demo

* [mixed_precison] add naive amp demo
pull/3780/head
jiangmingyan 2023-05-18 16:33:14 +08:00 committed by GitHub
parent 48bd056761
commit 2703a37ac9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/booster/mixed_precision/__init__.py
View File

@ -1,17 +1,19 @@
from .bf16 import BF16MixedPrecision from .bf16 import BF16MixedPrecision
from .fp8 import FP8MixedPrecision from .fp8 import FP8MixedPrecision
from .fp16_apex import FP16ApexMixedPrecision from .fp16_apex import FP16ApexMixedPrecision
from .fp16_naive import FP16NaiveMixedPrecision
from .fp16_torch import FP16TorchMixedPrecision from .fp16_torch import FP16TorchMixedPrecision
from .mixed_precision_base import MixedPrecision from .mixed_precision_base import MixedPrecision
__all__ = [ __all__ = [
'MixedPrecision', 'mixed_precision_factory', 'FP16_Apex_MixedPrecision', 'FP16_Torch_MixedPrecision', 'MixedPrecision', 'mixed_precision_factory', 'FP16_Apex_MixedPrecision', 'FP16_Torch_MixedPrecision',
'FP32_MixedPrecision', 'BF16_MixedPrecision', 'FP8_MixedPrecision' 'FP32_MixedPrecision', 'BF16_MixedPrecision', 'FP8_MixedPrecision', 'FP16NaiveMixedPrecision'
] ]
_mixed_precision_mapping = { _mixed_precision_mapping = {
'fp16': FP16TorchMixedPrecision, 'fp16': FP16TorchMixedPrecision,
'fp16_apex': FP16ApexMixedPrecision, 'fp16_apex': FP16ApexMixedPrecision,
'fp16_naive': FP16NaiveMixedPrecision,
'bf16': BF16MixedPrecision, 'bf16': BF16MixedPrecision,
'fp8': FP8MixedPrecision 'fp8': FP8MixedPrecision
} }

5
colossalai/booster/mixed_precision/fp16_naive.py Normal file
View File

@ -0,0 +1,5 @@
from .mixed_precision_base import MixedPrecision
class FP16NaiveMixedPrecision(MixedPrecision):
pass