Stochastic Operator Stability Monitor

(1) Integration point: wrap GEMM/linear layers, attention score products, softmax, layer normalization, reductions, and custom CUDA kernels with an optional StabilityMonitor hook. During calibration, retain the ordinary low-precision output y_v, then run K stochastic replicas for only a sampled tensor block or a random projection of the output. Do not replicate the entire model; monitor one operator per step in a rotating schedule, while always monitoring operators whose scores were previously high.

(2) Pseudocode:

for training step t:
    choose monitored node v from rotating schedule
    y = forward_operator_v(inputs, precision=bf16)
    if t % monitor_period == 0:
        ys = []
        for k in 1..K:
            perturbed_inputs = inputs * (1 + Uniform(-u, u))
            ys.append(forward_operator_v(perturbed_inputs, precision=bf16))
        mean = average(ys)
        S = sqrt(sum_k(norm(ys[k]-mean)^2)/(K-1)) / (norm(mean)+eps)
        score[v] = beta*score[v] + (1-beta)*S
        if score[v] > tau_warn:
            mode[v] = FP32_ACCUMULATION
        if score[v] > tau_safe:
            mode[v] = STABLE_KERNEL

For a linear layer, FP32_ACCUMULATION means BF16 inputs with FP32 accumulation. For softmax, use max subtraction and FP32 accumulation of exponentials. For normalization, compute mean and variance in FP32. Apply the selected mode on later steps and periodically recheck the node.

(3) Computed directly from the mathematical construction: graph-node identity, stochastic perturbed replicas, sample mean, relative output dispersion S_v, and threshold-based intervention. Estimate u from the actual hardware format, thresholds from clean-versus-injected-instability calibration, and the running score from minibatches. For large outputs, use a fixed random sign projection R and replace each Frobenius norm by ||R(y_v^(k)-mean)||.

(4) First cheap experiment: train a 6-layer Transformer on WikiText-2 with BF16 and compare uninstrumented BF16, all-FP32, and selective monitoring. Inject instability by forcing FP16 accumulation in a cancellation-heavy linear layer, removing max subtraction from softmax, or scaling activations toward overflow. Measure validation loss, NaN/Inf frequency, detection precision and recall, wall-clock overhead, and the fraction of operators promoted to FP32. Success means near-FP32 NaN prevention and stable loss with under 5-10% monitoring overhead and substantially fewer FP32 operations than the all-FP32 baseline.