Loss Scaling Skip Step

Problem

Handle one mixed-precision step with dynamic loss scaling. If a gradient overflowed, skip the step (True), back off the scale and reset the good-step counter. Otherwise count a good step; once good_steps reaches growth_interval, grow the scale and reset the counter. Return [skip, new_loss_scale, new_good_steps].

The signature is loss_scaling_skip_step(loss_scale, found_inf, good_steps, growth_interval, growth_factor=2.0, backoff_factor=0.5).

Hint

Skipping the optimizer step is essential because overflowed gradients would corrupt the weights.

Test cases

case 1

in: [1024,true,5,10]

expected: [true,512,0]

case 2

in: [1024,false,0,2]

expected: [false,1024,1]

case 3

in: [1024,false,1,2]

expected: [false,2048,0]

case 4

in: [100,true,0,5,1.5,0.1]

expected: [true,10,0]

Starter codePython

def loss_scaling_skip_step(loss_scale, found_inf, good_steps, growth_interval, growth_factor=2.0, backoff_factor=0.5):
    # Your code here
    pass

4 test cases · run your code against them in the editor