Hybrid Precision Computation for Deep Learning

In modern deep learning, training efficiency is paramount. Large-scale neural networks can have billions of parameters, making memory bandwidth and computational throughput critical bottlenecks. Hybrid Precision Computation is a powerful technique that strategically combines different numerical precisions—using lower precision (float16 or bfloat16) for speed-critical operations while maintaining higher precision (float32) where numerical stability is essential.

The core insight is that neural network forward passes are robust to reduced precision because the values involved are typically within a reasonable range, and small rounding errors don't significantly affect the outcome. However, gradients during backpropagation can be very small (especially in deep networks), risking underflow to zero in reduced precision formats.

Loss Scaling addresses this challenge elegantly:

1. Forward Pass: Convert inputs to float16, perform computations, then scale the loss by a large factor (e.g., 1024.0)
2. Backward Pass: Compute gradients normally, then unscale them by dividing by the same factor
3. Gradient Checking: Detect overflow (inf/nan values) that may occur; if detected, skip the parameter update

This approach preserves representable gradient magnitudes in float16 range while keeping the final gradient values accurate in float32.

Your Task: Implement a HybridPrecisionTrainer class in Python using NumPy that performs hybrid precision training:

1. __init__(self, loss_scale=1024.0): Initialize with a loss scaling factor

2. forward(self, weights, inputs, targets):

   • Convert inputs to float16 for computation
   • Compute predictions as the dot product of inputs and weights
   • Calculate Mean Squared Error (MSE) loss: $\text{MSE} = \frac{1}{n}\sum_{i=1}^{n}(\text{pred}_i - \text{target}_i)^2$
   • Scale the loss by the loss_scale factor
   • Return the scaled loss as a Python float (float32 precision)

3. backward(self, gradients):

   • Convert gradients to float32 if not already
   • Unscale gradients by dividing by loss_scale
   • Check for overflow (any inf or nan values in the result)
   • If overflow detected, return an array of zeros with the same shape
   • Return the unscaled gradients as a float32 numpy array

weights = [0.5, -0.3]
inputs = [[1.0, 2.0], [3.0, 4.0]]
targets = [1.0, 0.0]
gradients = [512.0, -256.0]
loss_scale = 1024.0
method = "combined"

{"loss": 665.0, "gradients": [0.5, -0.25]}

weights = [1.0, 1.0]
inputs = [[1.0, 0.0], [0.0, 1.0]]
targets = [1.0, 1.0]
loss_scale = 1024.0
method = "forward"

0.0

gradients = [1024.0, 2048.0, 512.0]
loss_scale = 1024.0
method = "backward"

[1.0, 2.0, 0.5]