SkillAgentSearch skills...

Heat

Distributed tensors and Machine Learning framework with GPU and MPI acceleration in Python

Install / Use

/learn @helmholtz-analytics/Heat

README

<div align="center"> <img src="https://raw.githubusercontent.com/helmholtz-analytics/heat/main/doc/source/_static/images/logo.png"> </div>

Project Status

CPU/CUDA/ROCm tests Documentation Status coverage license: MIT PyPI Version Downloads conda-forge OpenSSF Scorecard OpenSSF Best Practices DOI Benchmarks

Heat

High-performance data analytics in Python, at scale.

Getting Started | Tutorials | Docs | Contributing


Why Heat?

Heat is a distributed tensor framework built on PyTorch and mpi4py. It provides highly optimized algorithms and data structures for tensor computations using CPUs, GPUs (CUDA/ROCm), and distributed cluster systems. It is designed to handle massive arrays that exceed the memory and computational limits of a single machine.

  • Seamless integration: Port existing NumPy/SciPy code to multi-node clusters with minimal effort.
  • Hardware-agnostic: Supports CPUs and GPUs (CUDA, ROCm, Apple MPS).
  • Efficient scaling: Exploit the entire, cumulative RAM of your cluster for memory-intensive operations.

Requirements

  • Python: >= 3.11
  • MPI: OpenMPI, MPICH, or Intel MPI
  • Frameworks: mpi4py >= 3.1, pytorch >= 2.3

Installation

# Via pip (with optional I/O support)
pip install heat[hdf5,netcdf,zarr]

# Via conda-forge
conda install -c conda-forge heat

# Via easybuild (for HPC systems)
eb heat-<version>.eb --robot

# Via spack (for HPC systems)
spack install py-heat

Distributed Example

Heat handles inter-node communication automatically. Define how your data is partitioned across the cluster using the DNDarray.split attribute. Push computations to your GPUs with the DNDarray.device attribute. Heat will take care of the rest, ensuring efficient data movement and synchronization across nodes.

Here an example from our Linear Algebra tutorial:

<details> <summary><b>View Distributed Example (mpirun / srun)</b></summary>

1. Create your script (my_script.py):

import heat as ht

split_A=0
split_B=1
M = 10000
N = 10000
K = 10000
A = ht.random.randn(M, N, split=split_A, device="gpu")
B = ht.random.randn(N, K, split=split_B, device="gpu")
C = ht.matmul(A, B)
print(C)

2. Run with MPI:

On your laptop, e.g. with OpenMPI:

mpirun -np 4 python my_script.py

On an HPC cluster with SLURM:

srun --nodes=2 --ntasks-per-node=2 --gpus-per-node=2 python my_script.py
</details>

Contributing and Support

We welcome contributions from the community. Please see our Contribution Guidelines and the Code of Conduct.

For bug reports, feature requests, or general questions, please use GitHub Issues or Discussions.

Citations

Citations are essential for the sustainability of this project. If Heat supports your work, please cite our main paper:

Götz, M., et al. (2020). HeAT - a Distributed and GPU-accelerated Tensor Framework for Data Analytics. In 2020 IEEE International Conference on Big Data (Big Data) (pp. 276-287). IEEE. DOI: 10.1109/BigData50022.2020.9378050.

<details> <summary>BibTeX</summary>
@inproceedings{heatBigData2020,
  author={Götz, Markus and Debus, Charlotte and Coquelin, Daniel and Krajsek, Kai and Comito, Claudia and Knechtges, Philipp and Hagemeier, Björn and Tarnawa, Michael and Hanselmann, Simon and Siggel, Martin and Basermann, Achim and Streit, Achim},
  booktitle={2020 IEEE International Conference on Big Data (Big Data)},
  title={HeAT – a Distributed and GPU-accelerated Tensor Framework for Data Analytics},
  year={2020},
  volume={},
  number={},
  pages={276-287},
  keywords={Heating systems;Industries;Data analysis;Big Data;Parallel processing;Libraries;Arrays;HeAT;Tensor Framework;High-performance Computing;PyTorch;NumPy;Message Passing Interface;GPU;Big Data Analytics;Machine Learning;Dask;Model Parallelism;Parallel Application Frameworks},
  doi={10.1109/BigData50022.2020.9378050}}
</details>

Acknowledgments

This work was funded by the Helmholtz Association Initiative and Networking Fund (Project ZT-I-0003, "Helmholtz Analytics Framework"); the Helmholtz AI platform grant; the European Space Agency (ESA) (Programme 4000144045); the Helmholtz Association Science Serve call 2025 (Project DB002891, HeatHub); the Google Summer of Code 2022 program.

License

Heat is distributed under the MIT license. See the LICENSE file for details.

Related Skills

View on GitHub
GitHub Stars231
CategoryData
Updated2d ago
Forks61

Languages

Python

Security Score

100/100

Audited on Mar 25, 2026

No findings