Dynamic Loss Scaling (Mixed Precision) - Optimization & Training

Dynamic Loss Scaling for Mixed Precision

What you'll build

dynamic_loss_scale(overflows, ...): simulate how torch.amp.GradScaler adapts its loss scale over a training run, given which steps produced inf/NaN gradients.

Background

In fp16 the smallest normal number is about 6e-5 and anything below about 6e-8 rounds to 0, so many small gradients underflow to zero. The fix is to multiply the loss by a large scale S before backward(): every gradient is then S times larger and fits in fp16, and the gradients are divided by S again before the optimizer step. If S is too large, some gradient overflows to inf/NaN instead. Dynamic loss scaling keeps S as large as possible without overflowing: it backs off right after an overflow and grows again after a run of clean steps. (bf16 has fp32's exponent range, which is why bf16 training usually needs no loss scaling.)

Algorithm

S = init_scale; good_steps = 0
for i, overflow in enumerate(overflows):
    record S as the scale used at step i
    if overflow:                         # some gradient was inf/NaN
        skip the optimizer step (record i)
        S = S * backoff_factor           # e.g. halve
        good_steps = 0
    else:                                # optimizer step applied
        good_steps += 1
        if good_steps == growth_interval:
            S = S * growth_factor        # e.g. double
            good_steps = 0

PyTorch's defaults are init_scale=65536, growth_factor=2, backoff_factor=0.5 and growth_interval=2000.

Inputs

• overflows: list of bools; overflows[i] is True if step i's gradients contained inf/NaN (may be empty)
• init_scale, growth_factor, backoff_factor: floats
• growth_interval: positive int, the number of consecutive good steps before growing

Expected Return Format

Return a dictionary with these keys:

• 'scales': list of floats, the scale used at each step (before that step's update), same length as overflows
• 'skipped': list of the indices of the steps whose optimizer update was skipped, in increasing order
• 'final_scale': float, the scale after the last step (init_scale for an empty run)

Edge cases and conventions

• An overflow resets the good-step counter, so growth needs growth_interval clean steps in a row.
• After growing, the counter restarts from 0.
• Return floats for the scales even when init_scale is passed as an int.

Constraints

Plain Python; one pass over overflows.

Follow-up

• Why is skipping the step (rather than clipping the inf gradients) the right response to an overflow?
• If the scale keeps shrinking towards 0 during training, what does that tell you about the model?

Input: dynamic_loss_scale([False, True, False, False, False], init_scale=8.0, growth_interval=2)
Output: {'scales': [8.0, 8.0, 4.0, 4.0, 8.0], 'skipped': [1], 'final_scale': 8.0}

Input: dynamic_loss_scale([True, True, True], init_scale=1024.0)
Output: {'scales': [1024.0, 512.0, 256.0], 'skipped': [0, 1, 2], 'final_scale': 128.0}

def dynamic_loss_scale(overflows, init_scale=65536.0, growth_factor=2.0,
                       backoff_factor=0.5, growth_interval=2000) -> dict:
    """
    Simulate GradScaler's dynamic loss scale over a sequence of steps.

    Args:
        overflows: List of bools, True if that step's gradients had inf/nan
        init_scale: Starting scale
        growth_factor: Multiplier after growth_interval good steps in a row
        backoff_factor: Multiplier after an overflow
        growth_interval: Number of consecutive good steps before growing

    Returns:
        dict with 'scales' (scale used at each step), 'skipped' (indices of
        skipped steps) and 'final_scale'
    """
    # Your code here
    pass