GPU Transformer training: OpenNN vs PyTorch vs TensorFlow

OpenNN trains the tested encoder-decoder Transformer at 1.69 million tokens/s in fp32 and 2.67 million tokens/s in bf16 on an NVIDIA GeForce RTX 4080, leading both PyTorch and TensorFlow in the same throughput workload.

Contents

• Introduction
• Benchmark application
• Reference computer
• Methodology
• Results
• Discussion
• Conclusions
• Reproducing
• References

Introduction

Training adds backward propagation and the Adam parameter update to the Transformer forward pass. A fair throughput comparison must keep the corpus, sequence length, model shape, batch, optimizer work, warmup, and timing boundary aligned while allowing each framework to use its normal optimized GPU path.

This benchmark compares OpenNN, PyTorch, and TensorFlow in both fp32 and bf16. OpenNN leads both precision cells, with its largest relative advantage in fp32.

Benchmark application

Reference computer

Methodology

The three engines read the same corpus and train the same shape. One warmup epoch per engine is excluded from timing.

• OpenNN uses device-resident training, fused attention, and its CUDA graph path.
• PyTorch uses its optimized Transformer path, bf16 autocast where applicable, and fused Adam.
• TensorFlow uses compiled graph execution with XLA and mixed bf16 where applicable.
• Tokens per second are derived from the common sequence dimensions and measured samples per second.
• The saved artifact contains one successful run per engine and precision. Its median field therefore represents that one run.

The benchmark is MLPerf-inspired but is not an official MLPerf submission. Final loss is not gated to a common target in this throughput run; convergence and energy-to-quality belong in separate benchmarks.

Results

The corresponding samples-per-second values are:

Discussion

OpenNN’s fp32 result is the strongest relative win: 76.4% more throughput than PyTorch and 90.8% more than TensorFlow. In bf16 all three frameworks use tensor-core-friendly paths and the field becomes closer, but OpenNN still leads by 16.0% over PyTorch and 29.0% over TensorFlow.

Within OpenNN, bf16 raises throughput from 1.69 million to 2.67 million tokens/s for this exact training configuration, a 1.58x increase. That number is specific to this model and should not be generalized to every Transformer shape.

Conclusions

• OpenNN leads both tested training precisions.
• In fp32, OpenNN reaches 1.76x PyTorch and 1.91x TensorFlow throughput.
• In bf16, OpenNN reaches 1.16x PyTorch and 1.29x TensorFlow throughput.
• OpenNN reaches a peak of 2.67 million tokens/s in bf16.
• These are throughput results. They do not claim that the frameworks reach a fixed loss in the same number of epochs.

Reproducing

The canonical runner is docs/benchmarks/throughput/attention-speed/run_transformer_train.py:

python run_transformer_train.py --batch 32 --epochs 9 --runs 1 --precision both

The result artifact is docs/benchmarks/results/gpu-transformer-training-speed-20260710T084504Z.json.

References

• Attention Is All You Need.
• OpenNN source repository.
• PyTorch Transformer documentation.
• TensorFlow mixed precision guide.