Optimistic Verifiable Training by Controlling Hardware Nondeterminism

TL;DR

This work proposes a method that combines training in a higher precision than the target, rounding after intermediate computations, and sharing rounding decisions based on an adaptive thresholding procedure, to successfully control for nondeterminism in verifiable training.

Abstract

The increasing compute demands of AI systems have led to the emergence of services that train models on behalf of clients lacking necessary resources. However, ensuring correctness of training and guarding against potential training-time attacks, such as data poisoning and backdoors, poses challenges. Existing works on verifiable training largely fall into two classes: proof-based systems, which are difficult to scale, and ``optimistic'' methods that consider a third-party auditor who can replicate the training process and contest the trainer. A key challenge with the latter is that nondeterminism between GPU types during training prevents exact replication of the training process, resulting in schemes that are non-robust. We propose a method that combines training in a higher precision than the target, rounding after intermediate computations, and sharing rounding decisions based on an adaptive thresholding procedure, to successfully control for nondeterminism. Across three different NVIDIA GPUs (A40, Titan XP, RTX 2080 Ti), we achieve exact training replication at FP32 precision for both full-training and fine-tuning of ResNet-50 (23M) and GPT-2 (117M) models. Our verifiable training scheme significantly decreases the storage and time costs compared to proof-based systems, and is publicly released at https://github.com/meghabyte/verifiable-training.

Figures (5)

• Figure 1: Overview of our scheme. After an $\mathrm{auditor}$ challenges a $\mathrm{trainer}$, they train the model, storing weights in a Merkle tree, and enter a binary search procedure to identify the exact steps of the dispute. We show how to control GPU nondeterminism between $\mathrm{auditor}$ and $\mathrm{trainer}$, expanding the set of potential auditors.
• Figure 2: Even after ensuring the same software version, random seed, and use of deterministic algorithms via library flags, training nondeterminism persists between three GPU types.
• Figure 3: Divergence between outputs on two different GPUs (in FP64) for a given function and input can result in different rounding choices when rounding to the nearest FP32. We only wish to log rounding decisions for Case A, where the $\mathrm{auditor}$ should copy the $\mathrm{trainer}$'s rounding choice in order to reach the same value. This requires defining a logging region, determined by a threshold $\tau$,
• Figure 4: We successfully control for nondeterminism between GPU types for both ResNet-50 (a.) and GPT-2 (b.) tasks, while standard training and simple rounding without performing $\mathsf{rev}$ corrections result in model divergence over the course of training. Stronger rounding has minimal affect to model performance (c.), but at the cost of increasing time for $\mathrm{trainer}$ (d.).
• Figure 5: We define rounding to $b$ bits as rounding to the nearest 32-bit FP number that has 0s in the last $32-b$ bits of the mantissa, after accounting for the exponent.