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

[NFC] polish colossalai/amp/naive_amp/_fp16_optimizer.py code style (#1819)

pull/1849/head
Fazzie-Maqianli 2022-11-08 15:07:02 +08:00 committed by binmakeswell
parent 9623ec1b02
commit 399f84d8f6
1 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
colossalai/amp/naive_amp/_fp16_optimizer.py
View File

@ -9,14 +9,16 @@ try:
except: except:
print('Colossalai should be built with cuda extension to use the FP16 optimizer') print('Colossalai should be built with cuda extension to use the FP16 optimizer')
from torch.optim import Optimizer
from colossalai.core import global_context as gpc
from colossalai.context import ParallelMode
from colossalai.logging import get_dist_logger
from colossalai.utils import (copy_tensor_parallel_attributes, clip_grad_norm_fp32, multi_tensor_applier)
from torch.distributed import ProcessGroup from torch.distributed import ProcessGroup
from .grad_scaler import BaseGradScaler from torch.optim import Optimizer
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.logging import get_dist_logger
from colossalai.utils import clip_grad_norm_fp32, copy_tensor_parallel_attributes, multi_tensor_applier
from ._utils import has_inf_or_nan, zero_gard_by_list from ._utils import has_inf_or_nan, zero_gard_by_list
from .grad_scaler import BaseGradScaler
__all__ = ['FP16Optimizer'] __all__ = ['FP16Optimizer']