Treffer: A novel method of heterogeneous parallel machine learning by CPU–TPU for molecular dynamics.

In this paper, a heterogeneous parallel machine learning molecular dynamics (MLMD) calculation method based on both central processing unit (CPU) and SOPHON BM1684X tensor processing unit (TPU) is proposed. The method aims to offer a new hardware deployment approach for advanced MLMD algorithms, alleviating the constraints imposed by the severe "memory wall" and "power wall" bottlenecks caused by the separation of storage units and computing units inherent in von Neumann architecture-based machines at the hardware level. By decomposing complex MD simulation tasks into subtasks that can be processed in parallel on both CPU and TPU, this method enhances computational efficiency while maintaining high precision. Specifically, the potential energy surface fitting task in MD simulation is deployed on the TPU, leveraging its parallel processing capabilities to accelerate computations. Meanwhile, load balancing between the CPU and TPU is achieved by executing other computational tasks on the CPU. Experimental results demonstrate significant improvements in computational speed, energy efficiency, and the size of computable systems compared to the non-heterogeneous CPU-only system, indicating that heterogeneous parallel computing is an effective method for accelerating MD simulations. [ABSTRACT FROM AUTHOR]