Implement a mixed-precision loop (bf16 compute, fp32 master weights) with loss scaling on

Super-hard

Optimization & Training Dynamics5 / 5

Chapter 075 / 5

Optimization & Training Dynamics

Implement a mixed-precision loop (bf16 compute, fp32 master weights) with loss scaling on

Super-hard

Implement a mixed-precision loop (bf16 compute, fp32 master weights) with loss scaling on a toy MLP, including a deliberate overflow and recovery.

Target interface

forward_backward(master, x, y, loss_scale)

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.

solution.pyPython

Local draft

import torch

def forward_backward(master, x, y, loss_scale):
    raise NotImplementedError

def is_finite(grads):
    raise NotImplementedError

⌘/Ctrl↵to submit

AI review

No review yet

Submit your solution for a verdict, a score out of 100 and concrete feedback. Any correct approach passes.

Chapter 075 / 5

Optimization & Training Dynamics

Implement a mixed-precision loop (bf16 compute, fp32 master weights) with loss scaling on

Super-hard

Implement a mixed-precision loop (bf16 compute, fp32 master weights) with loss scaling on a toy MLP, including a deliberate overflow and recovery.

Target interface

forward_backward(master, x, y, loss_scale)

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.