CAREER: A Framework to Power Next-Generation AI-based Scientific Models on Hardware with Massive Parallelism

Recently, the advancement of Large Language Models has facilitated the development of AI models for scientific discovery. However, many scientists today face two major challenges in building high-quality AI-based scientific models: (i) it is very challenging to fully utilize the compute power offered by modern hardware with massive parallelism because many science-informed model architectures often exhibit high-order computation and memory complexity, and programming such hardware is extremely complex, requiring advanced system technologies to effectively utilize all available hardware resources simultaneously to achieve optimal performance; (ii) the challenge of the humans in the loop -- non-CS scientists. The project's novelties are in addressing these two major challenges: empowering non-CS researchers to harness the power of modern hardware with massive parallelism for training science-informed AI4Science models, and simplifying the complex programming required to achieve optimal performance. The project’s broader significance and importance are in making training advanced AI models more accessible and effective for scientific research, thereby accelerating scientific discovery and innovation. The project includes the following synergistic components. First, it enables capturing unprecedented high-order, extremely long-range, and high-volumetric interactions in scientific data by developing novel memory-efficient and hardware-friendly kernels on a single GPU, effectively changing the state of many advanced models from impossible to possible to train with limited GPU resources. Second, it creates novel hardware and model-conscious 4D parallelism, which further unlocks the performance potential of AI4Science models in multi-node multi-GPU distributed environments. Finally, the project builds an automated pipeline through new techniques in performance adviser, deep learning compilation, and auto-tuner, which lowers the burden on non-CS scientists from dealing with complex parallel hardware. Together, this project paves the path to generalize AI system technologies so they can broadly address major system pain points and promote progress in large-scale scientific discoveries. This award reflects NSF's statutory mission and has been deemed worthy of support through evaluation using the Foundation's intellectual merit and broader impacts review criteria.