Repository: BlackSamorez/tensor_parallel (language: python)

Issue:
torchrun support



Discussion / hints:

Fix patch:
diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000..b924f80
Binary files /dev/null and b/.DS_Store differ
diff --git a/.github/workflows/check_style.yaml b/.github/workflows/check_style.yaml
new file mode 100644
index 0000000..94b9517
--- /dev/null
+++ b/.github/workflows/check_style.yaml
@@ -0,0 +1,26 @@
+name: Check style
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+
+jobs:
+  black:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - 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
+        with:
+          python-version: 3.8
+      - uses: isort/isort-action@master
+        with:
+          isortVersion: "5.10.1"
diff --git a/LICENCE b/LICENCE
new file mode 100644
index 0000000..44747f2
--- /dev/null
+++ b/LICENCE
@@ -0,0 +1,19 @@
+Copyright (c) 2022 by Andrei Panferov, Iaroslav Lisniak and collaborators
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file
diff --git a/README.md b/README.md
index f3f0b2a..a892380 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,56 @@
 # petals_local_parallel
 YSDA project
+
+```python
+import torch, torch.nn as nn
+model = nn.Sequential(nn.Embedding(1337, 64), nn.LayerNorm(64), nn.ReLU(), nn.Linear(128, 10))
+
+from tensor_parallel import TensorParallel
+model = TensorParallel(model, device_ids=['cuda:0', 'cuda:1'])
+
+normal_outputs = model(**normal_inputs)  # forward and backward works just like in the base model
+```
+
+## Benchmarking tutorial
+
+You may either use manual benchmark (```benchmark_manual.py```) or auto (```markbench.py```) 
+
+#### Manual benchmark
+
+consider command line arguments:
+
+```-d | do_backward``` -- wether you need backward passes or not
+
+```-n | num_iter ```   -- number of iterations
+
+```-s | seq_length```  -- sequence length
+
+```-b | batch_size```  -- okay
+
+```-c | bloomconfig``` -- str used in BloomConfig().from_pretrained to specify the model you need
+
+```CUDA_VISIBLE_DEVICES``` -- gpus, you are using
+
+```nproc_per_node```       -- # of gpus/ processes
+
+Don't forget to set correct gpu ids: ```export CUDA_DEVICE_ORDER=PCI_BUS_ID```
+
+So the following command
+``` CUDA_VISIBLE_DEVICES=4,5 torchrun --nproc_per_node 2 benchmark.py -d 0 -n 100 -s 17 -b 16 -c bloom ```
+will run the manual benchmark with no backward pass, 100 iterations, sequence length of 17, batch size of 16 and "bloom" 176B model.
+
+
+#### Auto benchmark
+
+no command line arguments this time, just run ```markbench.py```
+
+The script will run several experiments in cycle. To see the parameters, check the experiment setting section in the ```markbench.py```.
+Models are tested both with and without backward passes. The results will be printed for all of the ranks. (MESS)
+
+#### TODO:
+
+- Decide which models are too big for backward passes and don't check them
+- Decide what to do if one of the experiments failed
+
+
+
diff --git a/notebooks/gpipe.ipynb b/notebooks/gpipe.ipynb
new file mode 100644
index 0000000..453c41d
--- /dev/null
+++ b/notebooks/gpipe.ipynb
@@ -0,0 +1,279 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import List, Tuple\n",
+    "\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "from transformers.models.bloom.modeling_bloom import BloomBlock, build_alibi_tensor, _make_causal_mask, _expand_mask\n",
+    "from transformers.models.bloom.configuration_bloom import BloomConfig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 56,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "BLOOMCONFIG='bloom-1b1'\n",
+    "PIPELINEPARALLELSIZE = 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 57,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "loading configuration file config.json from cache at /home/adpanferov/.cache/huggingface/hub/models--bigscience--bloom-1b1/snapshots/1e718be072aa40714c9dc34e35ea6b64979a65ad/config.json\n",
+      "Model config BloomConfig {\n",
+      "  \"apply_residual_connection_post_layernorm\": false,\n",
+      "  \"architectures\": [\n",
+      "    \"BloomModel\"\n",
+      "  ],\n",
+      "  \"attention_dropout\": 0.0,\n",
+      "  \"attention_softmax_in_fp32\": true,\n",
+      "  \"bias_dropout_fusion\": true,\n",
+      "  \"bos_token_id\": 1,\n",
+      "  \"eos_token_id\": 2,\n",
+      "  \"hidden_dropout\": 0.0,\n",
+      "  \"hidden_size\": 1536,\n",
+      "  \"initializer_range\": 0.02,\n",
+      "  \"layer_norm_epsilon\": 1e-05,\n",
+      "  \"masked_softmax_fusion\": true,\n",
+      "  \"model_type\": \"bloom\",\n",
+      "  \"n_head\": 16,\n",
+      "  \"n_inner\": null,\n",
+      "  \"n_layer\": 24,\n",
+      "  \"offset_alibi\": 100,\n",
+      "  \"pad_token_id\": 3,\n",
+      "  \"pretraining_tp\": 1,\n",
+      "  \"skip_bias_add\": true,\n",
+      "  \"skip_bias_add_qkv\": false,\n",
+      "  \"slow_but_exact\": false,\n",
+      "  \"transformers_version\": \"4.22.2\",\n",
+      "  \"unk_token_id\": 0,\n",
+      "  \"use_cache\": true,\n",
+      "  \"vocab_size\": 250880\n",
+      "}\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = BloomConfig()\n",
+    "config = config.from_pretrained(f\"bigscience/{BLOOMCONFIG}\")\n",
+    "assert config.num_hidden_layers % PIPELINEPARALLELSIZE == 0\n",
+    "config.num_hidden_layers = config.num_hidden_layers // PIPELINEPARALLELSIZE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 77,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MiddleBloom(nn.Module):\n",
+    "    def __init__(self, config: BloomConfig):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.num_heads = config.n_head\n",
+    "        self.h = nn.ModuleList([BloomBlock(config) for _ in range(config.num_hidden_layers)])\n",
+    "        self.gradient_checkpointing = False\n",
+    "\n",
+    "    def _prepare_attn_mask(\n",
+    "        self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int\n",
+    "    ) -> torch.BoolTensor:\n",
+    "        # create causal mask\n",
+    "        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]\n",
+    "        combined_attention_mask = None\n",
+    "        device = attention_mask.device\n",
+    "        _, src_length = input_shape\n",
+    "\n",
+    "        if src_length > 1:\n",
+    "            combined_attention_mask = _make_causal_mask(\n",
+    "                input_shape, device=device, past_key_values_length=past_key_values_length\n",
+    "            )\n",
+    "\n",
+    "        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]\n",
+    "        expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)\n",
+    "        combined_attention_mask = (\n",
+    "            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask\n",
+    "        )\n",
+    "\n",
+    "        return combined_attention_mask\n",
+    "\n",
+    "    def forward(self, hidden_states, attention_mask=None, past_key_values = None):\n",
+    "        batch_size, seq_length, _ = hidden_states.shape\n",
+    "\n",
+    "        if past_key_values is None:\n",
+    "            past_key_values = tuple([None] * len(self.h))\n",
+    "        \n",
+    "        if attention_mask is None:\n",
+    "            attention_mask = torch.ones((batch_size, seq_length), device=hidden_states.device)\n",
+    "        else:\n",
+    "            attention_mask = attention_mask.to(hidden_states.device)\n",
+    "\n",
+    "        alibi = build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)\n",
+    "\n",
+    "        causal_mask = self._prepare_attn_mask(\n",
+    "            attention_mask,\n",
+    "            input_shape=(batch_size, seq_length),\n",
+    "            past_key_values_length=0,\n",
+    "        )\n",
+    "\n",
+    "        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):\n",
+    "            if self.gradient_checkpointing and self.training:\n",
+    "                def create_custom_forward(module):\n",
+    "                    def custom_forward(*inputs):\n",
+    "                        # None for past_key_value\n",
+    "                        return module(*inputs, use_cache=False, output_attentions=False)\n",
+    "\n",
+    "                    return custom_forward\n",
+    "\n",
+    "                outputs = torch.utils.checkpoint.checkpoint(\n",
+    "                    create_custom_forward(block),\n",
+    "                    hidden_states,\n",
+    "                    alibi,\n",
+    "                    None,\n",
+    "                    None,\n",
+    "                )\n",
+    "            else:\n",
+    "                outputs = block(\n",
+    "                    hidden_states,\n",
+    "                    layer_past=layer_past,\n",
+    "                    attention_mask=causal_mask,\n",
+    "                    head_mask=None,\n",
+    "                    use_cache=False,\n",
+    "                    output_attentions=False,\n",
+    "                    alibi=alibi,\n",
+    "                )\n",
+    "\n",
+    "            hidden_states = outputs[0]\n",
+    "\n",
+    "        return hidden_states\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 78,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = MiddleBloom(config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 79,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([17, 9, 1536])"
+      ]
+     },
+     "execution_count": 79,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model(torch.ones((17, 9, config.hidden_size)), None).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 82,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torchgpipe import GPipe\n",
+    "\n",
+    "model = nn.Sequential(*[MiddleBloom(config) for _ in range(PIPELINEPARALLELSIZE)])\n",
+    "model = GPipe(model,\n",
+    "              balance=[1, 1],\n",
+    "              devices=[0, 1],  # Specify GPUs.\n",
+    "              chunks=8)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 84,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([17, 9, 1536])"
+      ]
+     },
+     "execution_count": 84,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "in_device = model.devices[0]\n",
+    "out_device = model.devices[-1]\n",
+    "\n",
+    "model(torch.ones((17, 9, config.hidden_size)).to(in_device)).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.6 ('dist')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.6"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "b4f9af776eba44096d2a2b91b2ae4dd4ecd10ec5726c41598312832383f4ab0e"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..9e2ec11
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,18 @@
+[build-system]
+requires = [
+    "setuptools>=42",
+    "wheel"
+]
+build-backend = "setuptools.build_meta"
+
+[tool.black]
+line-length = 120
+force-exclude = "petals_local_parallel"
+required-version = "22.3.0"
+
+[tool.isort]
+profile = "black"
+line_length = 120
+combine_as_imports = true
+combine_star = true
+known_local_folder = ["tests", "petals_local_parallel"]
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..04fb798
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,48 @@
+[metadata]
+name = tensor_parallel
+version = 1.0.20dev
+author = Andrei Panferov and Yaroslav Lisnyak
+author_email = yalisnyak@nes.com
+description = Automatically shard your large model between multiple GPUs, works without torch.distributed
+long_description = file: README.md
+long_description_content_type = text/markdown
+url = https://github.com/BlackSamorez/tensor_parallel
+project_urls =
+    Bug Tracker = https://github.com/BlackSamorez/tensor_parallel/issues
+classifiers =
+    Development Status :: 4 - Beta
+    Intended Audience :: Developers
+    Intended Audience :: Science/Research
+    License :: OSI Approved :: MIT License
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.7
+    Programming Language :: Python :: 3.8
+    Programming Language :: Python :: 3.9
+    Programming Language :: Python :: 3.10
+    Topic :: Scientific/Engineering
+    Topic :: Scientific/Engineering :: Mathematics
+    Topic :: Scientific/Engineering :: Artificial Intelligence
+    Topic :: Software Development
+    Topic :: Software Development :: Libraries
+    Topic :: Software Development :: Libraries :: Python Modules
+
+[options]
+package_dir =
+    = src
+packages = find:
+python_requires = >=3.7
+install_requires =
+    torch>=1.12
+    transformers>=4.25.1
+[options.extras_require]
+dev =
+    pytest==6.2.5
+    pytest-forked
+    pytest-asyncio==0.16.0
+    accelerate==0.15.0
+    black==22.3.0
+    isort==5.10.1
+    psutil
+
+[options.packages.find]
+where = src
\ No newline at end of file
diff --git a/src/tensor_parallel/__init__.py b/src/tensor_parallel/__init__.py
new file mode 100644
index 0000000..5ec7126
--- /dev/null
+++ b/src/tensor_parallel/__init__.py
@@ -0,0 +1,9 @@
+"""
+A prototype TensorParallel wrapper that works without torch.distributed.
+Splits linear, conv and some other layers between GPUs
+"""
+
+from tensor_parallel.factory import tensor_parallel
+from tensor_parallel.pretrained_model import TensorParallelPreTrainedModel
+from tensor_parallel.slicer_wrapper import Config
+from tensor_parallel.tensor_parallel import TensorParallel
diff --git a/src/tensor_parallel/communications.py b/src/tensor_parallel/communications.py
new file mode 100644
index 0000000..0ce7748
--- /dev/null
+++ b/src/tensor_parallel/communications.py
@@ -0,0 +1,187 @@
+"""
+An interface for thread-parallel collective operations for use in TensorParallel
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any, List, Optional
+
+import torch
+from torch.distributed import all_gather, all_reduce
+
+import tensor_parallel.cross_device_ops as cross_device_ops
+
+
+class CollectiveOpetationBase:
+    def __call__(self, x: torch.Tensor, rank: int):
+        raise NotImplementedError()
+
+
+class CollectiveOperation(CollectiveOpetationBase):
+    def __init__(self, world_size: int, func: callable, authoritative_rank: int = 0):
+        """
+        Apply user-defined collective function in a way that is compatible with TensorParallel
+        :param func: function(input0, input1, ..., input_worldsize) -> (output0, output1, ..., output_worldsize)
+        """
+        self.world_size = world_size
+        self.func = func
+        self.authoritative_rank = authoritative_rank
+        self.rank_inputs: List[Any] = [None for _ in range(world_size)]
+        self.rank_outputs: List[Any] = [None for _ in range(world_size)]
+        self.barrier = threading.Barrier(world_size)
+
+    def __call__(self, x: torch.Tensor, rank: int):
+        try:
+            self.rank_inputs[rank] = x
+            self.barrier.wait()
+            if rank == self.authoritative_rank:
+                try:
+                    result = self.func(*self.rank_inputs)
+                    for i in range(self.world_size):
+                        self.rank_outputs[i] = (result[i], None)
+                except Exception as e:
+                    for i in range(self.world_size):
+                        self.rank_outputs[i] = (None, e)
+            self.barrier.wait()
+            result, exception = self.rank_outputs[rank]
+            if exception:
+                raise exception
+            return result
+        finally:
+            self.rank_inputs[rank] = self.rank_outputs[rank] = None
+
+
+class NCCLAllReduce(CollectiveOperation):
+    def __init__(self, world_size: int):
+        super().__init__(world_size, func=cross_device_ops.NCCLAllReduceFunction.apply)
+
+
+class NCCLAllGather(CollectiveOperation):
+    def __init__(self, world_size: int, dim: int):
+        super().__init__(world_size, func=cross_device_ops.NCCLAllGatherFunction.apply)
+        self.tensor_lengths = [None for _ in range(world_size)]
+        self.dim = dim
+
+    def __call__(self, x: torch.Tensor, rank: int):
+        # note: all-gather deliberately makes all device-local computations on their ranks, because failing to do so
+        # (making one of the ranks operate on other's tensors) may deadlock cuda streams in parallel_apply;
+        # if you modify this code and need cross-rank computations, use parallel_apply_simple to avoid stream deadlocks
+        try:
+            gather_dim = self.dim % x.ndim
+            self.tensor_lengths[rank] = x.shape[gather_dim]
+            self.barrier.wait()
+            tensor_lengths = tuple(self.tensor_lengths)
+            max_length = max(tensor_lengths)
+            if x.shape[gather_dim] < max_length:
+                pad = [0] * (x.ndim * 2)
+                pad[2 * (x.ndim - gather_dim - 1) + 1] = max_length - x.shape[gather_dim]
+                # to understand the weird indexing (2 * (ndim - dim) ...), see F.pad documentation
+                x = torch.nn.functional.pad(x, pad)
+            gathered_tensor = super().__call__(x, rank=rank)
+            used_padding = any(length != max_length for length in tensor_lengths)
+
+            if not used_padding:
+                dim_indices = list(range(1, gathered_tensor.ndim))
+                dim_indices.insert(gather_dim, 0)
+                concatenated_shape = list(x.shape)
+                concatenated_shape[gather_dim] = -1
+                return gathered_tensor.permute(dim_indices).reshape(concatenated_shape)
+            else:
+                # restore original tensor lengths by slicing off padding
+                gathered_tensor_parts = []
+                for i, tensor_length in enumerate(tensor_lengths):
+                    slices_i = [slice(None) for _ in range(gathered_tensor.ndim)]
+                    slices_i[0] = i  # select i-th element from gathered tensors
+                    slices_i[gather_dim + 1] = slice(0, tensor_length)  # remove padding
+                    gathered_tensor_parts.append(gathered_tensor[slices_i])
+                return torch.cat(gathered_tensor_parts, dim=self.dim)
+        finally:
+            self.tensor_lengths[rank] = None
+
+
+class AllReduce(CollectiveOpetationBase):
+    def __init__(
+        self,
+        world_size: int,
+        reduce_op: callable = cross_device_ops.reduce_add,
+        gather_op: callable = cross_device_ops.gather,
+    ):
+        self.scatter_reduce = ScatterReduce(world_size, reduce_op)
+        self.all_gather = AllGather(world_size, gather_op, barrier=False)
+        # note: AllGather does not need barrier here because scatter_reduce's ready event serves as barrier
+
+    def __call__(self, x: torch.Tensor, rank: int):
+        reduced_part = self.scatter_reduce(x, rank)
+        return self.all_gather(reduced_part, rank).view_as(x)
+
+
+class ScatterReduce(CollectiveOpetationBase):
+    def __init__(self, world_size: int, reduce_op: callable = cross_device_ops.reduce_add):
+        self.world_size = world_size
+        self.tensor_parts = [[] for _ in range(world_size)]
+        self.parts_ready = [threading.Event() for _ in range(world_size)]
+        self.reduce_op = reduce_op
+
+    def __call__(self, x: torch.Tensor, rank: int):
+        try:
+            for i, part in enumerate(x.flatten().tensor_split(self.world_size)):
+                self.tensor_parts[i].append(part)  # append is thread-safe. thanks, GIL!
+                if len(self.tensor_parts[i]) == self.world_size:
+                    self.parts_ready[i].set()  # can be called more than once; we don't care
+
+            self.parts_ready[rank].wait()
+            reduced_part = self.reduce_op(self.tensor_parts[rank], x.device)
+            return reduced_part
+        finally:
+            # prepare for next forward; each rank clears its own data
+            self.tensor_parts[rank].clear()
+            self.parts_ready[rank].clear()
+
+
+class AllGather(CollectiveOpetationBase):
+    def __init__(self, world_size: int, gather_op: callable = cross_device_ops.gather, barrier: bool = True):
+        self.world_size = world_size
+        self.barrier = threading.Barrier(world_size) if barrier else None
+        self.parts: List[Optional[torch.Tensor]] = [None for _ in range(world_size)]
+        self.ranks_updated = []
+        self.parts_ready = threading.Event()
+        self.gather_op = gather_op
+
+    def __call__(self, x: torch.Tensor, rank: int):
+        if self.barrier is not None:
+            self.barrier.wait()  # if this code is ran multiple times in quick succession,
+        # this even will wait for the previous call to finish before starting a new one
+        parts, ranks_updated, parts_ready = self.parts, self.ranks_updated, self.parts_ready
+        # ^-- note: we copy properties to locals so that the "finally" clause is thread-safe
+        try:
+            parts[rank] = x  # no race b/c each rank writes to a separate location
+            ranks_updated.append(rank)  # append is thread-safe. thanks, GIL!
+            if len(ranks_updated) == self.world_size:
+                parts_ready.set()  # can be called more than once; we dont care
+            parts_ready.wait()
+            # note: for one of the parts with r == rank, part.to(device) is a no-op
+            return self.gather_op(parts, x.device)
+        finally:
+            if ranks_updated[-1] == rank:
+                self.parts = [None for _ in range(self.world_size)]
+                self.ranks_updated = []
+                self.parts_ready = threading.Event()
+            # note: we can safely update these properties because all ranks have
+            # copied self.parts_* to locals before passing parts_ready.wait
+
+
+class DistributedAllReduce(CollectiveOpetationBase):
+    def __call__(self, x: torch.Tensor, rank: int):
+        all_reduce(x)
+
+
+class DistributedAllGather(CollectiveOpetationBase):
+    def __init__(self, world_size: int, dim: int):
+        self.dim = dim
+        self.world_size = world_size
+
+    def __call__(self, x: torch.Tensor, rank: int):
+        gathered_tensors = [torch.empty_like(x) for _ in range(self.world_size)]
+        all_gather(gathered_tensors, x)
+        return torch.cat(gathered_tensors, dim=self.dim)
diff --git a/src/tensor_parallel/cross_device_ops.py b/src/tensor_parallel/cross_device_ops.py
new file mode 100644
index 0000000..b417458
--- /dev/null
+++ b/src/tensor_parallel/cross_device_ops.py
@@ -0,0 +1,98 @@
+"""
+Utility functions that can process tensors that are located on multiple devices, with autograd support
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Sequence
+
+import torch
+from torch.cuda import nccl
+from torch.nn.parallel import comm
+from torch.nn.parallel._functions import Broadcast, Gather
+
+
+def broadcast_coalesced(
+    tensors: Sequence[torch.Tensor], devices: Sequence[torch.device], *, all_cuda: bool = None, **kwargs
+) -> Sequence[Sequence[torch.Tensor]]:
+    if all_cuda is None:
+        all_cuda = all(device.type == "cuda" for device in devices)
+    if not all_cuda:
+        broadcasted = [list() for _ in devices]
+        for x in tensors:
+            for i, device in enumerate(devices):
+                broadcasted[i].append(x.to(device, non_blocking=True))
+        return broadcasted
+    flat_outputs = Broadcast.apply(devices, *tensors)
+    return [flat_outputs[i * len(tensors) : (i + 1) * len(tensors)] for i in range(len(devices))]
+
+
+def gather(
+    tensors: Sequence[torch.Tensor], dim: int = 0, destination: Optional[torch.device] = None, all_cuda: bool = None
+):
+    """Gather tensors from multiple devices; differentiable w.r.t. input tensors"""
+    if all_cuda is None:
+        all_cuda = all(x.device.type == "cuda" for x in tensors)
+    if destination is None:
+        destination = tensors[0].device
+    if not all_cuda:
+        return torch.cat([x.to(destination, non_blocking=True) for x in tensors], dim=dim)
+    return Gather.apply(destination, dim, *tensors)
+
+
+def reduce_add(tensors: Sequence[torch.Tensor], destination: Optional[torch.device] = None, all_cuda: bool = None):
+    if all_cuda is None:
+        all_cuda = all(x.device.type == "cuda" for x in tensors)
+    if destination is None:
+        destination = tensors[0].device
+    if not all_cuda:
+        return sum([tensor.to(destination, non_blocking=True) for tensor in tensors])
+    return _ReduceAdd.apply(destination, *tensors)
+
+
+class _ReduceAdd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, destination: torch.device, *tensors: torch.Tensor):
+        ctx.source_gpus = [tensor.get_device() for tensor in tensors]
+        return comm.reduce_add(tensors, destination)
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        return (None,) + Broadcast.apply(ctx.source_gpus, *grad_outputs)
+
+
+class NCCLAllReduceFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, *inputs: torch.Tensor):
+        inputs = tuple(map(torch.Tensor.contiguous, inputs))
+        assert nccl.is_available(inputs)
+        outputs = tuple(map(torch.empty_like, inputs))
+        nccl.all_reduce(inputs, outputs, op=nccl.SUM)
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *grad_outputs: torch.Tensor):
+        grad_outputs = tuple(map(torch.Tensor.contiguous, grad_outputs))
+        assert nccl.is_available(grad_outputs)
+        grad_inputs = tuple(map(torch.empty_like, grad_outputs))
+        nccl.all_reduce(grad_outputs, grad_inputs, op=nccl.SUM)
+        return grad_inputs
+
+
+class NCCLAllGatherFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, *inputs: torch.Tensor):
+        world_size = len(inputs)
+        inputs = tuple(map(torch.Tensor.contiguous, inputs))
+        assert nccl.is_available(inputs)
+        outputs = tuple(torch.empty((world_size,) + x.shape, device=x.device, dtype=x.dtype) for x in inputs)
+        nccl.all_gather(inputs, outputs)
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *grad_outputs: torch.Tensor):
+        grad_outputs = tuple(map(torch.Tensor.contiguous, grad_outputs))
+        assert nccl.is_available(grad_outputs)
+        grad_inputs = tuple(torch.empty(x.shape[1:], device=x.device, dtype=x.dtype) for x in grad_outputs)
+        nccl.reduce_scatter(grad_outputs, grad_inputs, op=nccl.SUM)
+        return grad_inputs
diff --git a/src/tensor_parallel/factory.py b/src/tensor_parallel/factory.py
new file mode 100644
index 0000000..63d5b7c
--- /dev/null
+++ b/src/tensor_parallel/factory.py
@@ -0,0 +1,33 @@
+import logging
+from typing import Optional, Sequence
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+
+from tensor_parallel.pretrained_model import TensorParallelPreTrainedModel
+from tensor_parallel.slicer_wrapper import Config
+from tensor_parallel.tensor_parallel import TensorParallel
+
+logger = logging.getLogger(__file__)
+
+
+def tensor_parallel(
+    module: nn.Module,
+    device_ids: Optional[Sequence[torch.device]] = None,
+    config: Optional[Config] = None,
+    distributed: Optional[bool] = None,
+    **kwargs
+) -> nn.Module:
+    distributed = distributed if distributed is not None else torch.distributed.is_initialized()
+    if distributed:
+        if config is None:
+            config = Config.get_default_config(module)
+            logger.info("Using automatic config: sharding individual linear/conv/emb layers")
+
+        return config.make_distributed_shard(module, device=torch.device(device_ids[0]), **kwargs)
+    else:
+        if isinstance(module, PreTrainedModel):
+            return TensorParallelPreTrainedModel(module, device_ids=device_ids, config=config, **kwargs)
+        else:
+            return TensorParallel(module, device_ids=device_ids, config=config, **kwargs)
diff --git a/src/tensor_parallel/pretrained_model.py b/src/tensor_parallel/pretrained_model.py
new file mode 100644
index 0000000..1bfa9e6
--- /dev/null
+++ b/src/tensor_parallel/pretrained_model.py
@@ -0,0 +1,67 @@
+"""
+The TensorParallel module wrapper for Hugging Face PreTrainedModel
+"""
+import logging
+from typing import Any, Dict, Optional, Sequence
+
+import torch
+import transformers
+from transformers import PretrainedConfig, PreTrainedModel
+
+from tensor_parallel.slicer_wrapper import Config
+from tensor_parallel.slicing_configs import PREDEFINED_CONFIGS
+from tensor_parallel.tensor_parallel import TensorParallel, check_device_ids
+
+logger = logging.getLogger(__file__)
+
+
+def find_predefined_tensor_parallel_config(
+    model_config: PretrainedConfig, device_ids: Optional[Sequence[torch.device]]
+) -> Optional[Config]:
+    device_ids = check_device_ids(device_ids)
+    if len(model_config.architectures) != 1:
+        logger.warning(
+            f"No tensor parallel config provided and model architectures list is ambigious: {model_config.architectures}. Using possible inefficient fallback"
+        )
+        return None
+
+    try:
+        return PREDEFINED_CONFIGS[model_config.architectures[0]](model_config, device_ids)
+    except KeyError:
+        logger.warning(
+            "No tensor parallel config provided and no predefined configs can be used. Using possible inefficient fallback"
+        )
+        return None
+
+
+class TensorParallelPreTrainedModel(PreTrainedModel):
+    def __init__(
+        self,
+        module: PreTrainedModel,
+        device_ids: Optional[Sequence[torch.device]] = None,
+        output_device: Optional[torch.device] = None,
+        output_device_index: Optional[int] = None,
+        config: Optional[Config] = None,
+    ):
+        super().__init__(module.config)  # Temporary empty config. Gets replaced in from_pretrained
+
+        if config is None:
+            config = find_predefined_tensor_parallel_config(module.config, device_ids)
+
+        self.tensor_parallel = TensorParallel(module, device_ids, output_device, output_device_index, config)
+
+    def forward(self, *args, **kwargs):
+        return self.tensor_parallel(*args, **kwargs)
+
+    def _validate_model_class(self):
+        return self.tensor_parallel.module_shards[0]._validate_model_class()
+
+    def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
+        return self.tensor_parallel.module_shards[0]._validate_model_kwargs(model_kwargs)
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        return self.tensor_parallel.module_shards[0].prepare_inputs_for_generation(*args, **kwargs)
+
+    def _reorder_cache(self, past, beam_idx):
+        for shard in self.tensor_parallel.module_shards:
+            shard._reorder_cache(past, beam_idx)
diff --git a/src/tensor_parallel/slicer_wrapper.py b/src/tensor_parallel/slicer_wrapper.py
new file mode 100644
index 0000000..47cad71
--- /dev/null
+++ b/src/tensor_parallel/slicer_wrapper.py
@@ -0,0 +1,311 @@
+"""
+Tensor parallelism config and functions for splitting model into shards
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import os
+import re
+from copy import deepcopy
+from functools import partial
+from itertools import chain
+from typing import Callable, Dict, Sequence, Union
+
+import torch
+from torch import nn
+from torch.nn.modules import conv
+
+import tensor_parallel.cross_device_ops as cross_device_ops
+from tensor_parallel.communications import (
+    AllGather,
+    AllReduce,
+    DistributedAllGather,
+    DistributedAllReduce,
+    NCCLAllGather,
+    NCCLAllReduce,
+)
+
+Arg = Union[int, str]
+Pattern = Union[str, re.Pattern]
+Action = Union[str, Callable]  # actions describe what to do with tensors
+StateRules = Dict[Pattern, Action]  # state rules are pattern-matched actions on module state dict
+ModuleRules = Dict[Pattern, Dict[Arg, Action]]  # module rules are pattern-matched actions on inputs/outputs/attrs
+
+logger = logging.getLogger(__file__)
+
+
+TENSOR_PARALLEL_USE_NATIVE = bool(os.environ.get("TENSOR_PARALLEL_USE_NATIVE"))
+
+
+@dataclasses.dataclass
+class Config:
+    state_rules: StateRules
+    input_rules: ModuleRules
+    output_rules: ModuleRules
+    attr_rules: ModuleRules
+
+    def __init__(
+        self, state_rules: StateRules, input_rules: ModuleRules, output_rules: ModuleRules, attr_rules: ModuleRules
+    ):
+        all_rules = [state_rules, input_rules, output_rules, attr_rules]
+        for i, rule_set in enumerate(all_rules):
+            all_rules[i] = dict((re.compile(pattern), actions) for pattern, actions in rule_set.items())
+        self.state_rules, self.input_rules, self.output_rules, self.attr_rules = all_rules
+
+    @classmethod
+    def get_default_config(cls, module: nn.Module) -> Config:
+        """Make a generic config that wraps individual linear, embedding and convolutional layers"""
+        config = cls({}, {}, {}, {})
+        emb_weights = {m.weight for m in module.modules() if isinstance(m, (nn.Embedding, nn.EmbeddingBag))}
+
+        for name, module in module.named_modules():
+            if isinstance(module, (nn.Embedding, nn.EmbeddingBag)):
+                assert module.max_norm is None or module.norm_type < 2
+                assert getattr(module, "bias", None) is None or module.bias.shape == module.embedding_dim
+                config.state_rules[f"^{name}.weight$"] = "split 1"
+                if hasattr(module, "bias"):
+                    config.state_rules[f"^{name}.bias$"] = "split 0"
+                config.output_rules[f"^{name}$"] = {0: "gather -1"}
+            elif isinstance(module, nn.Linear):
+                assert module.weight.shape == (module.out_features, module.in_features)
+                assert module.bias is None or module.bias.shape == (module.out_features,)
+                if module.weight not in emb_weights:  # regular linear layer
+                    config.state_rules[f"^{name}.(weight|bias)$"] = "split 0"
+                    config.output_rules[f"^{name}$"] = {0: "gather -1"}
+                else:
+                    # linear weight tied with embeddings; this is a popular special case for language models;
+                    # since embedding weight will be sliced over dim 1, we should adapt to the input-sliced weight
+                    config.input_rules[f"^{name}$"] = {0: "split -1"}
+                    config.output_rules[f"^{name}$"] = {0: "sum"}
+                    if module.bias is not None:
+                        config.state_rules[f"^{name}.bias$"] = "scale"
+            elif isinstance(module, conv._ConvNd) and module.groups == 1:
+                shape = [module.out_channels, module.in_channels] + list(module.kernel_size)
+                shape[:2] = shape[:2][::-1] if module.transposed else shape[:2]
+                shape = tuple(shape)
+                assert module.weight.shape == shape, f"{module.weight.shape} != {shape}"
+                assert module.bias is None or module.bias.shape == (module.out_channels,), module.bias.shape
+                config.state_rules[f"^{name}.weight$"] = "split 1" if module.transposed else "split 0"
+                config.state_rules[f"^{name}.bias$"] = "split 0"
+                config.output_rules[f"^{name}$"] = {0: "gather 1"}
+            elif isinstance(module, conv._ConvNd) and module.groups != 1:
+                logger.warning(
+                    f"AutoConfig does not support sharding convolution layers {name} with {module.groups} groups yet; "
+                    "If you're sure that you need it, implement sharding in place of this message and open a PR"
+                )
+        return config
+
+    def create_collective_ops(self, devices):
+        """
+        Return a copy of config with thread-parallel collective operations, such as AllGather and AllReduce
+
+        :note: this function should be called during TensorParallel init, before making shards
+        """
+        return dataclasses.replace(
+            self,
+            input_rules=create_collective_ops(self.input_rules, devices),
+            output_rules=create_collective_ops(self.output_rules, devices),
+        )
+
+    def make_shard(
+        self, module: nn.Module, device: torch.device, config: Config, *, rank: int, world_size: int
+    ) -> nn.Module:
+        """Apply a tensor-parallel config to a high-level module, return module shard for a given rank and device"""
+        assert (
+            len(list(module.children())) != 0
+        ), "Please ensure module is a container (e.g. Sequential), not a single layer"
+        shard = deepcopy(module).to(device)
+        del module
+
+        # convert parameters and buffers
+        process_state_(shard, config, rank=rank, world_size=world_size)
+
+        # convert or wrap intermediate modules
+        unique_wrappers = {}
+        with torch.no_grad():
+            for name, submodule in shard.named_modules():
+                if submodule in unique_wrappers:
+                    continue  # wrap a module at most once
+                maybe_wrapped = config._maybe_wrap_submodule(name, submodule, rank=rank, world_size=world_size)
+                if maybe_wrapped:  # apply the same wrapper to all occurrences of the given module
+                    unique_wrappers[submodule] = maybe_wrapped
+
+        for parent in list(shard.modules()):
+            for child_name, child in list(parent.named_children()):
+                if child in unique_wrappers:
+                    setattr(parent, child_name, unique_wrappers[child])
+
+        return unique_wrappers.get(shard, shard)  # wrap the root module if needed
+
+    def make_distributed_shard(self, module: nn.Module, device: torch.device):
+        config_with_ops = self.create_collective_ops([torch.device("cpu")] * torch.distributed.get_world_size())
+        return self.make_shard(
+            module,
+            device,
+            config_with_ops,
+            rank=torch.distributed.get_rank(),
+            world_size=torch.distributed.get_world_size(),
+        )
+
+    def _maybe_wrap_submodule(self, name: str, module: nn.Module, *, rank: int, world_size: int) -> nn.Module:
+        """
+        Apply the tensor parallelism config to the specified module, return wrapped module
+
+        :note: this function does NOT apply state dict changes (state_rules), these should be applied separately!
+        :note: this function will modify the original module in-place!
+        """
+        attr_actions = find_matching_actions("attr", name, self.attr_rules)
+        input_actions = find_matching_actions("input", name, self.input_rules)
+        output_actions = find_matching_actions("output", name, self.output_rules)
+        if attr_actions:
+            process_attrs_(module, attr_actions, rank=rank, world_size=world_size)
+        if input_actions or output_actions:
+            module = _TensorParallelWrapper(module, input_actions, output_actions, rank=rank, world_size=world_size)
+        return module
+
+
+def find_matching_actions(action_type: str, name: str, rules: ModuleRules) -> Dict[Arg, Action]:
+    found_actions = {}
+    for pattern, actions_for_pattern in rules.items():
+        if pattern.search(name) is not None:
+            for key, action in actions_for_pattern.items():
+                if isinstance(key, str) and key.strip().isdigit():
+                    key = int(key)  # canonoicalize int keys
+                if found_actions.get(key, action) != action:
+                    raise ValueError(
+                        f"Found conflicting {action_type} rule for module {name}, key {key}:"
+                        f" {found_actions[key]} vs {action}"
+                    )
+                found_actions[key] = action
+    return found_actions
+
+
+def apply_action(input: torch.Tensor, action: Action, *, rank: int, world_size: int):
+    if callable(action):
+        return action(input, rank=rank)  # allreduce/allgather or a custom user-defined callable
+    action_type, *opts = action.split()
+    if action_type == "split":
+        dim = int(opts[0])
+        return torch.tensor_split(input, world_size, dim=dim)[rank]
+    if action_type == "split_with_sizes":
+        assert len(opts) == 3 and all(map(str.isdigit, opts))
+        dim = int(opts[0])
+        return torch.tensor_split(input, world_size, dim=dim)[rank]
+    if action_type == "scale":
+        return input / world_size
+    if action_type == "scale_int":
+        assert input % world_size == 0
+        return input // world_size
+    raise Exception(f"unexpected action {action_type}; supported actions: split, scale, or custom user-defined")
+
+
+def create_collective_ops(rules: dict, devices: Sequence[torch.device]):
+    """Initialize collective thread-parallel operations from config rules"""
+    world_size = len(devices)
+    all_cuda = all(device.type == "cuda" for device in devices)
+    unique_output_transforms = {op for output_actions in rules.values() for op in output_actions.values()}
+    transform_map = {}
+    if torch.distributed.is_initialized():
+        make_allreduce, make_allgather = DistributedAllReduce, DistributedAllGather
+    elif all_cuda and not TENSOR_PARALLEL_USE_NATIVE:
+        make_allreduce, make_allgather = NCCLAllReduce, NCCLAllGather
+    else:
+        make_allreduce = partial(
+            AllReduce,
+            reduce_op=lambda xs, destination: cross_device_ops.reduce_add(xs, destination, all_cuda=all_cuda),
+            gather_op=lambda xs, destination: cross_device_ops.gather(
+                xs, dim=0, destination=destination, all_cuda=all_cuda
+            ),
+        )
+        make_allgather = lambda world_size, dim: AllGather(
+            world_size, gather_op=lambda xs, destination: cross_device_ops.gather(xs, dim=dim, all_cuda=all_cuda)
+        )
+
+    for transform in unique_output_transforms:
+        if callable(transform):
+            continue  # user-defined transform, no action needed
+        transform_type, *opts = transform.split()
+        if transform_type in ("split", "scale", "scale_int"):
+            continue  # not a collective op, no action needed
+
+        if transform_type == "sum":
+            transform_map[transform] = make_allreduce(world_size)
+        elif transform_type == "gather":
+            dim = int(opts[0]) if opts else -1
+            transform_map[transform] = make_allgather(world_size, dim)
+
+    initialized_output_rules = {}
+    for pattern, output_actions in rules.items():
+        output_actions = {key: transform_map.get(rule, rule) for key, rule in output_actions.items()}
+        initialized_output_rules[pattern] = output_actions
+    return initialized_output_rules
+
+
+@torch.no_grad()
+def process_state_(module: nn.Module, config: Config, *, rank: int, world_size: int):
+    """Modify module parameters and/or buffers in-place"""
+    unused_patterns = set(config.state_rules.keys())
+    for name, param in chain(module.named_parameters(), module.named_buffers()):
+        for pattern, action in config.state_rules.items():
+            if pattern.search(name) is not None:
+                param.data = apply_action(param.data, action, rank=rank, world_size=world_size).clone()
+                # note: .clone is required so that the resulting parameter owns its storage
+                unused_patterns.discard(pattern)
+
+    if unused_patterns:
+        logger.warning(f"The following patterns in state_rules were unused: {[str(p) for p in unused_patterns]}")
+
+
+def process_attrs_(module: nn.Module, actions: Dict[Arg, str], rank: int, world_size: int):
+    """Modify module properties in-place"""
+    assert not getattr(module, "__tensor_parallel_process_attrs", False), "process_attrs was called more than once"
+    for attr, action in actions.items():
+        setattr(module, attr, apply_action(getattr(module, attr), action, rank=rank, world_size=world_size))
+    module.__tensor_parallel_process_attrs = True
+
+
+def process_input(input_actions: Dict[Arg, str], rank: int, world_size: int, *args, **kwargs):
+    extended_kwargs = dict(kwargs)
+    extended_kwargs.update(enumerate(args))
+    for target, action in input_actions.items():
+        extended_kwargs[target] = apply_action(extended_kwargs.get(target), action, rank=rank, world_size=world_size)
+    args = [extended_kwargs.pop(i) for i in range(len(args))]
+    return args, extended_kwargs
+
+
+def process_output(
+    output, output_actions: Dict[Arg, Callable[[torch.Tensor, int], torch.Tensor]], *, rank: int, world_size: int
+):
+    if isinstance(output, torch.Tensor):
+        return process_output({0: output}, output_actions, rank=rank, world_size=world_size)[0]
+    if isinstance(output, Sequence):
+        output_dict = process_output(dict(enumerate(output)), output_actions, rank=rank, world_size=world_size)
+        return type(output)((output_dict[i] for i in range(len(output))))
+    for target, action in output_actions.items():
+        output[target] = apply_action(output.get(target), action, rank=rank, world_size=world_size)
+    return output
+
+
+class _TensorParallelWrapper(nn.Module):
+    """Wraps a single module, applies tensor parallelism actions to module inputs and outputs"""
+
+    def __init__(
+        self,
+        module: nn.Module,
+        input_actions: Dict[Arg, Action],
+        output_actions: Dict[Arg, Action],
+        *,
+        rank: int,
+        world_size: int,
+    ):
+        super().__init__()
+        self.module = module
+        self.input_actions, self.output_actions = input_actions, output_actions
+        self.rank, self.world_size = rank, world_size
+
+    def forward(self, *args, **kwargs):
+        args, kwargs = process_input(self.input_actions, self.rank, self.world_size, *args, **kwargs)
+        output = self.module(*args, **kwargs)
+        return process_output(output, self.output_actions, rank=self.rank, world_size=self.world_size)
diff --git a/src/tensor_parallel/slicing_configs.py b/src/tensor_parallel/slicing_configs.py
new file mode 100644
index 0000000..acef0f0
--- /dev/null
+++ b/src/tensor_parallel/slicing_configs.py
@@ -0,0 +1,90 @@
+"""
+Optimized configs for selected models. These configs are not necessary, but they can improve performance in some
+cases, e.g. training with very small batches or inference with long sequences.
+
+NB: some of these configs get fairly complicated in order to squeeze a bit of extra performance. When developing your
+  own config, you can get most of the performance benefits by using auto config -- and maybe splitting MLP layers.
+"""
+from functools import partial
+from itertools import chain
+from typing import Callable, Dict, Sequence
+
+import torch
+from transformers import BloomConfig, PretrainedConfig
+
+from tensor_parallel.communications import CollectiveOperation
+from tensor_parallel.slicer_wrapper import Config
+from tensor_parallel.tensor_parallel import PerDeviceTensors
+
+ConfigGetter = Callable[[PretrainedConfig, Sequence[torch.device]], Config]
+
+
+def get_bloom_config(model_config: BloomConfig, devices: Sequence[torch.device]) -> Config:
+    world_size = len(devices)
+    num_heads = model_config.n_head
+    head_dim = model_config.hidden_size // num_heads
+    gather_kv_across_ranks = CollectiveOperation(
+        world_size=world_size, func=lambda *kvs: [PerDeviceTensors(*chain(*kvs))] * world_size
+    )  # this operation ensures that we get attention cache for all heads on each device
+
+    select_kv_for_rank = lambda kvs, rank: (kvs[2 * rank], kvs[2 * rank + 1]) if kvs else None
+
+    _split_alibi = partial(split_alibi, num_heads=num_heads, world_size=world_size)
+
+    return Config(
+        state_rules={
+            r".*self_attention\.query_key_value\.(weight|bias)$": partial(
+                split_heads, dim=0, head_dim=head_dim * 3, world_size=world_size
+            ),
+            r".*self_attention\.dense\.weight$": partial(split_heads, dim=1, head_dim=head_dim, world_size=world_size),
+            r".*self_attention\.dense\.bias$": "scale",
+            r".*mlp\.dense_h_to_4h\.(weight|bias)$": "split 0",
+            r".*mlp\.dense_4h_to_h\.weight$": "split 1",
+            r".*mlp\.dense_4h_to_h\.bias$": "scale",
+            r".*word_embeddings.weight$": "split 1",
+            # note: ^-- lm_head.weight is tied with word_embeddings
+        },
+        input_rules={
+            r".*self_attention$": {"layer_past": select_kv_for_rank, "alibi": _split_alibi},
+            r".*lm_head$": {0: "split -1"},  # note: we need to split lm_head inputs because
+            # ... lm_head's weights (tied embeddings) are already split across input dimension
+        },
+        output_rules={
+            r".*self_attention$": {1: gather_kv_across_ranks},
+            r".*self_attention\.dense$": {0: "sum"},
+            r".*mlp\.dense_4h_to_h$": {0: "sum"},
+            r".*word_embeddings$": {0: "gather -1"},
+            r".*lm_head$": {0: "sum"},
+        },
+        attr_rules={r".*self_attention$": {"num_heads": partial(split_num_heads, world_size=world_size)}},
+    )
+
+
+def split_heads(tensor: torch.Tensor, *, dim: int, head_dim: int, rank: int, world_size: int, optional: bool = False):
+    """Split a tensor along dim such that each part size is divisible by head_dim"""
+    if tensor is None and optional:
+        return None
+    assert tensor.shape[dim] % head_dim == 0, tensor.shape
+    if dim < 0:
+        dim = (tensor.ndim + dim) % tensor.ndim
+    shape = list(tensor.shape)
+    shape[dim] //= head_dim
+    shape.insert(dim + 1, head_dim)
+    tensor_part = tensor.reshape(shape).tensor_split(world_size, dim=dim)[rank].flatten(dim, dim + 1)
+    return tensor_part
+
+
+def split_num_heads(num_heads: int, *, rank: int, world_size: int):
+    return torch.empty(num_heads, device="meta").tensor_split(world_size)[rank].numel()
+
+
+def split_alibi(alibi: torch.Tensor, *, rank: int, num_heads: int, world_size: int) -> torch.Tensor:
+    """split alibi tensor of shape [batch_size * num_heads, ...] over attention heads"""
+    alibi_expanded = alibi.reshape(-1, num_heads, *alibi.shape[1:])
+    alibi_part = alibi_expanded.tensor_split(world_size, dim=1)[rank]
+    return alibi_part.reshape(-1, *alibi.shape[1:])
+
+
+PREDEFINED_CONFIGS: Dict[str, ConfigGetter] = {
+    "BloomModel": get_bloom_config,
+}
diff --git a/src/tensor_parallel/tensor_parallel.py b/src/tensor_parallel/tensor_parallel.py
new file mode 100644
index 0000000..9d0cb32
--- /dev/null
+++ b/src/tensor_parallel/tensor_parallel.py
@@ -0,0 +1,199 @@
+"""
+The main TensorParallel module wrapper
+"""
+import logging
+import os
+import threading
+from contextlib import nullcontext
+from typing import Any, Optional, Sequence, Union
+
+import torch
+from torch import nn
+from torch._utils import ExceptionWrapper, _get_all_device_indices, _get_device_index
+from torch.cuda.amp import autocast
+from torch.nn.parallel import parallel_apply
+
+from tensor_parallel.cross_device_ops import broadcast_coalesced
+from tensor_parallel.slicer_wrapper import TENSOR_PARALLEL_USE_NATIVE, Config
+from tensor_parallel.utils import nested_flatten, nested_pack
+
+logger = logging.getLogger(__file__)
+
+
+class TensorParallel(nn.Module):
+    def __init__(
+        self,
+        module: nn.Module,
+        device_ids: Optional[Sequence[torch.device]] = None,
+        output_device: Optional[torch.device] = None,
+        output_device_index: Optional[int] = None,
+        config: Optional[Config] = None,
+    ):
+        super().__init__()
+        original_params = sum(p.numel() for p in module.parameters())
+        assert output_device is None or output_device_index is None, "please specify either device or index, not both"
+        device_ids = check_device_ids(device_ids)
+
+        if output_device is not None:
+            output_device = canonicalize_device(output_device)
+            assert output_device in device_ids, f"Output device {output_device} not in {device_ids}"
+            output_device_index = device_ids.index(output_device)
+            del output_device
+        elif output_device_index is None:
+            output_device_index = 0
+
+        self.module_shards = nn.ModuleList()
+
+        self.devices = device_ids
+        self.output_device_index = output_device_index
+        self.all_cuda = all(device.type == "cuda" for device in self.devices)
+        self.device_ids = [_get_device_index(x, optional=True, allow_cpu=True) for x in device_ids]
+        world_size = len(self.devices)
+
+        if len(device_ids) <= 1:
+            self.module_shards.append(module)
+            return
+
+        if config is None:
+            config = Config.get_default_config(module)
+            logger.info("Using automatic config: sharding individual linear/conv/emb layers")
+
+        config_with_ops = config.create_collective_ops(self.devices)
+        # ^-- creates a copy of comfig with collective op instances, such as AllReduce and AllGather
+
+        for rank, device in enumerate(self.devices):
+            self.module_shards.append(
+                config.make_shard(module, device, config_with_ops, rank=rank, world_size=world_size)
+            )
+
+        # self-diagnostics: check if the model was sharded properly
+
+        params_per_shard = [sum(p.numel() for p in shard.parameters()) for shard in self.module_shards]
+        assert sum(params_per_shard) >= original_params, "Internal assert failed: lost some parameters during sharding"
+        self.param_fractions = tuple(params_i / original_params for params_i in params_per_shard)
+        inefficiency_rate = (sum(self.param_fractions) - 1) / len(device_ids)  # extra params rate per GPU
+        log_level = logging.DEBUG if inefficiency_rate < 0.1 else logging.WARNING
+        logger.log(
+            log_level,
+            f"Inefficiency warning: model has {original_params} params but shards have {params_per_shard} params. "
+            f"This means that each device uses {inefficiency_rate * 100:.3f}% extra memory for parameters",
+        )
+
+    def forward(self, *args, **kwargs):
+        if len(self.module_shards) <= 1:
+            return [self.module_shards[0](*args, **kwargs)][self.output_device_index]
+        args_and_kwargs = (args, kwargs)
+        flat_tensors = [obj for obj in nested_flatten(args_and_kwargs) if isinstance(obj, torch.Tensor)]
+        flat_tensors_replicated = broadcast_coalesced(flat_tensors, self.devices, all_cuda=self.all_cuda)
+        next_tensor_index = 0
+        args_and_kwargs_replicated = [list() for _ in self.device_ids]
+        for obj in nested_flatten(args_and_kwargs):
+            if isinstance(obj, torch.Tensor):
+                for idx in range(len(self.module_shards)):
+                    args_and_kwargs_replicated[idx].append(flat_tensors_replicated[idx][next_tensor_index])
+                next_tensor_index += 1
+            else:
+                for idx in range(len(self.module_shards)):
+                    args_and_kwargs_replicated[idx].append(obj)
+        for idx in range(len(self.module_shards)):
+            args_and_kwargs_replicated[idx] = nested_pack(args_and_kwargs_replicated[idx], args_and_kwargs)
+        inputs, kwargs_tup = zip(*args_and_kwargs_replicated)
+        if self.all_cuda and not TENSOR_PARALLEL_USE_NATIVE:
+            return parallel_apply(self.module_shards, inputs, kwargs_tup, self.devices)[self.output_device_index]
+        else:
+            return parallel_apply_simple(self.module_shards, inputs, kwargs_tup, self.devices)[self.output_device_index]
+
+
+def parallel_apply_simple(
+    modules: Sequence[nn.Module],
+    inputs: Sequence[Sequence[torch.Tensor]],
+    kwargs_tup: Optional[Any],
+    devices: Sequence[torch.device],
+) -> Sequence[Sequence[torch.Tensor]]:
+    r"""a version of parallel_apply that does not use cuda streams; somewhat slower"""
+    assert len(modules) == len(inputs)
+    if kwargs_tup is not None:
+        assert len(modules) == len(kwargs_tup)
+    else:
+        kwargs_tup = ({},) * len(modules)
+    lock = threading.Lock()
+    results = {}
+    grad_enabled, autocast_enabled = torch.is_grad_enabled(), torch.is_autocast_enabled()
+
+    def _worker(i, module, input, kwargs, device=None):
+        torch.set_grad_enabled(grad_enabled)
+        if device is None:
+            device = get_a_var(input).get_device()
+        try:
+            device_ctx = torch.cuda.device(device) if device.type == "cuda" else nullcontext()
+            with device_ctx, autocast(enabled=autocast_enabled):
+                # this also avoids accidental slicing of `input` if it is a Tensor
+                if not isinstance(input, (list, tuple)):
+                    input = (input,)
+                output = module(*input, **kwargs)
+            with lock:
+                results[i] = output
+        except Exception:
+            with lock:
+                results[i] = ExceptionWrapper(where="in replica {} on device {}".format(i, device))
+
+    if len(modules) > 1:
+        threads = [
+            threading.Thread(target=_worker, args=(i, module, input, kwargs, device))
+            for i, (module, input, kwargs, device) in enumerate(zip(modules, inputs, kwargs_tup, devices))
+        ]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+    outputs = []
+    for i in range(len(inputs)):
+        output = results[i]
+        if isinstance(output, ExceptionWrapper):
+            output.reraise()
+        outputs.append(output)
+    return outputs
+
+
+def get_a_var(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj
+
+    if isinstance(obj, list) or isinstance(obj, tuple):
+        for result in map(get_a_var, obj):
+            if isinstance(result, torch.Tensor):
+                return result
+    if isinstance(obj, dict):
+        for result in map(get_a_var, obj.items()):
+            if isinstance(result, torch.Tensor):
+                return result
+    return None
+
+
+def canonicalize_device(device: Union[torch.device, str]) -> torch.device:
+    device = torch.device(device)
+    if device.type == "cuda" and device.index is None:
+        device = torch.device(device, index=0)
+    return device
+
+
+def check_device_ids(device_ids: Optional[Sequence[torch.device]]) -> Sequence[torch.device]:
+    if device_ids is None:
+        device_ids = _get_all_device_indices() if torch.cuda.is_available() else []
+    return tuple(map(canonicalize_device, device_ids))
+
+
+class PerDeviceTensors:
+    """tensors located on different deviecs that will *not* be broadcasted when passed to TensorParallel.forward"""
+
+    def __init__(self, *tensors: torch.Tensor):
+        # note: this will not be broadcasted because broadcast_coalesced does not broadcast class properties
+        self.tensors = tuple(tensors)
+
+    def __getitem__(self, i: int):
+        return self.tensors[i]
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.tensors})"
diff --git a/src/tensor_parallel/utils.py b/src/tensor_parallel/utils.py
new file mode 100644
index 0000000..d18d20e
--- /dev/null
+++ b/src/tensor_parallel/utils.py
@@ -0,0 +1,93 @@
+"""
+Utility functions that help you process nested dicts, tuples, lists and namedtuples.
+Based on: https://stackoverflow.com/questions/49739102/python-nested-dictionary-comparison
+"""
+
+
+def nested_compare(t, u):
+    """
+    Return whether nested structure of t1 and t2 matches.
+    """
+    if isinstance(t, (list, tuple)):
+        if not isinstance(u, type(t)):
+            return False
+        if len(t) != len(u):
+            return False
+        for a, b in zip(t, u):
+            if not nested_compare(a, b):
+                return False
+        return True
+
+    if isinstance(t, dict):
+        if not isinstance(u, dict):
+            return False
+        if set(t.keys()) != set(u.keys()):
+            return False
+        for k in t:
+            if not nested_compare(t[k], u[k]):
+                return False
+        return True
+
+    else:
+        return True
+
+
+def nested_flatten(t):
+    """
+    Turn nested list/tuple/dict into a flat iterator.
+    """
+    if isinstance(t, (list, tuple)):
+        for x in t:
+            yield from nested_flatten(x)
+    elif isinstance(t, dict):
+        for k, v in sorted(t.items()):
+            yield from nested_flatten(v)
+    else:
+        yield t
+
+
+def nested_pack(flat, structure):
+    """
+    Restore nested structure from flattened state
+    :param flat: result of nested_flatten
+    :param structure: used as example when recovering structure
+    :returns: nested structure like :structure: filled with elements of :flat:
+    """
+    return _nested_pack(iter(flat), structure)
+
+
+def _nested_pack(flat_iter, structure):
+    if is_namedtuple(structure):
+        return type(structure)(*[_nested_pack(flat_iter, x) for x in structure])
+    elif isinstance(structure, (list, tuple)):
+        return type(structure)(_nested_pack(flat_iter, x) for x in structure)
+    elif isinstance(structure, dict):
+        return {k: _nested_pack(flat_iter, v) for k, v in sorted(structure.items())}
+    else:
+        return next(flat_iter)
+
+
+def is_namedtuple(x):
+    """Checks if x is a namedtuple instance. Taken from https://stackoverflow.com/a/2166841 ."""
+    t = type(x)
+    b = t.__bases__
+    if len(b) != 1 or b[0] != tuple:
+        return False
+    f = getattr(t, "_fields", None)
+    if not isinstance(f, tuple):
+        return False
+    return all(type(n) == str for n in f)
+
+
+def nested_map(fn, *t):
+    # Check arguments.
+    if not t:
+        raise ValueError("Expected 2+ arguments, got 1")
+    for i in range(1, len(t)):
+        if not nested_compare(t[0], t[i]):
+            msg = "Nested structure of %r and %r differs"
+            raise ValueError(msg % (t[0], t[i]))
+
+    # Map.
+    flat = map(nested_flatten, t)
+    return nested_pack(map(fn, *flat), t[0])


Test patch:
diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml
new file mode 100644
index 0000000..e389711
--- /dev/null
+++ b/.github/workflows/run-tests.yaml
@@ -0,0 +1,36 @@
+name: "Tests (tensor_parallel)"
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+
+jobs:
+  run_tests:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [ 3.7, 3.8, 3.9 ]
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Cache dependencies
+        uses: actions/cache@v2
+        with:
+          path: ~/.cache/pip
+          key: Key-v1-${{ matrix.python-version }}-${{ hashFiles('setup.cfg') }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+      - name: Build tensor_parallel
+        run: |
+          pip install .[dev]
+      - name: Test
+        run: |
+          cd tests
+          pytest --durations=0 --durations-min=1.0 -v
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_basic.py b/tests/test_basic.py
new file mode 100644
index 0000000..2057fd9
--- /dev/null
+++ b/tests/test_basic.py
@@ -0,0 +1,64 @@
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.nn.modules.conv import _ConvTransposeNd
+
+from tensor_parallel import TensorParallel
+
+
+@pytest.mark.parametrize("devices", [None, ("cpu",), ("cpu", "cpu"), ("cpu", "cpu", "cpu")])
+def test_embeds_and_linear(devices):
+    for emb_cls in nn.Embedding, nn.EmbeddingBag:
+        model = nn.Sequential(
+            emb_cls(num_embeddings=1337, embedding_dim=64),
+            nn.LayerNorm(64),
+            nn.Linear(64, 128),
+            nn.ReLU(),
+            nn.Linear(128, 10),
+        )
+
+        inputs = torch.randint(1, 1000, size=(1, 10))
+        ref_out = model(inputs)
+        ref_out.norm().backward()
+
+        model_tp = deepcopy(model)  # deepcopy to avoid accidental grad spillage and false positives
+        model_tp = TensorParallel(model_tp, device_ids=devices)
+        out_ours = model_tp(inputs)
+        out_ours.norm().backward()
+        assert torch.allclose(ref_out, out_ours, atol=1e-6)
+        our_grad = torch.cat([next(shard[0].parameters()).grad for shard in model_tp.module_shards], dim=1)
+        assert torch.allclose(model[0].weight.grad, our_grad, atol=1e-6)
+
+
+@pytest.mark.parametrize("devices", [None, ("cpu",), ("cpu",) * 2, ("cpu",) * 3, ("cpu",) * 4])
+@pytest.mark.parametrize("extra_options", [{}, {"padding": "same"}, {"stride": 2}, {"dilation": 2}])
+def test_convs(devices, extra_options):
+    for Conv, nd in (
+        (nn.Conv1d, 1),
+        (nn.Conv2d, 2),
+        (nn.Conv3d, 3),
+        (nn.ConvTranspose1d, 1),
+        (nn.ConvTranspose2d, 2),
+        (nn.ConvTranspose3d, 3),
+    ):
+        if issubclass(Conv, _ConvTransposeNd) and "padding" in extra_options:
+            continue  # unsupported by pytorch
+        model = nn.Sequential(
+            Conv(32, 64, kernel_size=(3,) * nd, **extra_options),
+            nn.ReLU(),
+            Conv(64, 14, kernel_size=(3,) * nd, **extra_options),
+        )
+        inputs = torch.randn(3, 32, *[10 for _ in range(nd)])
+        ref_out = model(inputs)
+        ref_out.norm().backward()
+
+        model_tp = deepcopy(model)  # deepcopy to avoid accidental grad spillage and false positives
+        model_tp = TensorParallel(model_tp, device_ids=devices)
+        out_ours = model_tp(inputs)
+        out_ours.norm().backward()
+        assert torch.allclose(ref_out, out_ours, atol=1e-6)
+        dim = 1 if model[0].transposed else 0  # concat over output channels
+        our_grad = torch.cat([next(shard[0].parameters()).grad for shard in model_tp.module_shards], dim=dim)
+        assert torch.allclose(model[0].weight.grad, our_grad, atol=1e-6)
diff --git a/tests/test_factory.py b/tests/test_factory.py
new file mode 100644
index 0000000..0ca2949
--- /dev/null
+++ b/tests/test_factory.py
@@ -0,0 +1,28 @@
+import torch.nn as nn
+import transformers
+
+from tensor_parallel import TensorParallel, tensor_parallel
+
+
+def test_factory_nn_module():
+    model = nn.Sequential(
+        nn.Embedding(num_embeddings=1337, embedding_dim=64),
+        nn.LayerNorm(64),
+        nn.Linear(64, 128),
+        nn.ReLU(),
+        nn.Linear(128, 10),
+    )
+
+    assert isinstance(model, nn.Module)
+    model = tensor_parallel(model, device_ids=["cpu", "cpu"])
+    assert isinstance(model, TensorParallel)
+
+
+def test_factory_pretrainedmodel():
+    devices = ["cpu", "cpu"]
+    model_name = "bigscience/bloom-560m"
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True).float().to(devices[0])
+
+    assert isinstance(model, transformers.PreTrainedModel)
+    model = tensor_parallel(model, device_ids=["cpu", "cpu"])
+    assert isinstance(model, transformers.PreTrainedModel)
diff --git a/tests/test_integration.py b/tests/test_integration.py
new file mode 100644
index 0000000..01fc93e
--- /dev/null
+++ b/tests/test_integration.py
@@ -0,0 +1,88 @@
+from typing import Tuple
+
+import pytest
+import torch
+from transformers import BloomConfig
+from transformers.models.bloom.modeling_bloom import BloomBlock, _expand_mask, _make_causal_mask, build_alibi_tensor
+
+from tensor_parallel import Config, TensorParallel
+
+
+@pytest.mark.parametrize("custom_config", [True, False])
+@pytest.mark.parametrize("devices", [("cpu",) * 2, ("cpu",) * 3, ("cpu",) * 4])
+def test_tp_bloom_block(devices, custom_config):
+    bloom_config = BloomConfig.from_pretrained("bigscience/bloom-560m")
+    bloom_config.torch_dtype = torch.float32
+    block = BloomBlock(bloom_config)
+
+    tp_config = None
+    if custom_config:
+        tp_config = Config(
+            state_rules={
+                ".*self_attention\.query_key_value\.(weight|bias)": "split 0",
+                ".*self_attention\.dense\.(weight|bias)": "split 0",
+                ".*mlp\.dense_h_to_4h\.(weight|bias)": "split 0",
+                ".*mlp\.dense_4h_to_h\.weight": "split 1",
+                ".*mlp\.dense_4h_to_h\.bias": "scale",
+            },
+            input_rules={},
+            output_rules={
+                ".*self_attention\.query_key_value": {0: "gather -1"},
+                ".*self_attention\.dense": {0: "gather -1"},
+                ".*mlp\.dense_4h_to_h$": {0: "sum"},
+            },
+            attr_rules={},
+        )
+
+    test_inputs1 = torch.randn(2, 3, 1024, requires_grad=True, device=devices[0])
+    test_inputs2 = test_inputs1.detach().clone().requires_grad_(True)
+    batch_size = test_inputs1.shape[0]
+    head_dim = len(block.input_layernorm.weight) // block.num_heads
+    prefix_length = 5
+
+    attention_mask = torch.ones(test_inputs1.shape[0], test_inputs1.shape[1] + prefix_length)
+    alibi = build_alibi_tensor(attention_mask, block.num_heads, dtype=test_inputs1.dtype)
+    attention_mask = _prepare_attn_mask(attention_mask, test_inputs1.shape[:2], prefix_length)
+
+    layer_past = (
+        torch.randn(batch_size * block.num_heads, head_dim, prefix_length, device=devices[0]),
+        torch.randn(batch_size * block.num_heads, prefix_length, head_dim, device=devices[0]),
+    )
+
+    grad_proj = torch.rand_like(test_inputs1)
+    y_ref, cache_ref = block(
+        test_inputs1, use_cache=True, layer_past=layer_past, alibi=alibi, attention_mask=attention_mask
+    )
+    y_ref.backward(grad_proj)
+
+    block_tp = TensorParallel(block, devices, config=tp_config)
+    y_ours, cache_ours = block_tp(
+        test_inputs2, use_cache=True, layer_past=layer_past, alibi=alibi, attention_mask=attention_mask
+    )
+    y_ours.backward(grad_proj)
+
+    assert torch.allclose(y_ours, y_ref, atol=1e-6)
+    assert torch.allclose(test_inputs1.grad, test_inputs2.grad, atol=1e-6)
+    assert torch.allclose(cache_ref[0], cache_ours[0], atol=1e-6)
+    assert torch.allclose(cache_ref[1], cache_ours[1], atol=1e-6)
+
+
+def _prepare_attn_mask(attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int):
+    # create causal mask
+    # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
+    combined_attention_mask = None
+    device = attention_mask.device
+    _, src_length = input_shape
+
+    if src_length > 1:
+        combined_attention_mask = _make_causal_mask(
+            input_shape, device=device, past_key_values_length=past_key_values_length
+        )
+
+    # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
+    expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
+    combined_attention_mask = (
+        expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
+    )
+
+    return combined_attention_mask
diff --git a/tests/test_transformers.py b/tests/test_transformers.py
new file mode 100644
index 0000000..3664930
--- /dev/null
+++ b/tests/test_transformers.py
@@ -0,0 +1,61 @@
+import pytest
+import torch
+import transformers
+from transformers import AutoTokenizer
+
+from tensor_parallel import TensorParallel, TensorParallelPreTrainedModel
+from tensor_parallel.slicing_configs import get_bloom_config
+
+
+@pytest.mark.parametrize("use_config", [False, True])
+@pytest.mark.parametrize("devices", [("cpu",) * 2, ("cpu",) * 3])
+def test_bloom_inference(use_config, devices, model_name="bigscience/bloom-560m"):
+    model_config = transformers.AutoConfig.from_pretrained(model_name)
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True).float().to(devices[0])
+
+    inp1 = torch.randint(1, 1000, size=(2, 3), device=devices[0])
+    inp2 = torch.randint(1, 1000, size=(2, 1), device=devices[0])
+    inp3 = torch.randint(1, 1000, size=(2, 2), device=devices[0])
+
+    out1_ref = model(inp1, use_cache=True, output_hidden_states=True)
+    out2_ref = model(inp2, use_cache=True, past_key_values=out1_ref.past_key_values)
+    out3_ref = model(inp3, use_cache=True, past_key_values=out2_ref.past_key_values)
+
+    tp_config = None
+    if use_config:
+        tp_config = get_bloom_config(model_config, devices)
+    model_tp = TensorParallel(model, devices, config=tp_config)
+    del model
+
+    out1 = model_tp(inp1, use_cache=True, output_hidden_states=True)
+    out2 = model_tp(inp2, use_cache=True, past_key_values=out1.past_key_values)
+    out3 = model_tp(inp3, use_cache=True, past_key_values=out2.past_key_values)
+
+    assert torch.allclose(out1_ref.hidden_states[-1], out1.hidden_states[-1], atol=3e-3)
+    assert torch.allclose(out1_ref.logits, out1.logits, atol=3e-3)
+    assert torch.allclose(out2_ref.logits, out2.logits, atol=3e-3)
+    assert torch.allclose(out3_ref.logits, out3.logits, atol=3e-3)
+
+
+@pytest.mark.parametrize("num_beams", [1, 3])
+def test_bloom_generate(num_beams, model_name="bigscience/bloom-560m"):
+    devices = ["cpu"] * 2
+    model_config = transformers.AutoConfig.from_pretrained(model_name)
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True).float().to(devices[0])
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    prompt = "Hello there!"
+
+    gen_ref = tokenizer.decode(
+        model.generate(tokenizer([prompt], return_tensors="pt")["input_ids"].to(devices[0]), num_beams=num_beams)[0]
+    )
+    assert gen_ref != prompt, "Nothing is generated. This test is unreliable"
+
+    tp_config = get_bloom_config(model_config, devices)
+    model_tp = TensorParallelPreTrainedModel(model, devices, config=tp_config)
+    del model
+
+    gen = tokenizer.decode(
+        model_tp.generate(tokenizer([prompt], return_tensors="pt")["input_ids"].to(devices[0]), num_beams=num_beams)[0]
+    )
+
+    assert gen == gen_ref