VLDB 2026 Research / reviewers in the wild / expert
Weile Jia
dblp:41/10441
· DBLP profile ↗
30ranked-venue papers
2as first author
26since 2021 · last 2025
—ORCID · conflict
Domains — the database's venue-derived domains; a paper can count in several
Systems, architecture and hardware · 23 · 2 first-author · 20 since 2021Artificial intelligence and machine learning · 5 · 5 since 2021Software engineering, systems software and programming languages · 1 · 1 since 2021Graphics, computer vision, multimedia, augmented reality and games · 1 · 1 since 2021Applied, interdisciplinary, general and emerging computing · 1 · 1 since 2021
| Year | Publication | Venue | Position |
|---|---|---|---|
| 2025 | FastSpMM: Leveraging Tensor Cores for Sparse Matrix Multiplication
Mingzhen Li 0001, Weile Jia, Hailong Yang 0002, Guangming Tan |
CF | 3 |
| 2025 | ELoRA: Low-Rank Adaptation for Equivariant GNNsabstractPre-trained interatomic potentials have become a new paradigm for atomistic materials simulations, enabling accurate and efficient predictions across diverse chemical systems. Despite their promise, fine-tuning is often required for complex tasks to achieve high accuracy. Traditional parameter-efficient fine-tuning approaches are effective in NLP and CV. However, when applied to SO(3) equivariant pre-trained interatomic potentials, these methods will inevitably break equivariance—a critical property for preserving physical symmetries. In this paper, we introduce ELoRA (Equivariant Low-Rank Adaptation), a novel fine-tuning method designed specifically for SO(3) equivariant Graph Neural Networks (GNNs), the backbones in multiple pre-trained interatomic potentials. ELoRA adopts a path-dependent decomposition for weights updating which offers two key advantages: (1) it preserves SO(3) equivariance throughout the fine-tuning process, ensuring physically consistent predictions, and (2) it leverages low-rank adaptations to significantly improve data efficiency. We prove that ELoRA maintains equivariance and demonstrate its effectiveness through comprehensive experiments. On the rMD17 organic dataset, ELoRA achieves a 25.5% improvement in energy prediction accuracy and a 23.7% improvement in force prediction accuracy compared to full-parameter fine-tuning. Similarly, across 10 inorganic datasets, ELoRA achieves average improvements of 12.3% and 14.4% in energy and force predictions, respectively. Code will be made publicly available at https://github.com/hyjwpk/ELoRA. Chen Wang 0154, Siyu Hu, Guangming Tan, Weile Jia |
ICML | 4 |
| 2025 | Efficient Long Context Fine-tuning with Chunk FlowabstractLong context fine-tuning of large language models(LLMs) involves training on datasets that are predominantly composed of short sequences and a small proportion of longer sequences. However, existing approaches overlook this long-tail distribution and employ training strategies designed specifically for long sequences. Moreover, these approaches also fail to address the challenges posed by variable sequence lengths during distributed training, such as load imbalance in data parallelism and severe pipeline bubbles in pipeline parallelism. These issues lead to suboptimal training performance and poor GPU resource utilization. To tackle these problems, we propose a chunk-centric training method named ChunkFlow. ChunkFlow reorganizes input sequences into uniformly sized chunks by consolidating short sequences and splitting longer ones. This approach achieves optimal computational efficiency and balance among training inputs. Additionally, ChunkFlow incorporates a state-aware chunk scheduling mechanism to ensure that the peak memory usage during training is primarily determined by the chunk size rather than the maximum sequence length in the dataset. Integrating this scheduling mechanism with existing pipeline scheduling algorithms further enhances the performance of distributed training. Experimental results demonstrate that, compared with Megatron-LM, ChunkFlow can be up to 4.53x faster in the long context fine-tuning of LLMs. Furthermore, we believe that ChunkFlow serves as an effective solution for a broader range of scenarios, such as long context continual pre-training, where datasets contain variable-length sequences. Xiulong Yuan, Hongtao Xu, Wenting Shen, Ang Wang, Xiafei Qiu, Jie Zhang 0135, Yuqiong Liu, Bowen Yu 0002, Junyang Lin, Mingzhen Li 0001, Weile Jia, Yong Li 0045, Wei Lin 0016 |
ICML | 11 |
| 2025 | Cherry: Breaking the GPU Memory Wall for Large-Scale GNN Training via Micro-BatchingabstractGraph Neural Networks (GNNs) have shown remarkable performance across a variety of graph-related tasks.Recent efforts indicate that GNN performance can be enhanced through more sophisticated strategies, such as employing advanced aggregators, increasing aggregation depth, and utilizing larger sampling rates, etc.While these strategies yield promising results, it also incurs a significantly larger memory footprint that can easily surpass the GPU memory capacity.Micro-batching has emerged as a promising method to mitigate GPU memory bottleneck while preserving model accuracy.Nevertheless, integrating micro-batches into GNN Yan Wang 0022, Haoran Kong, Hao Chen 0002, Weile Jia, Dingwen Tao, Xin He 0054 |
ICS | 7 |
| 2025 | Large Scale Finite-Temperature Real-Time Time Dependent Density Functional Theory Calculation with Hybrid Functional on ARM and GPU SystemsabstractUltra-fast electronic phenomena originating from finite temperature, such as nonlinear optical excitation, can be simulated with high fidelity via real-time time dependent density functional theory (rt-TDDFT) calculations with hybrid functional. However, previous rt-TDDFT simulations of real materials using the optimal gauge-known as the parallel transport gauge-have been limited to low-temperature systems with band gaps. In this paper, we introduce the parallel transport-implicit midpoint (PT-IM) method, which significantly accelerates finitetemperature rt-TDDFT calculations of real materials with hybrid function. We first implement PT-IM with hybrid functional in our plane wave code PWDFT, and optimized it on both GPU and ARM platforms to build a solid baseline code. Next, we propose a diagonalization method to reduce computation and communication complexity, and then, we employ adaptively compressed exchange (ACE) method to reduce the frequency of the most expensive Fock exchange operator. Finally, we adopt the ring_based method and the shared memory mechanism to overlap computation and communication and alleviate memory consumption respectively. Numerical results show that our optimized code can reach 3072 atoms for rt-TDDFT simulation with hybrid functional at finite temperature on 192 computing nodes, the time-to-solution for one time step is 429.3 s, which is 41.4 times faster compared to the baseline. Zhuoqiang Guo, Qiuchen Sha, Haibo Li 0007, Wei Hu 0006, Guangming Tan, Weile Jia |
IPDPS | 9 |
| 2025 | FastCHGNet: Training One Universal Interatomic Potential to 1.5 Hours with 32 GPUsabstractGraph neural network universal interatomic potentials (GNN-UIPs) have demonstrated remarkable generalization and transfer capabilities in material discovery and property prediction. These models can accelerate molecular dynamics (MD) simulation by several orders of magnitude while maintaining$a b$initio accuracy, making them a promising new paradigm in material simulations. One notable example is Crystal Hamiltonian Graph Neural Network (CHGNet), pretrained on the energies, forces, stresses, and magnetic moments from the MPtrj dataset, representing a state-of-the-art GNN-UIP model for charge-informed MD simulations. However, training the CHGNet model is time-consuming (8.3 days on one A100 GPU) for three reasons: (i) requiring multi-layer propagation to reach more distant atom information, (ii) requiring secondorder derivatives calculation to finish weights updating and (iii) the implementation of reference CHGNet does not fully leverage the computational capabilities. This paper introduces FastCHGNet, an optimized CHGNet, with three contributions: Firstly, we design innovative Force/Stress Readout modules to decompose Force/Stress prediction. Secondly, we adopt massive optimizations such as kernel fusion, redundancy bypass, etc, to exploit GPU computation power sufficiently. Finally, we extend CHGNet to support multiple GPUs and propose a load-balancing technique to enhance GPU utilization. Numerical results show that FastCHGNet reduces memory footprint by a factor of 3.59. The final training time of FastCHGNet can be decreased to 1.53 hours on 32 GPUs without sacrificing model accuracy. Yuanchang Zhou, Siyu Hu, Chen Wang 0154, Lin-Wang Wang, Guangming Tan, Weile Jia |
IPDPS | 6 |
| 2025 | MD-pipe: A Strong Scaling Enhanced Pipeline Architecture for Ab Initio Accuracy Molecular DynamicsabstractMolecular Dynamics (MD) simulations with first-principles accuracy are widely applied in various fields, including materials science and molecular pharmacology.Current research focus on reducing the solution time of ab initio molecular dynamics (AIMD) from both Ning Kang 0007, Guojun Yuan, Beining Zhang, Guanglei Chen, Jiayi Rao, Zhan Wang 0003, Weile Jia, Ninghui Sun, Guangming Tan |
ISCA | 11 |
| 2025 | Skrull: Towards Efficient Long Context Fine-tuning through Dynamic Data SchedulingabstractLong-context supervised fine-tuning (Long-SFT) plays a vital role in enhancing the performance of large language models (LLMs) on long-context tasks. To smoothly adapt LLMs to long-context scenarios, this process typically entails training on mixed datasets containing both long and short sequences. However, this heterogeneous sequence length distribution poses significant challenges for existing training systems, as they fail to simultaneously achieve high training efficiency for both long and short sequences, resulting in sub-optimal end-to-end system performance in Long-SFT.
In this paper, we present a novel perspective on data scheduling to address the challenges posed by the heterogeneous data distributions in Long-SFT. We propose Skrull, a dynamic data scheduler specifically designed for efficient long-SFT. Through dynamic data scheduling, Skrull balances the computation requirements of long and short sequences, improving overall training efficiency. Furthermore, we formulate the scheduling process as a joint optimization problem and thoroughly analyze the trade-offs involved. Based on those analysis, Skrull employs a lightweight scheduling algorithm to achieve near-zero cost online scheduling in Long-SFT. Finally, we implement Skrull upon DeepSpeed, a state-of-the-art distributed training system for LLMs. Experimental results demonstrate that Skrull outperforms DeepSpeed by 3.76x on average (up to 7.54x) in real-world long-SFT scenarios. Hongtao Xu, Wenting Shen, Yuanxin Wei, Ang Wang, Guo Runfan, Tianxing Wang 0006, Yong Li 0045, Mingzhen Li 0001, Weile Jia |
NeurIPS | 9 |
| 2025 | Exploring Landscapes for Better Minima along ValleysabstractFinding lower and better-generalizing minima is crucial for deep learning. However, most existing optimizers stop searching the parameter space once they reach a local minimum. Given the complex geometric properties of the loss landscape, it is difficult to guarantee that such a point is the lowest or provides the best generalization. To address this, we propose an adaptor "E" for gradient-based optimizers. The adapted optimizer tends to continue exploring along landscape valleys (areas with low and nearly identical losses) in order to search for potentially better local minima even after reaching a local minimum. This approach increases the likelihood of finding a lower and flatter local minimum, which is often associated with better generalization. We also provide a proof of convergence for the adapted optimizers in both convex and non-convex scenarios for completeness. Finally, we demonstrate their effectiveness in an important but notoriously difficult training scenario, large-minibatch training, where Lamb is the benchmark optimizer. Our testing results show that the adapted Lamb, ALTO, increases the test accuracy (generalization) of the current state-of-the-art optimizer by an average of 2.5\% across a variety of large-batch training tasks. This work potentially opens a new research direction in the design of optimization algorithms. Yuanchang Zhou, Guangming Tan, Weile Jia |
NeurIPS | 5 |
| 2025 | Mario: Near Zero-cost Activation Checkpointing in Pipeline ParallelismabstractLarge language models have to be trained in parallel due to their large number of parameters and significant memory footprint. Among various parallelism techniques, pipeline parallelism is widely adopted in inter-node scenarios with minimal communication overhead. However, state-of-the-art pipeline schemes lead to extra and imbalanced memory footprints, leaving room for further improvement. In this paper, we propose Mario, a pipeline optimizer that automatically tessellates activation checkpointing to existing pipeline schemes, enabling training larger models (or longer sequences) with less and balanced memory footprint across GPUs and improved GPU utilization. First, the activation recomputation can be effectively overlapped in the bubbles by moving it earlier in the execution process, thereby improving overall efficiency. With eliminated memory footprint through checkpointing, Mario allows for preposing more forward computation into the pipeline bubbles, making more room for further overlapping with greater flexibility, and thus exploiting the bubbles. Then we design a lightweight pipeline simulator to model execution behavior w/o|w/ Mario. Finally, we introduce an automatic pipeline scheduler specifically for Mario, capable of searching for near optimal combination of checkpointing and pipeline configurations within minutes. Experimental results on GPT3 and LLaMA2 models show that Mario can speed up existing state-of-the-art pipeline schemes (w/o|w/ checkpointing) including 1F1B, Chimera, and Interleave by 1.16×|1.57× on average. This work paves a new direction for effective low-cost pipeline training. Mingzhen Li 0001, Guangming Tan, Weile Jia |
PPoPP | 4 |
| 2025 | Million-Atom Ab Initio Electron Dynamics: Discontinuous Galerkin Real-Time Time-Dependent Density Functional TheoryabstractOver the past decades, first-principles real-time time dependent density functional theory(rt-TDDFT) simulations have been limited to systems with only thousands of atoms. We propose a novel method based on the discontinuous Galerkin adaptive local basis, significantly reducing global communication in rt-TDDFT. We further introduce a tensor compression technique that leverages basis locality to avoid repeated evaluation of multi-center integrals in hybrid functionals, greatly reducing computational cost. To overcome the projection bottleneck in our basis sets, we design a fused Gemm-Reduce operation that achieves several times higher floating-point efficiency than standard BLAS combination. Our implementation reaches 34.8% of theoretical peak performance on 524,288 CGs of the New Sunway supercomputer and simulates electronic dynamics of systems with over one million atoms for both local-semi-local and hybrid functionals. This work improves computational scale by two orders of magnitude, opening new possibilities for exploring ultrafast dynamics in large-scale materials and nanophotonic devices. Junwei Feng, Junshi Chen 0003, Xinming Qin, Lingyun Wan, Wentiao Wu, Bingkun Hou, Yexuan Lin, Zechuan Zhang, Weile Jia, Hong An, Jinlong Yang 0003, Wei Hu 0006 |
SC | 14 |
| 2025 | Deep Learning-Enabled Supercritical Flame Simulation at Detailed Chemistry and Real-Fluid Accuracy Towards Trillion-Cell ScaleabstractFor decades, supercritical flame simulations incorporating detailed chemistry and real-fluid transport have been limited to millions of cells, constraining the resolved spatial and temporal scales of the physical system. We optimize the supercritical flame simulation software DeepFlame—which incorporates deep neural networks while retaining the real-fluid mechanical and chemical accuracy—from three perspectives: parallel computing, computational efficiency, and I/O performance. Our highly optimized DeepFlame achieves supercritical liquid oxygen/methane (LOX/\(\ce {CH4}\)) turbulent combustion simulation of up to 618 and 154 billion cells with unprecedented time-to-solution, attaining 439/1186 and 187/316 PFlop/s (32.3%/21.8% and 37.4%/31.8% of the peak) in FP32/mixed-FP16 precision on Sunway (98,304 nodes) and Fugaku (73,728 nodes) supercomputers, respectively. This computational capability surpasses existing capacities by three orders of magnitude, enabling the first practical simulation of rocket engine combustion with >100 LOX/\(\ce {CH4}\) injectors. This breakthrough establishes high-fidelity supercritical flame modeling as a critical design tool for next-generation rocket propulsion and ultra-high energy density systems. Zhuoqiang Guo, Runze Mao, Guangming Tan, Weile Jia, Zhi X. Chen |
SC | 5 |
| 2025 | Caracal: A GPU-Resident Sparse LU Solver with Lightweight Fine-Grained SchedulingabstractWe address inefficiencies in task scheduling, memory management, and scalability in GPU-resident sparse LU factorization with a two-level approach of sequentially scheduled coarse-grained blocks containing multiple fine-grained blocks managed with a lightweight static scheduler enabling multi-stream parallelism. Additionally, we design an intelligent memory caching mechanism for the fine-grained scheduler, which retains frequently accessed data in GPU memory. To further enhance scalability, we introduce a distributed memory design that partitions the input matrix using a 1D block-cyclic distribution and optimizes inter-GPU communication via NVLink. The multi-GPU design reaches a computational throughput of 6.46 TFLOP/s on four A100 GPUs, demonstrating promising scalability. This is up to 7x speedup over the latest SuperLU_DIST with 3D communication, 94x speedup over PanguLU, 16x speedup over PasTiX, and 10x speedup over our own coarse-grained dynamic scheduling implementation while reaching up to 21% of the A100’s theoretical peak performance. Tingxuan Zhong, Yuxi Hong 0001, Guofeng Feng, Weile Jia, Hatem Ltaief, David E. Keyes |
SC | 6 |
| 2025 | An interpretable DeePMD-kit performance model for emerging supercomputersabstractAbstract Deep potential (DP) scheme has increased the simulation temporal and spatial scales while maintaining the ab initio accuracy of the molecular dynamics. DeePMD-kit is an outstanding application that implements DP scheme efficiently. However, current performance model cannot accurately measure the resource utilization of DeePMD-kit operators and predict the execution time. We introduce DP-perf, an interpretable performance model for DeePMD-kit. DP-perf can accurately measure the resource utilization of the individual DeePMD-kit operators, communication pattern, and the overall application by exploiting physical system properties and machine configurations. It can be easily applied to mainstream supercomputers including Tianhe-3F, the new Sunway, Fugaku, and Summit. With DP-perf, users can select the optimal machine and decide the corresponding configuration for various purposes (e.g., lower cost, less time) without real runs. Evaluation of four top supercomputers shows that DP-perf can fit overall execution time with a low mean absolute percentage error of 5.7 %/8.1%/14.3%/13.1% on Tianhe-3F/new Sunway/Fugaku/Summit. On the prediction scenario, DP-perf can predict the total execution time with a mean absolute percentage error of less than 20%. Xiangyu Meng 0005, Xun Wang 0010, Mingzhen Li 0001, Guangming Tan, Weile Jia |
CCF Trans. High Perform. Comput. | 5 |
| 2025 | 29-Billion Atoms Molecular Dynamics Simulation With Ab Initio Accuracy on 35 Million Cores of New Sunway SupercomputerabstractPhysical phenomena such as bond breaking and phase transitions require molecular dynamics (MD) withab initioaccuracy, involving up to billions of atoms and over nanosecond timescales. Previous state-of-the-art work has demonstrated that neural network molecular dynamics (NNMD) like deep potential molecular dynamics (DeePMD), can successfully extend the temporal and spatial scales of MD withab initioaccuracy on both ARM and GPU platforms. However, the DeePMD-kit package is currently unable to fully exploit the computational potential of the new Sunway supercomputer due to its unique many-core architecture, memory hierarchy, and low precision capability. In this paper, we re-design the DeePMD-kit to harness the massive computing power of the new Sunway, enabling the MD with over ten billion atoms. We first design a large-scale parallelization scheme to exploit the massive parallelism of the new Sunway. Then we devise specialized optimizations for the time-consuming operators. Finally, we design a novel mixed precision method for DeePMD-kit customized operators to leverage the low precision computing power of the new Sunway. The optimized DeePMD-kit achieves 67.6 / 56.5$\boldsymbol{\times}$speedup for water / copper systems on the new Sunway. Meanwhile, it can perform 29 billion atoms simulation for the water system on 35 million cores (i.e., 90,000 computing nodes, around 84% of the whole supercomputer) with a peak performance of 57.1 PFLOPs, which is 7.9$\boldsymbol{\times}$bigger and 1.2$\boldsymbol{\times}$faster than state-of-the-art results. This paves the way for investigating more realistic scenarios, such as studying the mechanical properties of metals, semiconductor devices, batteries, and other materials and physical systems. Xun Wang 0010, Xiangyu Meng 0005, Zhuoqiang Guo, Mingzhen Li 0001, Mingfan Li, Ninghui Sun, Guangming Tan, Weile Jia |
IEEE Trans. Computers | 11 |
| 2024 | Accelerating Large-Scale Sparse LU Factorization for RF Circuit Simulation
Guofeng Feng, Zhuoqiang Guo, Mingzhen Li 0001, Zhou Jin 0001, Weile Jia, Guangming Tan, Ninghui Sun |
Euro-Par (3) | 7 |
| 2024 | POSTER: Optimizing Sparse Tensor Contraction with Revisiting Hash Table DesignabstractSparse tensor contraction (SpTC) serves as an essential operation in high-performance applications. The high dimensionality of sparse tensors makes SpTC fundamentally challenging in aspects such as costly multidimensional index search, extensive intermediate output data, and indirect addressing. Previous state-of-the-art work addresses some of these challenges through hash-table implementation. In this paper, we propose a hash-table based and fully optimized SpTC by providing a more carefully designed customized hash table design, proposing an architecture-aware algorithm for hash table selection with size prediction, applying cross-stage optimizations to exploit shared information and avoid redundant operations. Evaluating on a set of tensors extracted from the real world, our method can achieve superior speedup and reduce the memory footprint substantially compared to the current state-of-the-art work. Guofeng Feng, Weile Jia, Ninghui Sun, Guangming Tan, Jiajia Li 0001 |
PPoPP | 2 |
| 2024 | Training one DeePMD Model in Minutes: a Step towards Online LearningabstractNeural Network Molecular Dynamics (NNMD) has become a major approach in material simulations, which can speedup the molecular dynamics (MD) simulation for thousands of times, while maintaining ab initio accuracy, thus has a potential to fundamentally change the paradigm of material simulations. However, there are two time-consuming bottlenecks of the NNMD developments. One is the data access of ab initio calculation results. The other, which is the focus of the current work, is reducing the training time of NNMD model. The training of NNMD model is different from most other neural network training because the atomic force (which is related to the gradient of the network) is an important physical property to be fit. Tests show the traditional stochastic gradient methods, like the Adam algorithms, cannot efficiently deploy the multisample minibatch algorithm. As a result, a typical training (taking the Deep Potential Molecular Dynamics (DeePMD) as an example) can take many hours. In this work, we designed a heuristic minibatch quasi-Newtonian optimizer based on Extended Kalman Filter method. An early reduction of gradient and error is adopted to reduce memory footprint and communication. The memory footprint, communication and settings of hyper-parameters of this new method are analyzed in detail. Computational innovations such as customized kernels of the symmetry-preserving descriptor are applied to exploit the computing power of the heterogeneous architecture. Experiments are performed on 8 different datasets representing different real case situations, and numerical results show that our new method has an average speedup of 32.2 compared to the Reorganized Layer-wised Extended Kalman Filter with 1 GPU, reducing the absolute training time of one DeePMD model from hours to several minutes, making it one step toward online training. Siyu Hu, Qiuchen Sha, Enji Li, Xiangyu Meng 0005, Lin-Wang Wang, Guangming Tan, Weile Jia |
PPoPP | 9 |
| 2024 | Scaling Molecular Dynamics with ab initio Accuracy to 149 Nanoseconds per DayabstractPhysical phenomena such as chemical reactions, bond breaking, and phase transition require molecular dynamics (MD) simulation with ab initio accuracy ranging from milliseconds to microseconds. However, previous state-of-the-art neural network based MD packages such as DeePMD-kit can only reach 4.7 nanoseconds per day on the Fugaku supercomputer. In this paper, we present a novel node-based parallelization scheme to reduce communication by 81%, then optimize the computationally intensive kernels with sve-gemm and mixed precision. Finally, we implement intra-node load balance to further improve the scalability. Numerical results on the Fugaku supercomputer show that our work has significantly improved the time-to-solution of the DeePMD-kit by a factor of 31.7 x, reaching 149 nanoseconds per day on 12,000 computing nodes. This work has opened the door for millisecond simulation with ab initio accuracy within one week for the first time. Zhuoqiang Guo, Mingzhen Li 0001, Enji Li, Guojun Yuan, Zhan Wang 0003, Guangming Tan, Weile Jia |
SC | 10 |
| 2024 | Mille-feuille: A Tile-Grained Mixed Precision Single-Kernel Conjugate Gradient Solver on GPUsabstractConjugate gradient (CG) and biconjugate gradient stabilized (BiCGSTAB) are effective methods used for solving sparse linear systems. We in this paper propose Mille-feuille, a new solver for accelerating CG and BiCGSTAB on GPUs. We first analyze the two methods and list three findings related to the use of mixed precision, the reduction of kernel synchronization costs, and the awareness of partial convergence during the iteration steps. Then, (1) to enable tile-grained mixed precision, we develop a tiled sparse format; (2) to reduce synchronization costs, we leverage atomic operations that make the whole solving procedure work within a single GPU kernel; (3) to support a partial convergence-aware mixed precision strategy, we enable tile-wise on-chip dynamic precision conversion within the single kernel at runtime. The experimental results on an NVIDIA A100 and an AMD MI210 show that the Mille-feuille solver outperforms baseline implementations using the vendor-support cuSPARSE/hipSPARSE as well as two state-of-the-art libraries PETSc and Ginkgo by a factor of on average 3.03x/2.68x, 5.37 x, 4.36x (up to $8.77 \mathrm{x} / 7.14 x$, 16.54x, 15.69x) in CG, on average 2.65x/2.32x, 3.57x, 3.78x (up to 7.51x/6.63x, 16.64x, 11.73x) in BiCGSTAB, on average 3.82x/3.47x (up to 40.38x/47.75x) in preconditioned CG (PCG), on average 1.79x/1.63x (up to 45.63x/44.34x) in preconditioned BiCGSTAB (PBiCGSTAB), respectively. Dechuang Yang, Yiduo Niu, Weile Jia, En Shao, Weifeng Liu 0002, Guangming Tan, Zhou Jin 0001 |
SC | 4 |
| 2024 | 10-Million Atoms Simulation of First-Principle Package LS3DF
Yujin Yan, Haibo Li 0007, Lin-Wang Wang, Guangming Tan, Weile Jia, Ninghui Sun |
J. Comput. Sci. Technol. | 8 |
| 2023 | RLEKF: An Optimizer for Deep Potential with Ab Initio AccuracyabstractIt is imperative to accelerate the training of neural network force field such as Deep Potential, which usually requires thousands of images based on first-principles calculation and a couple of days to generate an accurate potential energy surface. To this end, we propose a novel optimizer named reorganized layer extended Kalman filtering (RLEKF), an optimized version of global extended Kalman filtering (GEKF) with a strategy of splitting big and gathering small layers to overcome the O(N^2) computational cost of GEKF. This strategy provides an approximation of the dense weights error covariance matrix with a sparse diagonal block matrix for GEKF. We implement both RLEKF and the baseline Adam in our alphaDynamics package and numerical experiments are performed on 13 unbiased datasets. Overall, RLEKF converges faster with slightly better accuracy. For example, a test on a typical system, bulk copper, shows that RLEKF converges faster by both the number of training epochs (x11.67) and wall-clock time (x1.19). Besides, we theoretically prove that the updates of weights converge and thus are against the gradient exploding problem. Experimental results verify that RLEKF is not sensitive to the initialization of weights. The RLEKF sheds light on other AI-for-science applications where training a large neural network (with tons of thousands parameters) is a bottleneck. Siyu Hu, Qiuchen Sha, Lin-Wang Wang, Weile Jia, Guangming Tan |
AAAI | 6 |
| 2023 | Enhance the Strong Scaling of LAMMPS on FugakuabstractPhysical phenomenon such as protein folding requires simulation up to microseconds of physical time, which directly corresponds to the strong scaling of molecular dynamics(MD) on modern supercomputers. In this paper, we present a highly scalable implementation of the state-of-the-art MD code LAMMPS on Fugaku by exploiting the 6D mesh/torus topology of the TofuD network. Based on our detailed analysis of the MD communication pattern, we first adapt coarse-grained peer-to-peer ghost-region communication with uTofu interface, then further improve the scalability via fine-grained thread pool. Finally, Remote direct memory access (RDMA) primitives are utilized to avoid buffer overhead. Numerical results show that our optimized code can reduce 77% of the communication time, improving the performance of baseline LAMMPS by a factor of 2.9x and 2.2x for Lennard-Jones and embedded-atom method potentials when scaling to 36, 846 computing nodes. Our optimization techniques can also benefit other applications with stencil or domain decomposition methods. Zhuoqiang Guo, Shunchen Shi, Guangming Tan, Weile Jia, Guojun Yuan, Zhan Wang 0003 |
SC | 7 |
| 2022 | Extending the limit of molecular dynamics with ab initio accuracy to 10 billion atomsabstractHigh-performance computing, together with a neural network model trained from data generated with first-principles methods, has greatly boosted applications of ab initio molecular dynamics in terms of spatial and temporal scales on modern supercomputers. Previous state-of-the-art can achieve 1 -- 2 nanoseconds molecular dynamics simulation per day for 100-million atoms on the entire Summit supercomputer. In this paper, we have significantly reduced the memory footprint and computational time by a comprehensive approach with both algorithmic and system innovations. The neural network model is compressed by model tabulation, kernel fusion, and redundancy removal. Then optimizations such as acceleration of customized kernel, tabulation of activation function, MPI+OpenMP parallelization are implemented on GPU and ARM architectures. Testing results of the copper system show that the optimized code can scale up to the entire machine of both Fugaku and Summit, and the corresponding system size can be extended by a factor of 134 to an unprecedented 17 billion atoms. The strong scaling of a 13.5-million atom copper system shows that the time-to-solution can be 7 times faster, reaching 11.2 nanoseconds per day. This work opens the door for unprecedentedly large-scale molecular dynamics simulations based on ab initio accuracy and can be potentially utilized in studying more realistic applications such as mechanical properties of metals, semiconductor devices, batteries, etc. The optimization techniques detailed in this paper also provide insight for relevant high-performance computing applications. Zhuoqiang Guo, Denghui Lu, Yujin Yan, Siyu Hu, Guangming Tan, Ninghui Sun, Wanrun Jiang, Linfeng Zhang 0002, Mohan Chen 0002, Han Wang 0006, Weile Jia |
PPoPP | 14 |
| 2022 | 2.5 Million-Atom Ab Initio Electronic-Structure Simulation of Complex Metallic Heterostructures with DGDFTabstractOver the past three decades, ab initio electronic structure calculations of large, complex and metallic systems are limited to tens of thousands of atoms in computational accuracy and efficiency on leadership supercomputers. We present a massively parallel discontinuous Galerkin density functional theory (DGDFT) implementation, which adopts adaptive local basis functions to discretize the Kohn-Sham equation, resulting in a block-sparse Hamiltonian matrix. A highly efficient pole expansion and selected inversion (PEXSI) sparse direct solver is implemented in DGDFT to achieve O(N1.5) scaling for quasi two-dimensional systems. DGDFT allows us to compute the electronic structures of complex metallic heterostructures with 2.5 million atoms (17.2 million electrons) using 35.9 million cores on the new Sunway supercomputer. The peak performance of PEXSI can achieve 64 PFLOPS (~5% of theoretical peak), which is un-precedented for sparse direct solvers. This accomplishment paves the way for quantum mechanical simulations into mesoscopic scale for designing next-generation electronic devices. Wei Hu 0006, Hong An, Zhuoqiang Guo, Qingcai Jiang, Xinming Qin, Junshi Chen 0003, Weile Jia, Chao Yang 0001, Zhaolong Luo, Jielan Li, Wentiao Wu, Guangming Tan, Dongning Jia, Qinglin Lu, Yeqi Huang, Liyi Wang, Jinlong Yang 0003 |
SC | 7 |
| 2021 | Editorial for the special issue on large-scale AI in classical HPC environment and AI for science
Wei Xue 0003, Haohuan Fu, Weile Jia, Guangming Tan |
CCF Trans. High Perform. Comput. | 3 |
| 2020 | Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learningabstractFor 35 years, ab initio molecular dynamics (AIMD) has been the method of choice for modeling complex atomistic phenomena from first principles. However, most AIMD applications are limited by computational cost to systems with thousands of atoms at most. We report that a machine learning based simulation protocol (Deep Potential Molecular Dynamics), while retaining ab initio accuracy, can simulate more than 1 nanosecond-long trajectory of over 100 million atoms per day, using a highly optimized code (GPU DeePMD-kit) on the Summit supercomputer. Our code can efficiently scale up to the entire Summit supercomputer, attaining 91 PFLOPS in double precision (45.5% of the peak) and 162/275 PFLOPS in mixed-single/half precision. The great accomplishment of this work is that it opens the door to simulating unprecedented size and time scales with ab initio accuracy. It also poses new challenges to the next-generation supercomputer for a better integration of machine learning and physical modeling. Weile Jia, Han Wang 0006, Mohan Chen 0002, Denghui Lu, Lin Lin 0001, Roberto Car, Weinan E, Linfeng Zhang 0002 |
SC | 1 |
| 2019 | Parallel transport time-dependent density functional theory calculations with hybrid functional on summitabstractReal-time time-dependent density functional theory (rt-TDDFT) with hybrid exchange-correlation functional has wide-ranging applications in chemistry and material science simulations. However, it can be thousands of times more expensive than a conventional ground state DFT simulation, and hence is limited to small systems. In this paper, we accelerate hybrid functional rt-TDDFT calculations using the parallel transport gauge formalism, and the GPU implementation on Summit. Our implementation can efficiently scale to 786 GPUs for a large system with 1536 silicon atoms, and the wall clock time is only 1.5 hours per femtosecond. This unprecedented speed enables the simulation of large systems with more than 1000 atoms using rt-TDDFT and hybrid functional. Weile Jia, Lin-Wang Wang, Lin Lin 0001 |
SC | 1 |
| 2018 | A Left-Looking Selected Inversion Algorithm and Task Parallelism on Shared Memory SystemsabstractGiven a sparse matrix A, the selected inversion algorithm is an efficient method for computing certain selected elements of A-1. These selected elements correspond to all or some nonzero elements of the LU factors of A. In many ways, the types of matrix updates performed in the selected inversion algorithm are similar to those performed in the LU factorization, although the sequence of operations is different. Mathias Jacquelin, Lin Lin 0001, Weile Jia, Yonghua Zhao, Chao Yang 0001 |
HPC Asia | 3 |
| 2011 | Large scale plane wave pseudopotential density functional theory calculations on GPU clustersabstractIn this work, we present our implementation of the density functional theory (DFT) plane wave pseudopotential (PWP) calculations on GPU clusters. This GPU version is developed based on a CPU DFT-PWP code: PEtot, which can calculate up to a thousand atoms on thousands of processors. Our test indicates that the GPU version can have a ~10 times speed-up over the CPU version. A detail analysis of the speed-up and the scaling on the number of CPU/GPU computing units (up to 256) are presented. The success of our speed-up relies on the adoption a hybrid reciprocal-space and band-index parallelization scheme. As far as we know, this is the first GPU DFT-PWP code scalable to large number of CPU/GPU computing units. We also outlined the future work, and what is needed to further increase the computational speed by another factor of 10. Long Wang 0008, Weile Jia, Weiguo Gao, Xuebin Chi, Lin-Wang Wang |
SC | 3 |