ColossalAI/colossalai/booster/mixed_precision/fp16_naive.py

from .mixed_precision_base import MixedPrecision


class FP16NaiveMixedPrecision(MixedPrecision):
    """
    Precision for mixed precision training in FP16 using naive AMP.

    Args:
    log_num_zeros_in_grad(bool): return number of zeros in the gradients.
    initial_scale(int): initial scale of gradient scaler.
    growth_factor(int): the growth rate of loss scale.
    backoff_factor(float): the decrease rate of loss scale.
    hysteresis(int): delay shift in dynamic loss scaling.
    max_scale(int): maximum loss scale allowed.
    verbose(bool): if set to `True`, will print debug info.
    """

    def __init__(self,
                 log_num_zeros_in_grad: bool,
                 initial_scale: int,
                 growth_factor: int,
                 backoff_factor: float,
                 hysteresis: int,
                 max_scale: int,
                 verbose: bool = None) -> None:
        pass