MegaFold: System-Level Optimizations for Accelerating Protein Structure Prediction Models

• URL: http://arxiv.org/abs/2506.20686v1
• Date: Tue, 24 Jun 2025 23:30:49 GMT
• Title: MegaFold: System-Level Optimizations for Accelerating Protein Structure Prediction Models
• Authors: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang,
• Abstract summary: We present MegaFold, a cross-platform system to accelerate AF3 training.<n>We show that MegaFold reduces peak memory usage of AF3 training by up to 1.23$times$ and improves per-iteration training time by up-to 1.73$times$ and 1.62$times$ respectively.
• Score: 17.994632753972958
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Protein structure prediction models such as AlphaFold3 (AF3) push the frontier of biomolecular modeling by incorporating science-informed architectural changes to the transformer architecture. However, these advances come at a steep system cost, introducing: compute- and memory-intensive operators, 2D attention mechanisms, and retrieval-augmented data pipelines, which collectively hinder the scalability of AF3 training. In this work, we present MegaFold, a cross-platform system to accelerate AF3 training. MegaFold tackles key bottlenecks through ahead-of-time caching to eliminate GPU idle time from the retrieval-augmented data pipeline, Triton-based kernels for memory-efficient EvoAttention on heterogeneous devices, and deep fusion for common and critical small operators in AF3. Evaluation on both NVIDIA H200 and AMD MI250 GPUs shows that MegaFold reduces peak memory usage of AF3 training by up to 1.23$\times$ and improves per-iteration training time by up-to 1.73$\times$ and 1.62$\times$ respectively. More importantly, MegaFold enables training on 1.35$\times$ longer sequence lengths compared to PyTorch baselines without running out-of-memory, significantly improving the scalability of modern protein folding models. We open source our code at https://github.com/Supercomputing-System-AI-Lab/MegaFold/.

Related papers

• Flexiffusion: Training-Free Segment-Wise Neural Architecture Search for Efficient Diffusion Models [50.260693393896716]
  Diffusion models (DMs) are powerful generative models capable of producing high-fidelity images but constrained by high computational costs.<n>We propose Flexiffusion, a training-free NAS framework that jointly optimize generation schedules and model architectures without modifying pre-trained parameters.<n>Our work pioneers a resource-efficient paradigm for searching high-speed DMs without sacrificing quality.
  arXiv  Detail & Related papers  (2025-06-03T06:02:50Z)
• Ultra Memory-Efficient On-FPGA Training of Transformers via Tensor-Compressed Optimization [9.555456615472512]
  This paper presents the first on-FPGA accelerator for end-to-end transformer training.<n>On the algorithm side, we present a bi-directional contraction flow for tensorized transformer training.<n>On the hardware side, we store all highly compressed model parameters and gradient information on chip.
  arXiv  Detail & Related papers  (2025-01-11T23:29:51Z)
• FAMOUS: Flexible Accelerator for the Attention Mechanism of Transformer on UltraScale+ FPGAs [0.0]
  Transformer neural networks (TNNs) are being applied across a widening range of application domains, including natural language processing (NLP), machine translation, and computer vision (CV)<n>This paper proposes textitFAMOUS, a flexible hardware accelerator for dense multi-head attention computation of TNNs on field-programmable gate arrays (FPGAs)<n>It is optimized for high utilization of processing elements and on-chip memories to improve parallelism and reduce latency.
  arXiv  Detail & Related papers  (2024-09-21T05:25:46Z)
• TrIM, Triangular Input Movement Systolic Array for Convolutional Neural Networks: Architecture and Hardware Implementation [0.0]
  TrIM is an innovative dataflow based on a triangular movement of inputs.<n> TrIM can reduce the number of memory accesses by one order of magnitude when compared to state-of-the-art systolic arrays.<n> architecture achieves a peak throughput of 453.6 Giga Operations per Second.
  arXiv  Detail & Related papers  (2024-08-05T10:18:00Z)
• fVDB: A Deep-Learning Framework for Sparse, Large-Scale, and High-Performance Spatial Intelligence [50.417261057533786]
  fVDB is a novel framework for deep learning on large-scale 3D data. Our framework is fully integrated with PyTorch enabling interoperability with existing pipelines.
  arXiv  Detail & Related papers  (2024-07-01T20:20:33Z)
• AI and Memory Wall [81.06494558184049]
  We show how memory bandwidth can become the dominant bottleneck for decoder models. We argue for a redesign in model architecture, training, and deployment strategies to overcome this memory limitation.
  arXiv  Detail & Related papers  (2024-03-21T04:31:59Z)
• Gated Linear Attention Transformers with Hardware-Efficient Training [60.670102007737476]
  This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability. We then generalize this algorithm to a more expressive variant of linear attention with data-dependent gates. When used as a replacement for the standard attention layer in Transformers, the resulting gated linear attention Transformer is found to perform competitively.
  arXiv  Detail & Related papers  (2023-12-11T18:51:59Z)
• FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores [18.016204763652553]
  Convolution models with long filters have demonstrated state-of-the-art reasoning abilities in many long-sequence tasks. Fast Fourier Transform (FFT) allows long convolutions to run in $O(N logN)$ time in sequence length $N$ but has poor hardware utilization. In this paper, we study how to optimize the FFT convolution.
  arXiv  Detail & Related papers  (2023-11-10T07:33:35Z)
• Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions [101.08706223326928]
  Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In this paper, we seek to enable $mathcal O(1)$ compute and memory cost per token in any pre-trained long convolution architecture.
  arXiv  Detail & Related papers  (2023-10-28T18:40:03Z)
• Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity [12.663030430488922]
  We propose Flash-LLM for enabling low-cost and highly-efficient large generative model inference on high-performance Cores. At SpMM kernel level, Flash-LLM significantly outperforms the state-of-the-art library, i.e., Sputnik and SparTA by an average of 2.9x and 1.5x, respectively.
  arXiv  Detail & Related papers  (2023-09-19T03:20:02Z)
• GLEAM: Greedy Learning for Large-Scale Accelerated MRI Reconstruction [50.248694764703714]
  Unrolled neural networks have recently achieved state-of-the-art accelerated MRI reconstruction. These networks unroll iterative optimization algorithms by alternating between physics-based consistency and neural-network based regularization. We propose Greedy LEarning for Accelerated MRI reconstruction, an efficient training strategy for high-dimensional imaging settings.
  arXiv  Detail & Related papers  (2022-07-18T06:01:29Z)
• On-Device Training Under 256KB Memory [62.95579393237751]
  We propose an algorithm-system co-design framework to make on-device training possible with only 256KB of memory. Our framework is the first solution to enable tiny on-device training of convolutional neural networks under 256KB and 1MB Flash.
  arXiv  Detail & Related papers  (2022-06-30T17:59:08Z)
• Siamese Transformer Pyramid Networks for Real-Time UAV Tracking [3.0969191504482243]
  We introduce the Siamese Transformer Pyramid Network (SiamTPN), which inherits the advantages from both CNN and Transformer architectures. Experiments on both aerial and prevalent tracking benchmarks achieve competitive results while operating at high speed. Our fastest variant tracker operates over 30 Hz on a single CPU-core and obtaining an AUC score of 58.1% on the LaSOT dataset.
  arXiv  Detail & Related papers  (2021-10-17T13:48:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.