Feature/gcp tpu

.github/workflows/check-style.yml

@@ -9,18 +9,27 @@ jobs:
   black:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: psf/black@stable
         with:
           options: "--check --diff"
           version: "22.3.0"
   isort:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v3
         with:
           python-version: 3.8
       - uses: isort/isort-action@master
         with:
           isortVersion: "5.10.1"
+
+  codespell:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: codespell-project/actions-codespell@v1
+        with:
+          only_warn: 1
+          ignore_words_list: ibrary,nd

.github/workflows/push-docker-image.yml

@@ -14,7 +14,7 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
 
       - name: Docker meta
         id: meta

.github/workflows/run-benchmarks.yml

@@ -11,13 +11,13 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: 3.9
       - name: Cache dependencies
-        uses: actions/cache@v2
+        uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: Key-v1-3.9-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}
@@ -26,6 +26,9 @@ jobs:
           python -m pip install --upgrade pip
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
+      - name: Build bitsandbytes
+        run: |
+          pip install bitsandbytes==0.41.1
       - name: Build hivemind
         run: |
           pip install .

.github/workflows/run-tests.yml

@@ -11,16 +11,16 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ 3.7, 3.8, 3.9 ]
+        python-version: [ '3.7', '3.8', '3.9', '3.10', '3.11' ]
     timeout-minutes: 15
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache dependencies
-        uses: actions/cache@v2
+        uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: Key-v1-${{ matrix.python-version }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}
@@ -29,6 +29,9 @@ jobs:
           python -m pip install --upgrade pip
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
+      - name: Build bitsandbytes
+        run: |
+          pip install bitsandbytes==0.41.1
       - name: Build hivemind
         run: |
           pip install .
@@ -41,23 +44,23 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: actions/setup-go@v3
         with:
-          go-version: '1.16'
+          go-version: '1.20.11'
           check-latest: true
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: '3.8'
       - name: Cache dependencies
-        uses: actions/cache@v2
+        uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: Key-v1-3.8-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
+          python -m pip install --upgrade pip setuptools wheel
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
       - name: Build hivemind
@@ -73,27 +76,30 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 15
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: '3.8'
       - name: Cache dependencies
-        uses: actions/cache@v2
+        uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: Key-v1-3.8-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
+          python -m pip install --upgrade pip setuptools wheel
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
+      - name: Build bitsandbytes
+        run: |
+          pip install bitsandbytes==0.41.1
       - name: Build hivemind
         run: |
           pip install -e . --no-use-pep517
       - name: Test
         run: |
           export HIVEMIND_MEMORY_SHARING_STRATEGY=file_descriptor
-          pytest --cov hivemind -v tests
+          pytest --cov hivemind --cov-config=pyproject.toml -v tests
       - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v1
+        uses: codecov/codecov-action@v3

.readthedocs.yml

@@ -4,9 +4,13 @@ sphinx:
   fail_on_warning: true
 
 python:
-  version: 3.7
   install:
     - requirements: requirements.txt
     - requirements: requirements-docs.txt
     - method: pip
       path: .
+
+build:
+   os: ubuntu-22.04
+   tools:
+     python: "3.7"

CONTRIBUTING.md

@@ -38,7 +38,8 @@ with the following rules:
   cannot be longer than 119 characters.
 * We use [black](https://github.com/psf/black) for code formatting and [isort](https://github.com/PyCQA/isort) for 
   import sorting. Before submitting a PR, make sure to install and run `black .` and `isort .` in the root of the
-  repository.
+  repository. Also, you may want to check your code for typos by running `codespell --skip=".git"`, though there
+  might be false positives.
 * We highly encourage the use of [typing](https://docs.python.org/3/library/typing.html) where applicable.
 * Use `get_logger` from `hivemind.utils.logging` to log any information instead of `print`ing directly to standard
   output/error streams.

Dockerfile

@@ -9,11 +9,14 @@ RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment
 # Install packages
 RUN apt-get update && apt-get install -y --no-install-recommends --force-yes \
   build-essential \
+  curl \
   wget \
   git \
   vim \
   && apt-get clean autoclean && rm -rf /var/lib/apt/lists/{apt,dpkg,cache,log} /tmp/* /var/tmp/*
 
+RUN curl https://sh.rustup.rs -sSf | sh -s -- -y
+ENV PATH="/root/.cargo/bin:${PATH}"
 RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O install_miniconda.sh && \
   bash install_miniconda.sh -b -p /opt/conda && rm install_miniconda.sh
 ENV PATH="/opt/conda/bin:${PATH}"

MANIFEST.in

@@ -0,0 +1 @@
+include requirements*

README.md

@@ -12,10 +12,6 @@ large model on hundreds of computers from different universities, companies, and
 
 ![img](https://i.imgur.com/GPxolxb.gif)
 
-## Live Demo
-
-Check out our NeurIPS 2021 demonstration ["Training Transformers Together"](https://training-transformers-together.github.io/) to see hivemind in action, join an ongoing collaborative experiment, and learn more about the technologies behind it!
-
 ## Key Features
 
 * Distributed training without a master node: Distributed Hash Table allows connecting computers in a decentralized
@@ -28,12 +24,24 @@ Check out our NeurIPS 2021 demonstration ["Training Transformers Together"](http
   Decentralized Mixture-of-Experts ([paper](https://arxiv.org/abs/2002.04013)).
 
 To learn more about the ideas behind this library,
-see the [full list](https://github.com/learning-at-home/hivemind/tree/refer-to-discord-in-docs#citation) of our papers below.
+see the [full list](#citation) of our papers below.
+
+## Example Use Cases
+
+This section lists projects that leverage hivemind for decentralized training. 
+If you have successfully trained a model or created a downstream repository with the help of our library, 
+feel free to submit a pull request that adds your project to this list.
+
+* **Petals** ([webpage](https://petals.dev), [code](https://github.com/bigscience-workshop/petals)) — a decentralized platform for inference and fine-tuning of 100B+ language models.
+* **Training Transformers Together** ([webpage](https://training-transformers-together.github.io/), [code](https://github.com/learning-at-home/dalle-hivemind)) — a NeurIPS 2021 demonstration that trained a collaborative text-to-image Transformer model.
+* **CALM** ([webpage](https://huggingface.co/CALM), [code](https://github.com/NCAI-Research/CALM)) — a masked language model trained on a combination of Arabic datasets.
+* **sahajBERT** ([blog post](https://huggingface.co/blog/collaborative-training), [code](https://github.com/tanmoyio/sahajbert)) — a collaboratively pretrained ALBERT-xlarge for the Bengali language.
+* **HivemindStrategy** ([docs](https://lightning.ai/docs/pytorch/stable/advanced/third_party/hivemind.html?highlight=hivemindstrategy)) for PyTorch Lightning allows adapting your existing pipelines to training over slow network with unreliable peers.
 
 ## Installation
 
 Before installing, make sure that your environment has Python 3.7+
-and [PyTorch](https://pytorch.org/get-started/locally/#start-locally) 1.6.0 or newer. They can be installed either
+and [PyTorch](https://pytorch.org/get-started/locally/#start-locally) 1.9.0 or newer. They can be installed either
 natively or with [Anaconda](https://www.anaconda.com/products/individual).
 
 You can get [the latest release](https://pypi.org/project/hivemind) with pip or build hivemind from source.
@@ -46,6 +54,10 @@ If your versions of Python and PyTorch match the requirements, you can install h
 pip install hivemind
 ```
 
+Also, if you want to use blockwise 8-bit compression from [bitsandbytes](https://github.com/TimDettmers/bitsandbytes) 
+during data transfer, you can install it with `pip install hivemind[bitsandbytes]`. 
+After that, you can use the `BlockwiseQuantization` class in [hivemind.compression](./hivemind/compression)
+
 ### From source
 
 To install hivemind from source, simply run the following:
@@ -69,7 +81,8 @@ of [Go toolchain](https://golang.org/doc/install) (1.15 or 1.16 are supported).
 
 - __Linux__ is the default OS for which hivemind is developed and tested. We recommend Ubuntu 18.04+ (64-bit), but
   other 64-bit distros should work as well. Legacy 32-bit is not recommended.
-- __macOS 10.x__ can run hivemind using [Docker](https://docs.docker.com/desktop/mac/install/).
+- __macOS__ is partially supported.
+  If you have issues, you can run hivemind using [Docker](https://docs.docker.com/desktop/mac/install/) instead.
   We recommend using [our Docker image](https://hub.docker.com/r/learningathome/hivemind).
 - __Windows 10+ (experimental)__ can run hivemind
   using [WSL](https://docs.microsoft.com/ru-ru/windows/wsl/install-win10). You can configure WSL to use GPU by
@@ -111,10 +124,10 @@ If you found hivemind or its underlying algorithms useful for your research, ple
 
 ```bibtex
 @misc{hivemind,
-  author = {Learning{@}home team},
   title = {{H}ivemind: a {L}ibrary for {D}ecentralized {D}eep {L}earning},
+  author = {Learning{@}home team},
   year = 2020,
-  howpublished = {\url{https://github.com/learning-at-home/hivemind}},
+  howpublished = {\url{https://github.com/learning-at-home/hivemind}}
 }
 ```
 
@@ -124,15 +137,12 @@ at [mryab/learning-at-home](https://github.com/mryab/learning-at-home)):
 
 ```bibtex
 @inproceedings{ryabinin2020crowdsourced,
+  title = {Towards Crowdsourced Training of Large Neural Networks using Decentralized Mixture-of-Experts},
   author = {Ryabinin, Max and Gusev, Anton},
+  year = 2020,
   booktitle = {Advances in Neural Information Processing Systems},
-  editor = {H. Larochelle and M. Ranzato and R. Hadsell and M. F. Balcan and H. Lin},
-  pages = {3659--3672},
-  publisher = {Curran Associates, Inc.},
-  title = {Towards Crowdsourced Training of Large Neural Networks using Decentralized Mixture-of-Experts},
-  url = {https://proceedings.neurips.cc/paper/2020/file/25ddc0f8c9d3e22e03d3076f98d83cb2-Paper.pdf},
-  volume = {33},
-  year = {2020}
+  volume = 33,
+  url = {https://proceedings.neurips.cc/paper/2020/file/25ddc0f8c9d3e22e03d3076f98d83cb2-Paper.pdf}
 }
 ```
 
@@ -142,39 +152,53 @@ at [mryab/learning-at-home](https://github.com/mryab/learning-at-home)):
 ["Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices"](https://arxiv.org/abs/2103.03239)
 
 ```bibtex
-@misc{ryabinin2021moshpit,
+@inproceedings{ryabinin2021moshpit,
   title = {Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices},
-  author = {Max Ryabinin and Eduard Gorbunov and Vsevolod Plokhotnyuk and Gennady Pekhimenko},
-  year = {2021},
-  eprint = {2103.03239},
-  archivePrefix = {arXiv},
-  primaryClass = {cs.LG}
+  author = {Ryabinin, Max and Gorbunov, Eduard and Plokhotnyuk, Vsevolod and Pekhimenko, Gennady},
+  year = 2021,
+  booktitle = {Advances in Neural Information Processing Systems},
+  volume = 34,
+  url = {https://proceedings.neurips.cc/paper/2021/file/97275a23ca44226c9964043c8462be96-Paper.pdf}
 }
 ```
 
 ["Distributed Deep Learning in Open Collaborations"](https://arxiv.org/abs/2106.10207)
 
 ```bibtex
-@misc{diskin2021distributed,
-  title = {Distributed Deep Learning in Open Collaborations},
-  author = {Michael Diskin and Alexey Bukhtiyarov and Max Ryabinin and Lucile Saulnier and Quentin Lhoest and Anton Sinitsin and Dmitry Popov and Dmitry Pyrkin and Maxim Kashirin and Alexander Borzunov and Albert Villanova del Moral and Denis Mazur and Ilia Kobelev and Yacine Jernite and Thomas Wolf and Gennady Pekhimenko},
-  year = {2021},
-  eprint = {2106.10207},
-  archivePrefix = {arXiv},
-  primaryClass = {cs.LG}
+@inproceedings{diskin2021distributed,
+  title = {Distributed Deep Learning In Open Collaborations},
+  author = {Michael Diskin and Alexey Bukhtiyarov and Max Ryabinin and Lucile Saulnier and Quentin Lhoest and Anton Sinitsin and Dmitry Popov and Dmitriy Pyrkin and Maxim Kashirin and Alexander Borzunov and Albert Villanova del Moral and Denis Mazur and Ilia Kobelev and Yacine Jernite and Thomas Wolf and Gennady Pekhimenko},
+  year = 2021,
+  booktitle = {Advances in Neural Information Processing Systems},
+  url = {https://openreview.net/forum?id=FYHktcK-7v}
 }
 ```
 
 ["Secure Distributed Training at Scale"](https://arxiv.org/abs/2106.11257)
 
 ```bibtex
-@misc{gorbunov2021secure,
+@inproceedings{gorbunov2022secure,
   title = {Secure Distributed Training at Scale},
-  author = {Eduard Gorbunov and Alexander Borzunov and Michael Diskin and Max Ryabinin},
-  year = {2021},
-  eprint = {2106.11257},
-  archivePrefix = {arXiv},
-  primaryClass = {cs.LG}
+  author = {Gorbunov, Eduard and Borzunov, Alexander and Diskin, Michael and Ryabinin, Max},
+  year = 2022,
+  month = {17--23 Jul},
+  booktitle = {Proceedings of the 39th International Conference on Machine Learning},
+  series = {Proceedings of Machine Learning Research},
+  volume = 162,
+  url = {https://proceedings.mlr.press/v162/gorbunov22a.html}
+}
+```
+
+["Training Transformers Together"](https://arxiv.org/abs/2207.03481)
+
+```bibtex
+@misc{borzunov2022training,
+  title = {Training Transformers Together},
+  author = {Alexander Borzunov and Max Ryabinin and Tim Dettmers and Quentin Lhoest and Lucile Saulnier and Michael Diskin and Yacine Jernite and Thomas Wolf},
+  year = 2022,
+  eprint = {2207.03481},
+  archiveprefix = {arXiv},
+  primaryclass = {cs.LG}
 }
 ```
 

benchmarks/benchmark_dht.py

@@ -20,7 +20,7 @@ class NodeKiller:
     """Auxiliary class that kills dht nodes over a pre-defined schedule"""
 
     def __init__(self, shutdown_peers: list, shutdown_timestamps: list):
-        self.shutdown_peers = set(shutdown_peers)
+        self.shutdown_peers = shutdown_peers
         self.shutdown_timestamps = shutdown_timestamps
         self.current_iter = 0
         self.timestamp_iter = 0
@@ -51,7 +51,7 @@ async def store_and_get_task(
     latest: bool,
     node_killer: NodeKiller,
 ) -> Tuple[list, list, list, list, int, int]:
-    """Iteratively choose random peers to store data onto the dht, then retreive with another random subset of peers"""
+    """Iteratively choose random peers to store data onto the dht, then retrieve with another random subset of peers"""
 
     total_stores = total_gets = 0
     successful_stores = []