colossal AI使用amp方式

有三种方式可以选择：

torch amp

native amp

apex amp

from colossalai.amp import AMP_TYPE

# use Torch AMP
fp16=dict(
    mode = AMP_TYPE.TORCH  
)

# use naive AMP
fp16=dict(
    mode = AMP_TYPE.NAIVE  
)

# use NVIDIA Apex AMP
fp16=dict(
    mode = AMP_TYPE.APEX  
)

torch amp config

from colossalai.amp import AMP_TYPE

fp16=dict(
    mode=AMP_TYPE.TORCH,

    # below are default values for grad scaler
    init_scale=2.**16,
    growth_factor=2.0,
    backoff_factor=0.5,
    growth_interval=2000,
    enabled=True
)

native amp config

from colossalai.amp import AMP_TYPE

fp16 = dict(
    mode=AMP_TYPE.NAIVE,
    
    # below are the default values
    log_num_zeros_in_grad=False,
    initial_scale=2 ** 32,
    min_scale=1,
    growth_factor=2,
    backoff_factor=0.5,
    growth_interval=1000,
    hysteresis=2
)

apex amp

from colossalai.amp import AMP_TYPE

fp16 = dict(
    mode=AMP_TYPE.APEX,
    
    # below are the default values
    enabled=True, 
    opt_level='O1', 
    cast_model_type=None, 
    patch_torch_functions=None, 
    keep_batchnorm_fp32=None, 
    master_weights=None, 
    loss_scale=None, 
    cast_model_outputs=None,
    num_losses=1, 
    verbosity=1, 
    min_loss_scale=None, 
    max_loss_scale=16777216.0
)

参考：

https://colossalai.org/docs/features/mixed_precision_training/