JanusMLIP

JanusPipe: Efficient Pipeline Parallel Training for Machine Learning Interatomic Potentials.

JanusMLIP is the open-source repository supporting this ICML 2026 work. The current release covers the pipeline-parallel training system described in the paper. We plan to extend the project with more distributed training and memory saving features for MLIP model training.

Current Scope

The first release adds pipeline parallel training support for conservative MLIPs. The runtime code lives under janus/pipeline, with common distributed utilities under janus/distributed.

The runnable UMA pipeline-parallel example is under:

examples/pipeline_parallel/uma/

Run the quick PP=4 example with:

bash examples/pipeline_parallel/uma/scripts/run_quick.sh

To adapt JanusPipe to another MLIP model, see docs/how_to_use_januspipe_with_your_model.md.

Installation

Create the environment and install the package in editable mode:

conda env create -f environment/environment.yml
conda activate janusmlip
pip install -e .

For existing environments, install the Python dependencies with:

pip install -r environment/requirements.txt
pip install -e .

GPU runs require a working PyTorch, CUDA, and DeepSpeed installation that matches the target cluster.

Roadmap

• Add FSDP support.
• Add pipeline parallelism support.
• Add graph parallelism support.
• Add expert parallelism support.
• Add hybrid parallel training examples.
• Add activation recomputation.

References and Thanks

JanusPipe builds upon several key components from the open-source community.

• FAIRChem (fairchem-core) for providing a centralized repository of data, models, demos, and application efforts in materials science and quantum chemistry.
• DeepSpeed for enabling extreme speed and scale in deep learning training, particularly for large language models (LLMs).

We extend our special thanks to the contributors for providing open-source resources that were invaluable to the development and execution of this project.