{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/high-performance-training-and-inference-for","title":"High-performance training and inference for deep equivariant interatomic potentials","arxiv_id":"2504.16068","date":"2025-04-22","proceeding":null,"authors":["Chuin Wei Tan","Marc L. Descoteaux","Mit Kotak","Gabriel de Miranda Nascimento","Seán R. Kavanagh","Laura Zichi","Menghang Wang","Aadit Saluja","Yizhong R. Hu","Tess Smidt","Anders Johansson","William C. Witt","Boris Kozinsky","Albert Musaelian"],"abstract":"Machine learning interatomic potentials, particularly those based on deep equivariant neural networks, have demonstrated state-of-the-art accuracy and computational efficiency in atomistic modeling tasks like molecular dynamics and high-throughput screening. The size of datasets and demands of downstream workflows are growing rapidly, making robust and scalable software essential. This work presents a major overhaul of the NequIP framework focusing on multi-node parallelism, computational performance, and extensibility. The redesigned framework supports distributed training on large datasets and removes barriers preventing full utilization of the PyTorch 2.0 compiler at train time. We demonstrate this acceleration in a case study by training Allegro models on the SPICE 2 dataset of organic molecular systems. For inference, we introduce the first end-to-end infrastructure that uses the PyTorch Ahead-of-Time Inductor compiler for machine learning interatomic potentials. Additionally, we implement a custom kernel for the Allegro model's most expensive operation, the tensor product. Together, these advancements speed up molecular dynamics calculations on system sizes of practical relevance by up to a factor of 18.","url_abs":"https://arxiv.org/abs/2504.16068v1","url_pdf":"https://arxiv.org/pdf/2504.16068v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"high-performance-training-and-inference-for","repo_url":"https://github.com/mir-group/allegro","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":null},{"paper_slug":"high-performance-training-and-inference-for","repo_url":"https://github.com/mir-group/nequip","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"MIT"}},{"paper_slug":"high-performance-training-and-inference-for","repo_url":"https://github.com/mir-group/pair_allegro","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"MIT"}},{"paper_slug":"high-performance-training-and-inference-for","repo_url":"https://github.com/teddykoker/nequip-eqx","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":0,"framework":"jax","reach":null}],"tasks":[{"task_slug":"computational-efficiency","task_name":"Computational Efficiency"}],"methods":[{"method_slug":"speed","method_name":"SPEED"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2504.16068","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2504.16068"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/mir-group/nequip","reach":{"status":"ok","spdx":"MIT"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/mir-group/pair_allegro","reach":{"status":"ok","spdx":"MIT"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/mir-group/allegro","reach":null},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/teddykoker/nequip-eqx","reach":null}],"summary":{"unverified":6},"by_repo_kind":{"official":{"samples":6,"ran":0,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"f72f41271352db69","entry":"MuonParamGroups","repo":"mir-group/nequip","repo_kind":"official","path":"nequip/model/param_groups.py","file_url":"https://github.com/mir-group/nequip/blob/HEAD/nequip/model/param_groups.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"f72f41271352db69"}},{"code_sha256_prefix":"2c174e3cbc8cfaca","entry":"find_first_of_type","repo":"mir-group/nequip","repo_kind":"official","path":"nequip/utils/modules.py","file_url":"https://github.com/mir-group/nequip/blob/HEAD/nequip/utils/modules.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"2c174e3cbc8cfaca"}},{"code_sha256_prefix":"448f60f789fcdb8f","entry":"get_current_compile_mode","repo":"mir-group/nequip","repo_kind":"official","path":"nequip/model/utils.py","file_url":"https://github.com/mir-group/nequip/blob/HEAD/nequip/model/utils.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"448f60f789fcdb8f"}},{"code_sha256_prefix":"c96e9a2889f4c105","entry":"is_model_modifier","repo":"mir-group/nequip","repo_kind":"official","path":"nequip/nn/model_modifier_utils.py","file_url":"https://github.com/mir-group/nequip/blob/HEAD/nequip/nn/model_modifier_utils.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"c96e9a2889f4c105"}},{"code_sha256_prefix":"f8b26386fa20b570","entry":"is_persistent_model_modifier","repo":"mir-group/nequip","repo_kind":"official","path":"nequip/nn/model_modifier_utils.py","file_url":"https://github.com/mir-group/nequip/blob/HEAD/nequip/nn/model_modifier_utils.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"f8b26386fa20b570"}},{"code_sha256_prefix":"2fa3fa4c4b39e903","entry":"model_modifier","repo":"mir-group/nequip","repo_kind":"official","path":"nequip/nn/model_modifier_utils.py","file_url":"https://github.com/mir-group/nequip/blob/HEAD/nequip/nn/model_modifier_utils.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"2fa3fa4c4b39e903"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}