The Pulse
TurboPairFormer speeds OpenFold3 training up to 1.73× in tests
Researchers report that TurboPairFormer reduced GPU computation time for OpenFold3 training compared with its original Triton backend. The open-source kernels also produced more accurate and reproducible gradients under the tested conditions.
TurboPairFormer ran OpenFold3 training up to 1.73× faster in GPU computation time per optimizer step than the model’s original Triton backend, researchers report in a preprint. The open-source kernels were built by researchers from Lambda, Stevens Institute of Technology and the OpenFold team to accelerate operations used in biomolecular structure prediction.
Lambda described the results in a post on November 13, following the preprint dated October 5. The researchers also report more accurate gradients and repeatable results in tests, although those findings are specific to the configurations they evaluated.
Where the time goes
OpenFold3 supports training and fine-tuning of models that predict structures involving proteins, nucleic acids and small molecules. TurboPairFormer implements triangle attention and triangle multiplication, operations used in those models.
Triangle attention lets a model use information about a third molecular component when reasoning about a pair. Its computational cost grows cubically with the number of tokens: doubling the training crop size makes the operation about eight times as expensive, according to the researchers.
During the backward pass, training calculates gradients for queries, keys, values and pair biases. TurboPairFormer recomputes each attention probability tile once and reuses it for all four gradients, rather than repeatedly calculating it for different gradients.
Training benchmarks
The researchers integrated TurboPairFormer into OpenFold3 alongside their triangle multiplication kernels and compared it with seven other backend configurations. Tests ran on 16 NVIDIA H100 GPUs across two nodes, using crop sizes of 384, 640 and 768 tokens.
At a crop size of 768, TurboPairFormer took 12.58 seconds of GPU computation per optimizer step, compared with 21.72 seconds for OpenFold3 Triton and 14.25 seconds for NVIDIA cuEquivariance, according to the reported results. The researchers found that it had the lowest computation time per step among the tested configurations at all three crop sizes.
The measurements covered steps 21 through 200 and excluded data loading, host-to-device transfers and EMA updates. On H200 GPUs at crop size 768, Lambda reported a 1.54× speedup over OpenFold3 Triton, with time per step falling from 22.53 to 14.67 seconds.
Accuracy and repeatability
The kernels also address rounding error when values are converted from FP32 to BF16 precision. TurboPairFormer saves an approximation of the rounding residual and uses it in the backward calculation without changing the BF16 output.
In controlled experiments, that compensation reduced mean gradient error by 28–47% for the query, key and pair-bias gradients, the preprint reports. Against an FP64 reference, the implementation had the lowest mean errors for those gradients among the eight attention implementations evaluated; value-gradient accuracy was essentially tied, according to Lambda.
The team also tested whether repeated calls produced identical gradients. On a single H100, all four gradients matched bit for bit across five runs in each of 600 fixed-input cases, under the test conditions.
A training continuation using an earlier build of the kernels kept losses finite, while the OpenFold3 Triton run produced NaN losses within 10 updates, the researchers report. TurboPairFormer is available on PyPI and currently supports H100 and H200 GPUs.