Estimating LLM Training FLOPs on the Nvidia Jetson Orin Nano

This is a research summary for an ongoing project I am working on as part of the UChicago Existential Risks Laboratory Summer Research Fellowship. I would really appreciate any feedback.

In want of a quantifiable way to decide what counts as a frontier AI model, compute thresholds have emerged as the standard for AI policy: California's SB 53 uses 10^26 floating-point operations (FLOPs) in the training run as the threshold for what counts as a frontier model and the EU AI Act applies the same categorization at 10^25. Proposals for international AI agreements (example 1, example 2, example 3) extend the use of training FLOPs to determine part or all of the threshold for what counts as a frontier model under the agreement. Current AI laws have no way of actually verifying AI companies' claims about the number of FLOPs used in training and instead just rely on self-reports, but an international AI agreement can't assume compliance from each involved party. As such, we'd like to verify the number of FLOPs used in LLM training runs through side-channel GPU readings. This allows AI developers' code and data to remain hidden from the verifiers of the AI agreement, but allow verification of training FLOPs even under conditions where the model training might be adversarially changed to circumvent them. My work builds a Minimum Viable Product (MVP) for how this verification could work on an Nvidia Jetson Orin Nano.

EpochAI has done work on estimating the number of training FLOPs using (1) insider knowledge about the model architecture and (2) open-source reports of the number of accelerators, amount of time used in training runs, and expected utilization rate. My estimator extends their method (2) to include power-monitoring and real-time data on utilization.

Chaudhuri et al. demonstrate how thermal and power side-channels, along with assumptions about model architecture, can be used to extract information about transformer model weights during a training run. While an explicit goal of AI treaty verification is to NOT expose model secrets to the verifiers, this shows that extracting information about transformer training runs through side-channels is possible.[1]

Recently, Rahman and Tajdari showed how GPU system-level readings could be used to classify AI workloads as training vs. inference.

All experiments are run on an Nvidia Jetson Orin Nano 8 GB developer kit. This is essentially Nvidia's version of the Raspberry Pi, a single-board computer used for small-scale experiments and edge computing, but with an Ampere GPU. To estimate FLOPs, we run a sample LLM training script on the Nano and then observe the following signals:

Power (W): The Nano comes with an onboard INA3221 sensor which measures the power consumed by the board in three channels: VDD_IN - the total input power to the board, VDD_CPU_GPU_CV - the power consumed by the shared CPU/GPU/CV Accelerator rail, and VDD_SOC - the power draw of the memory subsystem. For the estimator, I use VDD_CPU_GPU_CV which means we have no way of separating CPU vs GPU power draw, however, we can expect GPU power to dominate the reading.

Memory Bandwidth Utilization (%): Actmon is a hardware activity monitor which reports whether the External Memory Controller (EMC) is active. The EMC is the hardware block which manages the transfer of data between the GPU and DRAM. What this reading does not tell us is how much data is transferred, just whether or not the controller is active. Still, a higher EMC utilization suggests there is a lot of data that needs to be transferred between the GPU from DRAM which can be indicative of a training run happening.

The Nano is not a device that is suitable for running an actual LLM training workload, so instead I use a script that mimics the behavior of a real workload. It initializes a transformer and then runs a training loop for a preset number of steps on randomly generated tokens and can be configured with a range of hyperparameter settings explained below.

Model Width (d_model): the size of the vector representing each token as it flows through the model. A larger model width will quadratically increase the size of matrix multiplications that happen during training.

Depth (n_layers): how many identical transformer blocks are there in total? The total number of operations in the training run will scale linearly with the number of blocks.

Attention heads (n_heads): How many parallel attention operations are used on the input token vector? This doesn't impact the total number of operations.

Sequence length (seq_len): how many tokens are processed together in parallel? This increases the computation done by the MLP linearly and the attention heads quadratically.

Batch size: how many sequences of tokens are processed in a single step? In a real LLM training setup, changing this would not be expected to have a big impact on FLOPs because although more sequences are processed in parallel, there is typically only a fixed number of tokens for the training run to get through. However, in this experimental setup, the tokens are made up as the training run goes for a fixed number of steps, so we should expect FLOPs to increase linearly with batch size.

Optimizer (Adam vs. SGD): The algorithm used for updating model weights during training. The size of these operations should be dwarfed by the much larger matrix multiplications.

Numeric Precision (FP32/TF32/FP16/BF16): What precision to represent every number in during the training run? Changing this shouldn't affect the ground truth FLOPs, but might confuse the estimator since an FP16 operation will be less intensive than an FP32 operation. For keeping this experiment simple, I am only testing on FP32, but future work will involve expanding this to work across multiple types of precision.

While the sample training run script can simulate a workload with any value of hyperparameters, we only care about the FLOP estimator being accurate on workloads which emulate a realistic training run that an AI developer might actually do. In practice, we'd expect an AI developer to want to fully utilize their GPUs as best as they can, so I define "frontier" sample training runs as ones in which the GPU is active at least 80% of the time, as reported by polling tegrastats, a profiling tool on the Jetson . Note, this is different from the EMC bandwidth utilization percentage which is measuring how often the memory controller is running. Since I was unsure which hyperparameters to use to be above the 80% threshold, I ran the script under a bunch of hyperparameter configurations and only kept the ones which utilized the GPU above the 80% threshold on average.

We can get the ground truth for the number of FLOPs by looking at the shapes of tensors involved in the training algorithm. For example, multiplying an (m×k) matrix by a (k×n) matrix requires m·n·k multiplications and roughly as many additions, so we could count this as 2·m·k·n FLOPs towards the ground truth, regardless of how the hardware executes this at the lowest level. To get the ground truth FLOPs for the various sample training runs in my experiments, I used PyTorch's FlopCounterMode library. This library hooks onto PyTorch to intercept every tensor shape as it is used in low-level operations during the training run and dynamically compute the algorithmic FLOPs as the program executes.

The estimator's algorithm is shown below in equation (1).

I explain terms below, but first, let me motivate on a high level what this equation is doing. In the simplest case, we would like to just figure out how much energy the GPU is using, calibrate our Energy_Per_TFLOP parameter based on a bunch of workloads with known total FLOP counts, and then predict on held-out