Repository: databricks/megablocks (language: python)

Issue:
Current installation instructions don't quite work
`docker build . -t megablocks-dev` and then `bash docker.sh` gives a docker container with the wrong version of `stanford-stk`. We also need to still run `python setup.py install` from `/mount/megablocks` to install `megablocks_ops`.


Discussion / hints:
Fixed ([1](https://github.com/stanford-futuredata/megablocks/commit/cd8a7fbcc26d875924d1f849bb3b2cc46e9b97fe), [2](https://github.com/stanford-futuredata/megablocks/commit/a1226898d93ea4667d5be531552c2576d1f73c26)).

Fix patch:
diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
new file mode 100644
index 00000000..0221752d
--- /dev/null
+++ b/.github/workflows/pr-gpu.yaml
@@ -0,0 +1,45 @@
+name: PR GPU tests
+on:
+  push:
+    branches:
+    - main
+    - release/*
+  pull_request_target:
+    branches:
+    - main
+    - release/**
+  workflow_dispatch:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main'}}
+jobs:
+  pytest-gpu:
+    name: ${{ matrix.name }}
+    runs-on: linux-ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+        - name: "python3.11-pytorch2.3.1-gpus1"
+          gpu_num: 1
+          python_version: 3.11
+          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+        - name: "python3.11-pytorch2.3.1-gpus2"
+          gpu_num: 2
+          python_version: 3.11
+          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+    steps:
+    - name: Run PR GPU tests
+      uses: mosaicml/ci-testing/.github/actions/pytest-gpu@v0.1.0
+      with:
+        name: ${{ matrix.name }}
+        container: ${{ matrix.container }}
+        python_version: ${{ matrix.python_version }}
+        gpu_num: ${{ matrix.gpu_num }}
+        git_repo: databricks/megablocks
+        pip_deps: "[all,testing]"
+        pytest_command: "coverage run -m pytest tests" # todo: remove tests from pytest tests when we delete all tests outside of MegaBlocks repo
+        pytest_markers: "gpu"
+        composer_package_name: mosaicml  # Required as Composer is built from source
+        mcloud_timeout: 3600
+        mcloud_api_key: ${{ secrets.MCLOUD_API_KEY }}
diff --git a/.gitignore b/.gitignore
index 32ec9ad7..000c47a9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,8 +1,165 @@
-*~
-gpt2-merges.txt
-gpt2-vocab.json
-core
-__pycache__
-*.egg-info
-dist
-build
\ No newline at end of file
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+docs/path/
+docs/source/_build/
+docs/source/api_reference/
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Virtual Env
+venv/
+
+# WandB
+wandb/
+
+# Neptune
+.neptune/
+
+# Spacemacs
+._#*
+.#*
+.vscode/
+
+# Vim
+*.swp
+
+# README backups
+README.md.bkp
+
+# OS X
+.DS_Store
+
+# Sphinx
+api_reference/
+
+# Node (for pyright)
+node_modules/
+package.json
+package-lock.json
+
+# pycharm
+.idea/
diff --git a/Makefile b/Makefile
index b8f859cd..5c2ab2f4 100644
--- a/Makefile
+++ b/Makefile
@@ -1,11 +1,23 @@
-clean:
-	rm -rf dist/*
+# several pytest settings
+WORLD_SIZE ?= 1  # world size for launcher tests
+MASTER_PORT ?= 26000 # port for distributed tests
+PYTHON ?= python3  # Python command
+PYTEST ?= pytest  # Pytest command
+PYRIGHT ?= pyright  # Pyright command. Pyright must be installed seperately -- e.g. `node install -g pyright`
+EXTRA_ARGS ?=  # extra arguments for pytest
+EXTRA_LAUNCHER_ARGS ?= # extra arguments for the composer cli launcher
 
-dist: clean
-	python3 setup.py sdist
+test:
+	LOCAL_WORLD_SIZE=1 $(PYTHON) -m $(PYTEST) $(EXTRA_ARGS)
 
-upload: dist
-	twine upload dist/*
+test-gpu:
+	LOCAL_WORLD_SIZE=1 $(PYTHON) -m $(PYTEST) -m gpu $(EXTRA_ARGS)
 
-upload-test: dist
-	twine upload --repository-url https://test.pypi.org/legacy/ dist/*
+# runs tests with the launcher
+test-dist:
+	$(PYTHON) -m composer.cli.launcher -n $(WORLD_SIZE) --master_port $(MASTER_PORT) $(EXTRA_LAUNCHER_ARGS) -m $(PYTEST) $(EXTRA_ARGS)
+
+test-dist-gpu:
+	$(PYTHON) -m composer.cli.launcher -n $(WORLD_SIZE) --master_port $(MASTER_PORT) $(EXTRA_LAUNCHER_ARGS) -m $(PYTEST) -m gpu $(EXTRA_ARGS)
+
+.PHONY: test test-gpu test-dist test-dist-gpu
diff --git a/README.md b/README.md
index e9083652..ee3628f0 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,8 @@ NOTE: This assumes you have `numpy` and `torch` installed.
 
 Installing `megablocks[gg]` enables dMoE computation with grouped GEMM. This feature is enabled by setting the `mlp_impl` argument to `grouped`. This is currently our recommended path for Hopper-generation GPUs.
 
+Installing `megablocks[dev]` allows you to contribute to MegaBlocks and test locally. Installing `megablocks[testing]` allows you to test via Github Actions.
+
 MegaBlocks can be installed with all dependencies via the `megablocks[all]` package.
 
 # :steam_locomotive: Usage
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..b4b90ecf
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,46 @@
+# build requirements
+[build-system]
+requires = ["setuptools < 70.0.0", "torch >= 2.3.0, < 2.4"]
+build-backend = "setuptools.build_meta"
+
+# Pytest
+[tool.pytest.ini_options]
+# By default, run tests in README codeblocks, raise typos in fixtures as errors,
+# have a short traceback, skip gpu tests, and display in color
+addopts = "--codeblocks --strict-markers --tb=short -m 'not gpu' --color=yes"
+pythonpath = [
+  "."
+]
+
+markers = [
+    # For distributed testing
+    "world_size(val)",
+    # Should be run during daily regression
+    "daily",
+    # Whether the test will be reading data from a remote source, and may require credentials
+    "remote",
+    # whether the test requires a gpu
+    "gpu",
+]
+
+filterwarnings = [
+    # "error",  # warnings should be treated like errors, but still need to fix some warnings
+    'ignore:ExtraArgumentWarning',  # extra arguments originate from pytest-specific CLI args
+    'ignore:DistributedDefaultValueWarning',  # default distributed values are fine
+    'ignore:NoDistributedWarning',  # running without distributed is fine
+    'ignore:Deterministic mode is activated:UserWarning',  # all tests run with deterministic mode
+    'ignore:SubsetNumBatchesWarning',  # different subsets OK for testing
+    'ignore:No optimizer:UserWarning',  # testing defaults
+    'ignore:No scheduler:UserWarning',  # testing defaults
+    'ignore::DeprecationWarning:tensorboard',  # ignore tensorboard
+]
+
+# Coverage
+[tool.coverage.run]
+parallel = true
+branch = true
+relative_files = true
+concurrency = ["thread"]
+include = [
+    "megablocks/*"
+]
diff --git a/setup.py b/setup.py
index 95c4a0e8..ac1b43fc 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,7 @@
-from setuptools import setup, find_packages
 import os
+
 import torch
+from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 if os.environ.get("TORCH_CUDA_ARCH_LIST"):
@@ -28,22 +29,37 @@
     )
 ]
 
-install_requires=[
-    "triton>=2.1.0",
-    "stanford-stk==0.7.0",
+install_requires = [
+    'numpy>=1.21.5,<2.1.0',
+    'torch>=2.3.0,<2.4',
+    'triton>=2.1.0',
+    'stanford-stk @ git+https://git@github.com/stanford-futuredata/stk.git@a1ddf98466730b88a2988860a9d8000fd1833301',
+    'packaging>=21.3.0,<24.2',
 ]
 
 extra_deps = {}
 
 extra_deps["gg"] = [
-    "grouped_gemm==0.1.4",
+    'grouped_gemm @ git+https://git@github.com/tgale96/grouped_gemm.git@66c7195e35e8c4f22fa6a014037ef511bfa397cb',
+]
+
+extra_deps['dev'] = [
+    'absl-py', # todo: delete when finish removing all absl tests
+    'coverage[toml]==7.4.4',
+    'pytest_codeblocks>=0.16.1,<0.17',
+    'pytest-cov>=4,<5',
+    'pytest>=7.2.1,<8',
+    'pre-commit>=3.4.0,<4',
 ]
 
-extra_deps["dev"] = [
-    "absl-py",
+extra_deps['testing'] = [
+    'mosaicml>=0.22.0',
 ]
 
-extra_deps['all'] = set(dep for deps in extra_deps.values() for dep in deps)
+extra_deps['all'] = list({
+    dep for key, deps in extra_deps.items() for dep in deps
+    if key not in {'testing'}
+})
 
 setup(
     name="megablocks",


Test patch:
diff --git a/megablocks/layers/dmoe_test.py b/megablocks/layers/dmoe_test.py
deleted file mode 100644
index dfba7935..00000000
--- a/megablocks/layers/dmoe_test.py
+++ /dev/null
@@ -1,166 +0,0 @@
-import unittest
-from functools import partial
-
-from absl.testing import parameterized
-from megablocks import grouped_gemm_util as gg
-from megablocks.layers.arguments import Arguments
-from megablocks.layers import dmoe
-from megablocks.layers import moe
-from megablocks.layers import testing
-import torch
-
-
-def test_modules(
-        hidden_size,
-        ffn_hidden_size,
-        moe_num_experts=1,
-        moe_capacity_factor=1,
-        moe_top_k=1,
-        mlp_impl='sparse'):
-    init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
-    args = Arguments(
-        hidden_size=hidden_size,
-        ffn_hidden_size=ffn_hidden_size,
-        moe_num_experts=moe_num_experts,
-        moe_capacity_factor=moe_capacity_factor,
-        moe_top_k=moe_top_k,
-        init_method=init_method,
-        memory_optimized_mlp=True,
-        mlp_type='mlp',
-        mlp_impl=mlp_impl,
-        fp16=False,
-        bf16=True)
-
-    mlp = testing.FFN(args)
-    moe_mlp = moe.MoE(args)
-    dmoe_mlp = dmoe.dMoE(args)
-
-    mlp.cuda(torch.cuda.current_device()).to(torch.bfloat16)
-    moe_mlp.cuda(torch.cuda.current_device()).to(torch.bfloat16)
-    dmoe_mlp.cuda(torch.cuda.current_device()).to(torch.bfloat16)
-
-    # Set the baseline parameters to match exactly.
-    with torch.no_grad():
-        ne, hs, fhs = moe_mlp.experts.mlp.w1.size()
-        w1 = dmoe_mlp.experts.mlp.w1.view([ne, fhs, hs])
-        moe_mlp.experts.mlp.w1.copy_(torch.transpose(w1, 1, 2).contiguous())
-        moe_mlp.experts.mlp.w2.copy_(dmoe_mlp.experts.mlp.w2.view([ne, fhs, hs]))
-        moe_mlp.router.layer.weight.copy_(dmoe_mlp.router.layer.weight)
-        if moe_num_experts == 1:
-            mlp.w1.copy_(moe_mlp.experts.mlp.w1.squeeze())
-            mlp.w2.copy_(moe_mlp.experts.mlp.w2.squeeze())
-    return args, mlp, moe_mlp, dmoe_mlp
-
-# min size: (1, 2, 128, 2, 1)
-_FORWARD_TESTS_DEFAULT = (
-    (16, 1024, 512, 1, 1),
-    (16, 1024, 512, 2, 1),
-    (16, 1024, 512, 4, 1),
-    (16, 1024, 512, 8, 1),
-    (8, 2048, 512, 1, 1),
-    (8, 2048, 512, 2, 1),
-    (8, 2048, 512, 4, 1),
-    (16, 1024, 512, 2, 2),
-    (16, 1024, 512, 4, 2),
-    (16, 1024, 512, 4, 4),
-    (16, 1024, 512, 8, 2),
-    (16, 1024, 512, 8, 4),
-    (16, 1024, 512, 8, 8),
-    (16, 1024, 128, 1, 1),
-)
-
-_FORWARD_TESTS_GROUPED_MLP = tuple([
-    p + ('grouped',) for p in _FORWARD_TESTS_DEFAULT
-]) if gg.grouped_gemm_is_available() else ()
-
-_FORWARD_TESTS = (_FORWARD_TESTS_DEFAULT + _FORWARD_TESTS_GROUPED_MLP)
-
-
-_DENSE_TESTS = (
-    (16, 1024, 512),
-    (8, 2048, 512),
-)
-
-
-class dMoETest(parameterized.TestCase):
-
-    @staticmethod
-    def tearDown():
-        moe.clear_load_balancing_loss()
-
-    @parameterized.parameters(*_FORWARD_TESTS)
-    def testdMoE_Forward(self, bs, sl, hs, num_experts, top_k,
-                         mlp_impl='sparse'):
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-
-        _, _, _, layer = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            moe_num_experts=num_experts,
-            moe_top_k=top_k,
-            mlp_impl=mlp_impl)
-
-        out, _ = layer(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-
-    @parameterized.parameters(*_FORWARD_TESTS)
-    def testdMoE_ForwardBackward(
-            self, bs, sl, hs, num_experts, top_k,
-            mlp_impl='sparse'):
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-        x.requires_grad_(True)
-
-        args, _, _, layer = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            moe_num_experts=num_experts,
-            moe_top_k=top_k,
-            mlp_impl=mlp_impl)
-
-        out, _ = layer(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-        loss = out.sum() + moe.batched_load_balancing_loss(args)
-        loss.backward()
-        self.assertTrue(x.grad is not None)
-        layer.zero_grad(set_to_none=True)
-        x.grad = None
-        moe.clear_load_balancing_loss()
-
-    @parameterized.parameters(*_DENSE_TESTS)
-    def testdMoE_ForwardVersusBaseline(self, bs, sl, hs):
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-
-        _, mlp, _, dmoe_mlp = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2)
-
-        expected_out = mlp(x)
-        out, _ = dmoe_mlp(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-        self.assertSequenceEqual(expected_out.shape, x.shape)
-        self.assertTrue(testing.allclose(out, expected_out))
-
-    @parameterized.parameters(*_FORWARD_TESTS)
-    def testdMoE_ForwardVersusMoE(
-            self, bs, sl, hs, num_experts, top_k,
-            mlp_impl='sparse'):
-        torch.manual_seed(42)
-
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-
-        _, _, moe_mlp, dmoe_mlp = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs,
-            moe_num_experts=num_experts,
-            moe_capacity_factor=0,
-            mlp_impl=mlp_impl)
-
-        expected_out, _= moe_mlp(x)
-        out, _ = dmoe_mlp(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-        self.assertSequenceEqual(expected_out.shape, x.shape)
-        self.assertTrue(testing.allclose(out, expected_out))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/layers/glu_test.py b/megablocks/layers/glu_test.py
deleted file mode 100644
index 7422340f..00000000
--- a/megablocks/layers/glu_test.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import unittest
-from functools import partial
-
-from absl.testing import parameterized
-from megablocks.layers.arguments import Arguments
-from megablocks.layers.glu import SparseGLU, GroupedGLU
-from megablocks.layers import dmlp_registry
-from megablocks.layers import testing
-
-import torch
-import stk
-import numpy as np
-
-def test_modules(
-        hidden_size,
-        ffn_hidden_size,
-        mlp_impl='sparse',
-        memory_optimized_mlp=False):
-    init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
-    args = Arguments(
-        hidden_size=hidden_size,
-        ffn_hidden_size=ffn_hidden_size,
-        moe_num_experts=1,
-        moe_top_k=1,
-        init_method=init_method,
-        memory_optimized_mlp=memory_optimized_mlp,
-        mlp_type='glu',
-        mlp_impl=mlp_impl,
-        fp16=False,
-        bf16=True)
-
-    glu = testing.GLU(args)
-    dmoe_glu = dmlp_registry.get(args)
-
-    dmoe_glu.cuda(torch.cuda.current_device()).to(torch.bfloat16)
-    glu.cuda(torch.cuda.current_device()).to(torch.bfloat16)
-
-    with torch.no_grad():
-        glu.w1.copy_(dmoe_glu.w1.T)
-        glu.v1.copy_(dmoe_glu.v1.T)
-        glu.w2.copy_(dmoe_glu.w2)
-
-    return args, glu, dmoe_glu
-
-_DENSE_TESTS = (
-    (16, 1024, 512),
-    (8, 2048, 512),
-)
-
-class GLUTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_DENSE_TESTS)
-    def testGLU_ForwardGroupedMLP(self, bs, sl, hs):
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-
-        _, glu, dmoe_glu = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            mlp_impl='grouped')
-
-        expected_out = glu(x)
-        tokens_per_expert = torch.tensor([bs * sl]).cuda()
-        out = dmoe_glu(x.view(bs * sl, hs), tokens_per_expert)
-        out = out.view(sl, bs, hs)
-
-        self.assertSequenceEqual(out.shape, x.shape)
-        self.assertSequenceEqual(expected_out.shape, x.shape)
-        self.assertTrue(testing.allclose(out, expected_out))
-
-    @parameterized.parameters(*_DENSE_TESTS)
-    def testGLU_ForwardGroupedMLP_MemOpt(self, bs, sl, hs):
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-
-        _, glu, dmoe_glu = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            mlp_impl='grouped',
-            memory_optimized_mlp=True)
-
-        expected_out = glu(x)
-        tokens_per_expert = torch.tensor([bs * sl]).cuda()
-        out = dmoe_glu(x.view(bs * sl, hs), tokens_per_expert)
-        out = out.view(sl, bs, hs)
-
-        self.assertSequenceEqual(out.shape, x.shape)
-        self.assertSequenceEqual(expected_out.shape, x.shape)
-        self.assertTrue(testing.allclose(out, expected_out))
-
-    @parameterized.parameters(*_DENSE_TESTS)
-    def testGLU_ForwardSparseMLP(self, bs, sl, hs):
-        x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
-
-        _, glu, dmoe_glu = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            mlp_impl='sparse')
-
-        expected_out = glu(x)
-        with torch.no_grad():
-            topo = stk.random.mask(bs * sl, hs * 2, 0, blocking=128).cuda()
-        out = dmoe_glu(x.view(bs * sl, hs), topo)
-        out = out.view(sl, bs, hs)
-
-        self.assertSequenceEqual(out.shape, x.shape)
-        self.assertSequenceEqual(expected_out.shape, x.shape)
-        self.assertTrue(testing.allclose(out, expected_out))
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/layers/moe_test.py b/megablocks/layers/moe_test.py
deleted file mode 100644
index 08f6bc8a..00000000
--- a/megablocks/layers/moe_test.py
+++ /dev/null
@@ -1,151 +0,0 @@
-import unittest
-from functools import partial
-
-from absl.testing import parameterized
-from megablocks.layers.arguments import Arguments
-from megablocks.layers import moe
-from megablocks.layers import testing
-import torch
-
-
-def test_modules(
-        hidden_size,
-        ffn_hidden_size,
-        moe_num_experts=1,
-        moe_capacity_factor=1,
-        moe_top_k=1):
-    init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
-    args = Arguments(
-        hidden_size=hidden_size,
-        ffn_hidden_size=ffn_hidden_size,
-        moe_num_experts=moe_num_experts,
-        moe_capacity_factor=moe_capacity_factor,
-        moe_top_k=moe_top_k,
-        init_method=init_method)
-
-    mlp = testing.FFN(args)
-    moe_mlp = moe.MoE(args)
-
-    mlp.cuda(torch.cuda.current_device()).half()
-    moe_mlp.cuda(torch.cuda.current_device()).half()
-
-    # Set the baseline parameters to match exactly.
-    if moe_num_experts == 1:
-        with torch.no_grad():
-            mlp.w1.copy_(moe_mlp.experts.mlp.w1.squeeze())
-            mlp.w2.copy_(moe_mlp.experts.mlp.w2.squeeze())
-    return args, mlp, moe_mlp
-
-
-_FORWARD_TESTS = (
-    (16, 1024, 512, 1, 1),
-    (16, 1024, 512, 2, 1),
-    (16, 1024, 512, 4, 1),
-    (16, 1024, 512, 8, 1),
-    (8, 2048, 512, 1, 1),
-    (8, 2048, 512, 2, 1),
-    (8, 2048, 512, 4, 1),
-    (16, 1024, 512, 2, 2),
-    (16, 1024, 512, 4, 2),
-    (16, 1024, 512, 4, 4),
-    (16, 1024, 512, 8, 2),
-    (16, 1024, 512, 8, 4),
-    (16, 1024, 512, 8, 8),
-)
-
-
-_DENSE_TESTS = (
-    (16, 1024, 512),
-    (8, 2048, 512),
-)
-
-
-class MoETest(parameterized.TestCase):
-
-    @staticmethod
-    def tearDown():
-        moe.clear_load_balancing_loss()
-
-    @parameterized.parameters(*_FORWARD_TESTS)
-    def testMoE_Forward(
-            self, bs, sl, hs, num_experts, top_k):
-        x = torch.randn(sl, bs, hs).half().cuda()
-
-        _, _, layer = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            moe_num_experts=num_experts,
-            moe_top_k=top_k)
-
-        out, _ = layer(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-
-    @parameterized.parameters(*_FORWARD_TESTS)
-    def testMoE_ForwardBackward(
-            self, bs, sl, hs, num_experts, top_k):
-        x = torch.randn(sl, bs, hs).half().cuda()
-        x.requires_grad_(True)
-
-        args, _, layer = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2,
-            moe_num_experts=num_experts,
-            moe_top_k=top_k)
-
-        out, _ = layer(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-        loss = out.sum() + moe.batched_load_balancing_loss(args)
-        loss.backward()
-        layer.zero_grad(set_to_none=True)
-        x.grad = None
-        moe.clear_load_balancing_loss()
-
-    @parameterized.parameters(*_DENSE_TESTS)
-    def testMoE_ForwardVersusDense(self, bs, sl, hs):
-        x = torch.randn(sl, bs, hs).half().cuda()
-
-        _, mlp, moe_mlp = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2)
-
-        expected_out = mlp(x)
-        out, _ = moe_mlp(x)
-        self.assertSequenceEqual(out.shape, x.shape)
-        self.assertSequenceEqual(expected_out.shape, x.shape)
-        self.assertTrue(testing.allclose(out, expected_out))
-
-    @parameterized.parameters(*_DENSE_TESTS)
-    def testMoE_ForwardBackwardVersusDense(self, bs, sl, hs):
-        x = torch.randn(sl, bs, hs).half().cuda()
-        x.requires_grad_(True)
-
-        _, mlp, moe_mlp = test_modules(
-            hidden_size=hs,
-            ffn_hidden_size=hs * 2)
-
-        out, _ = moe_mlp(x)
-        loss = out.sum()
-        loss.backward()
-        w1_grad = moe_mlp.experts.mlp.w1.grad.detach().squeeze()
-        w2_grad = moe_mlp.experts.mlp.w2.grad.detach().squeeze()
-        moe_mlp.zero_grad(set_to_none=True)
-        x.grad = None
-        moe.clear_load_balancing_loss()
-
-        expected_out = mlp(x)
-        expected_loss = expected_out.sum()
-        expected_loss.backward()
-        expected_w1_grad = mlp.w1.grad.detach()
-        expected_w2_grad = mlp.w2.grad.detach()
-        mlp.zero_grad(set_to_none=True)
-        x.grad = None
-
-        # Verify the gradients match.
-        self.assertSequenceEqual(w1_grad.shape, expected_w1_grad.shape)
-        self.assertTrue(testing.allclose(w1_grad, expected_w1_grad))
-        self.assertSequenceEqual(w2_grad.shape, expected_w2_grad.shape)
-        self.assertTrue(testing.allclose(w2_grad, expected_w2_grad))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/layers/parallelism_test.sh b/megablocks/layers/parallelism_test.sh
deleted file mode 100644
index e71560a0..00000000
--- a/megablocks/layers/parallelism_test.sh
+++ /dev/null
@@ -1,12 +0,0 @@
-#!/bin/bash
-
-DISTRIBUTED_ARGUMENTS="\
---nproc_per_node 8 \
---nnodes 1 \
---node_rank 0 \
---master_addr localhost \
---master_port 6000"
-
-python -m torch.distributed.launch \
-       ${DISTRIBUTED_ARGUMENTS} \
-       megablocks/layers/parallelism_test.py
diff --git a/megablocks/ops/binned_gather_test.py b/megablocks/ops/binned_gather_test.py
deleted file mode 100644
index bddc3ac7..00000000
--- a/megablocks/ops/binned_gather_test.py
+++ /dev/null
@@ -1,69 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import numpy as np
-import torch
-
-
-_BINNED_GATHER_TESTS = (
-    (4, 2, 2, 1),
-    (4, 2, 2, 2),
-    (4, 2, 2, 4),
-    (1024, 1536, 4, 1),
-    (1024, 1536, 4, 2),
-    (1024, 1536, 4, 4),
-    (1024, 1536, 64, 1),
-    (1024, 1536, 64, 2),
-    (1024, 1536, 64, 4),
-    (1024, 1536, 128, 1),
-    (1024, 1536, 128, 2),
-    (1024, 1536, 128, 4),
-    (16384, 768, 4, 1),
-    (16384, 768, 4, 2),
-    (16384, 768, 4, 4),
-    (16384, 768, 64, 1),
-    (16384, 768, 64, 2),
-    (16384, 768, 64, 4),
-    (16384, 768, 128, 1),
-    (16384, 768, 128, 2),
-    (16384, 768, 128, 4),
-)
-
-
-class BinnedGatherTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_BINNED_GATHER_TESTS)
-    def testBinnedGather(self, sl, hs, ne, top_k):
-        # NOTE: Capacity factor == 1.
-        ec = (sl * top_k) // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        _, indices = ops.sort(top_expert)
-        bins = ops.inclusive_cumsum(ops.histogram(top_expert, ne), 0)
-
-        def binned_gather(x, indices, bins, ec, top_k):
-            x = x.cpu().numpy()
-            indices = indices.cpu().numpy()
-            bins = bins.cpu().numpy()
-            start = 0
-            out = np.zeros((ne, ec, hs))
-            for i in range(ne):
-                end = bins[i]
-                for j in range(min(ec, end - start)):
-                    index = indices[start + j] // top_k
-                    out[i, j, :] = x[index, :]
-                start = end
-            return torch.from_numpy(out).cuda().half()
-
-        out = ops.binned_gather(x, indices, bins, ec, top_k)
-        expected_out = binned_gather(x, indices, bins, ec, top_k)
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/binned_scatter_test.py b/megablocks/ops/binned_scatter_test.py
deleted file mode 100644
index 63904c0c..00000000
--- a/megablocks/ops/binned_scatter_test.py
+++ /dev/null
@@ -1,81 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import numpy as np
-import torch
-
-
-_BINNED_SCATTER_TESTS = (
-    (4, 2, 2, 1),
-    (4, 2, 2, 2),
-    (4, 2, 2, 4),
-    (1024, 1536, 4, 1),
-    (1024, 1536, 4, 2),
-    (1024, 1536, 4, 4),
-    (1024, 1536, 64, 1),
-    (1024, 1536, 64, 2),
-    (1024, 1536, 64, 4),
-    (1024, 1536, 128, 1),
-    (1024, 1536, 128, 2),
-    (1024, 1536, 128, 4),
-    (16384, 768, 4, 1),
-    (16384, 768, 4, 2),
-    (16384, 768, 4, 4),
-    (16384, 768, 64, 1),
-    (16384, 768, 64, 2),
-    (16384, 768, 64, 4),
-    (16384, 768, 128, 1),
-    (16384, 768, 128, 2),
-    (16384, 768, 128, 4),
-)
-
-
-class BinnedScatterTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_BINNED_SCATTER_TESTS)
-    def testBinnedScatter(self, sl, hs, ne, top_k):
-        # NOTE: Capacity factor == 1.
-        ec = (sl * top_k) // ne
-
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        _, indices = ops.sort(top_expert)
-        bins = ops.inclusive_cumsum(ops.histogram(top_expert, ne), 0)
-
-        # Sample weights for the scatter reduce.
-        weights = torch.rand((sl * top_k,)).cuda().half()
-
-        x = ops.binned_gather(x, indices, bins, ec, top_k)
-
-        def binned_scatter(x, indices, weights, bins, top_k):
-            x = x.cpu().numpy()
-            indices = indices.cpu().numpy()
-            weights = weights.cpu().numpy()
-            bins = bins.cpu().numpy()
-            start = 0
-            out = np.zeros((sl, hs))
-            for i in range(ne):
-                end = bins[i]
-                for j in range(min(ec, end - start)):
-                    index = indices[start + j]
-                    scale = weights[index]
-                    index //= top_k
-
-                    out[index, :] += scale * x[i, j, :]
-                start = end
-            return torch.from_numpy(out).cuda().half()
-        out = ops.binned_scatter(x, indices, weights, bins, top_k)
-        expected_out = binned_scatter(x, indices, weights, bins, top_k)
-
-        # NOTE: We need to check approximate equality because the
-        # scatter reduce uses atomics.
-        np.testing.assert_allclose(
-            out.cpu(), expected_out.cpu(), rtol=5e-3)
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/cumsum_test.py b/megablocks/ops/cumsum_test.py
deleted file mode 100644
index a69d47b2..00000000
--- a/megablocks/ops/cumsum_test.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import torch
-
-
-_CUMSUM_TESTS = (
-    (1, 32),
-    (2, 32),
-    (2, 1024),
-    (4, 1024),
-    (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (2, 16384),
-    (4, 16384),
-    (8, 16384),
-    (16, 16384),
-    (32, 16384),
-    (64, 16384),
-    (128, 16384),
-)
-
-
-class CumsumTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_CUMSUM_TESTS)
-    def testExclusiveCumsum(self, n, m):
-        x = torch.randint(0, 2, (n, m)).long().cuda()
-        out = ops.exclusive_cumsum(x, 1) * x
-        expected_out = (torch.cumsum(x, dim=1) - 1) * x
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-    @parameterized.parameters(*_CUMSUM_TESTS)
-    def testInclusiveCumsum(self, n, m):
-        x = torch.randint(0, 2, (n, m)).long().cuda()
-        out = ops.inclusive_cumsum(x, 1)
-        expected_out = torch.cumsum(x, dim=1)
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/padded_gather_test.py b/megablocks/ops/padded_gather_test.py
deleted file mode 100644
index 290bc1ab..00000000
--- a/megablocks/ops/padded_gather_test.py
+++ /dev/null
@@ -1,91 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import numpy as np
-import torch
-
-
-_PADDED_GATHER_TESTS = (
-    (4, 2, 2, 1),
-    (4, 2, 2, 2),
-    (1024, 1, 4, 1),
-    (1024, 1, 4, 2),
-    (1024, 1, 4, 4),
-    (1024, 1, 64, 1),
-    (1024, 1, 64, 2),
-    (1024, 1, 64, 4),
-    (1024, 1, 128, 1),
-    (1024, 1, 128, 2),
-    (1024, 1, 128, 4),
-    (1024, 1536, 4, 1),
-    (1024, 1536, 4, 2),
-    (1024, 1536, 4, 4),
-    (1024, 1536, 64, 1),
-    (1024, 1536, 64, 2),
-    (1024, 1536, 64, 4),
-    (1024, 1536, 128, 1),
-    (1024, 1536, 128, 2),
-    (1024, 1536, 128, 4),
-    (16384, 768, 4, 1),
-    (16384, 768, 4, 2),
-    (16384, 768, 4, 4),
-    (16384, 768, 64, 1),
-    (16384, 768, 64, 2),
-    (16384, 768, 64, 4),
-    (16384, 768, 128, 1),
-    (16384, 768, 128, 2),
-    (16384, 768, 128, 4),
-    (16384, 1, 4, 1),
-    (16384, 1, 4, 2),
-    (16384, 1, 4, 4),
-    (16384, 1, 64, 1),
-    (16384, 1, 64, 2),
-    (16384, 1, 64, 4),
-    (16384, 1, 128, 1),
-    (16384, 1, 128, 2),
-    (16384, 1, 128, 4),
-)
-
-
-class PaddedGatherTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_PADDED_GATHER_TESTS)
-    def testPaddedGather(self, sl, hs, ne, top_k):
-        # Create the data and indices.
-        x = torch.randn((sl, hs)).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        def padded_gather(x, indices, bin_ids, bins, padded_bins, top_k):
-            x = x.cpu().numpy()
-            indices = indices.cpu().numpy()
-            bin_ids = bin_ids.cpu().numpy()
-            bins = bins.cpu().numpy()
-            padded_bins = padded_bins.cpu().numpy()
-
-            out = np.zeros((padded_bins[-1], hs))
-            in_idx = 0
-            for i in range(len(bins)):
-                out_idx = 0 if i == 0 else padded_bins[i - 1]
-                end = bins[i]
-                while in_idx < end:
-                    load_idx = indices[in_idx] // top_k
-                    out[out_idx, :] = x[load_idx, :]
-                    in_idx += 1
-                    out_idx += 1
-            return torch.from_numpy(out).cuda().half()
-
-        out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
-        expected_out = padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/padded_scatter_test.py b/megablocks/ops/padded_scatter_test.py
deleted file mode 100644
index 6b7928b3..00000000
--- a/megablocks/ops/padded_scatter_test.py
+++ /dev/null
@@ -1,128 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import numpy as np
-import torch
-
-
-_PADDED_SCATTER_TESTS = (
-    (4, 2, 2, 1),
-    (4, 2, 2, 2),
-    (4, 2, 2, 1),
-    (4, 2, 2, 1),
-    (4, 2, 2, 2),
-    (4, 2, 2, 2),
-    (1024, 1, 4, 1),
-    (1024, 1, 4, 2),
-    (1024, 1, 4, 4),
-    (1024, 1, 4, 1),
-    (1024, 1, 4, 2),
-    (1024, 1, 4, 4),
-    (1024, 1, 4, 1),
-    (1024, 1, 4, 2),
-    (1024, 1, 4, 4),
-    (1024, 1, 64, 1),
-    (1024, 1, 64, 2),
-    (1024, 1, 64, 4),
-    (1024, 1, 128, 1),
-    (1024, 1, 128, 2),
-    (1024, 1, 128, 4),
-    (1024, 1536, 4, 1),
-    (1024, 1536, 4, 2),
-    (1024, 1536, 4, 4),
-    (1024, 1536, 4, 4),
-    (1024, 1536, 4, 4),
-    (1024, 1536, 64, 1),
-    (1024, 1536, 64, 2),
-    (1024, 1536, 64, 4),
-    (1024, 1536, 128, 1),
-    (1024, 1536, 128, 2),
-    (1024, 1536, 128, 4),
-    (1024, 1536, 128, 1),
-    (1024, 1536, 128, 1),
-    (16384, 768, 4, 1),
-    (16384, 768, 4, 2),
-    (16384, 768, 4, 4),
-    (16384, 768, 64, 1),
-    (16384, 768, 64, 2),
-    (16384, 768, 64, 4),
-    (16384, 768, 128, 1),
-    (16384, 768, 128, 2),
-    (16384, 768, 128, 4),
-    (16384, 1, 4, 1),
-    (16384, 1, 4, 2),
-    (16384, 1, 4, 4),
-    (16384, 1, 64, 1),
-    (16384, 1, 64, 2),
-    (16384, 1, 64, 4),
-    (16384, 1, 128, 1),
-    (16384, 1, 128, 2),
-    (16384, 1, 128, 4),
-    (16384, 1, 128, 2),
-    (16384, 1, 128, 2),
-)
-
-
-def _to_numpy(x: torch.Tensor) -> np.ndarray:
-    return x.detach().cpu().numpy()
-
-class PaddedScatterTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_PADDED_SCATTER_TESTS)
-    def testPaddedScatter(self, sl, hs, ne, top_k):
-        # Create the data and indices.
-        x = torch.randn((sl, hs), requires_grad=True).cuda().half()
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
-        bin_ids, indices = ops.sort(top_expert)
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-        bins = ops.inclusive_cumsum(tokens_per_expert, 0)
-
-        # Sample weights for the scatter reduce.
-        weights = torch.rand((sl * top_k,), requires_grad=True).cuda().half()
-
-        # Gather the data to prepare for backwards.
-        x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
-
-        def padded_scatter(x, indices, bin_ids, weights, bins, padded_bins, top_k):
-            x = x.detach().cpu().numpy()
-            indices = _to_numpy(indices)
-            bin_ids = _to_numpy(bin_ids)
-            weights = _to_numpy(weights)
-            bins = _to_numpy(bins)
-            padded_bins = _to_numpy(padded_bins)
-
-            out = np.zeros((indices.shape[0] // top_k, hs))
-            out_idx = 0
-            for i in range(len(bins)):
-                in_idx = 0 if i == 0 else padded_bins[i - 1]
-                end = bins[i]
-                while out_idx < end:
-                    store_idx = indices[out_idx]
-                    scale = weights[store_idx]
-                    store_idx //= top_k
-
-                    out[store_idx, :] += scale * x[in_idx, :]
-                    out_idx += 1
-                    in_idx += 1
-            return torch.from_numpy(out).cuda().half()
-
-        out = ops.padded_scatter(
-            x, indices, bin_ids, weights, bins, padded_bins, top_k)
-        expected_out = padded_scatter(
-            x, indices, bin_ids, weights, bins, padded_bins, top_k)
-
-        out.backward(torch.randn_like(out)) # sanity check backward pass
-
-        # NOTE: We need to check approximate equality because the
-        # scatter reduce uses atomics.
-        np.testing.assert_allclose(
-            _to_numpy(out), _to_numpy(expected_out), rtol=5e-3)
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/replicate_test.py b/megablocks/ops/replicate_test.py
deleted file mode 100644
index 507d092b..00000000
--- a/megablocks/ops/replicate_test.py
+++ /dev/null
@@ -1,107 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import megablocks_ops as backend
-import numpy as np
-import torch
-
-
-_REPLICATE_TESTS = (
-    (8, 1, 1),
-    (8, 2, 1),
-    (8, 4, 1),
-    (8, 8, 1),
-    (8, 2, 2),
-    (8, 4, 2),
-    (8, 8, 2),
-    (8, 2, 4),
-    (8, 4, 4),
-    (8, 8, 4),
-    (8, 2, 8),
-    (8, 4, 8),
-    (8, 8, 8),
-    (16384, 2, 1),
-    (16384, 4, 1),
-    (16384, 8, 1),
-    (16384, 16, 1),
-    (16384, 32, 1),
-    (16384, 64, 1),
-    (16384, 128, 1),
-    (16384, 2, 2),
-    (16384, 4, 2),
-    (16384, 8, 2),
-    (16384, 16, 2),
-    (16384, 32, 2),
-    (16384, 64, 2),
-    (16384, 128, 2),
-    (16384, 2, 4),
-    (16384, 4, 4),
-    (16384, 8, 4),
-    (16384, 16, 4),
-    (16384, 32, 4),
-    (16384, 64, 4),
-    (16384, 128, 4),
-    (16384, 2, 8),
-    (16384, 4, 8),
-    (16384, 8, 8),
-    (16384, 16, 8),
-    (16384, 32, 8),
-    (16384, 64, 8),
-    (16384, 128, 8),
-)
-
-
-def promote_scalar(x):
-    return x.view(1) if not len(x.size()) else x
-
-
-class ReplicateTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_REPLICATE_TESTS)
-    def testReplicate(self, tokens, num_centers, top_k):
-        tokens_to_centers = torch.randint(0, num_centers, (tokens,)).cuda().int()
-        tokens_per_center = ops.histogram(tokens_to_centers, num_centers)
-        bins = ops.inclusive_cumsum(tokens_per_center, 0)
-        bins = promote_scalar(bins)
-        center_weights = torch.randn(top_k, num_centers).cuda().half()
-
-        def replicate(x, bins, num_outputs):
-            x = x.cpu().numpy()
-            bins = bins.cpu().numpy()
-            out = np.zeros((x.shape[0], num_outputs))
-            for batch_idx in range(x.shape[0]):
-                start = 0
-                for i in range(len(bins)):
-                    end = bins[i]
-                    value = x[batch_idx, i]
-                    while start < end:
-                        out[batch_idx, start] = value
-                        start += 1
-            return torch.from_numpy(out).cuda().half()
-
-        out = ops.replicate(center_weights, bins, tokens)
-        expected_out = replicate(center_weights, bins, tokens)
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-    @parameterized.parameters(*_REPLICATE_TESTS)
-    def testReplicate_Backward(self, tokens, num_centers, top_k):
-        tokens_to_centers = torch.randint(0, num_centers, (tokens,)).cuda().int()
-        tokens_per_center = ops.histogram(tokens_to_centers, num_centers)
-        bins = ops.inclusive_cumsum(tokens_per_center, 0)
-        bins = promote_scalar(bins)
-        center_weights = torch.randn(top_k, num_centers).cuda().half()
-
-        grad = ops.replicate(center_weights, bins, tokens)
-
-        out = torch.empty_like(center_weights)
-        backend.replicate_backward(grad, bins, out)
-        expected_out = center_weights * tokens_per_center.view([1, num_centers])
-
-        # NOTE: This floating-point reduction could be a problem for
-        # training stability and accuracy.
-        self.assertTrue(torch.allclose(out, expected_out, rtol=1e-2))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/sort_test.py b/megablocks/ops/sort_test.py
deleted file mode 100644
index 9b7694b3..00000000
--- a/megablocks/ops/sort_test.py
+++ /dev/null
@@ -1,64 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import numpy as np
-import torch
-
-
-_SORT_TESTS = (
-    (32, torch.int16, None),
-    (1024, torch.int16, None),
-    (16384, torch.int16, None),
-    (32, torch.int32, None),
-    (1024, torch.int32, None),
-    (16384, torch.int32, None),
-    (32, torch.int64, None),
-    (1024, torch.int64, None),
-    (16384, torch.int64, None),
-    (32, torch.int16, 128),
-    (1024, torch.int16, 128),
-    (16384, torch.int16, 128),
-    (32, torch.int32, 128),
-    (1024, torch.int32, 128),
-    (16384, torch.int32, 128),
-    (32, torch.int64, 128),
-    (1024, torch.int64, 128),
-    (16384, torch.int64, 128),
-)
-
-
-def numpy_dtype(dtype):
-    types = {
-        torch.int16: np.int16,
-        torch.int32: np.int32,
-        torch.int64: np.int64
-    }
-    return types[dtype]
-
-
-class SortTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_SORT_TESTS)
-    def testSort(self, n, dtype, max_val):
-        if max_val is None:
-            max_val = np.iinfo(numpy_dtype(dtype)).max
-        end_bit = int(np.ceil(np.log2(max_val)))
-        x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
-
-        out, indices = ops.sort(x, end_bit)
-        expected_out, expected_indices = torch.sort(x)
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-        # NOTE: The indices can be in different order depending
-        # on sort stability if multiple values in the array are
-        # equal.
-        data = torch.empty_like(x)
-        data.scatter_(0, indices.long(), out)
-        expected_data = torch.empty_like(x)
-        expected_data.scatter_(0, expected_indices, expected_out)
-        self.assertTrue(torch.all(torch.eq(data, expected_data)))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/megablocks/ops/topology_test.py b/megablocks/ops/topology_test.py
deleted file mode 100644
index 0ff3491e..00000000
--- a/megablocks/ops/topology_test.py
+++ /dev/null
@@ -1,77 +0,0 @@
-import unittest
-
-from absl.testing import parameterized
-from megablocks import ops
-import math
-import numpy as np
-import torch
-
-
-_TOPOLOGY_TESTS = (
-    (1024, 1536, 2),
-    (1024, 1536, 4),
-    (1024, 1536, 8),
-    (1024, 1536, 16),
-    (1024, 1536, 32),
-    (1024, 1536, 64),
-    (1024, 1536, 128),
-    (1024, 1536, 256),
-    (1024, 1536, 512),
-    (16384, 768, 2),
-    (16384, 768, 4),
-    (16384, 768, 8),
-    (16384, 768, 16),
-    (16384, 768, 32),
-    (16384, 768, 64),
-    (16384, 768, 128),
-    (16384, 768, 256),
-    (16384, 768, 512),
-    (16384, 768, 1024),
-    (8, 14336, 8),
-)
-
-
-class TopologyTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_TOPOLOGY_TESTS)
-    def testTopology(self, sl, hs, ne):
-        # Create the data and indices.
-        blocking = 128
-        assert hs % blocking == 0
-
-        # Randomly assign tokens to experts.
-        top_expert = torch.randint(0, ne, (sl,)).cuda().int()
-        tokens_per_expert = ops.histogram(top_expert, ne)
-        padded_tokens_per_expert = ops.round_up(tokens_per_expert, blocking)
-        padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
-
-        # Dimensions for the output indices.
-        output_block_rows = int(padded_bins[-1]) // blocking
-        output_block_columns = hs // blocking
-
-        def topology(padded_bins, blocking, rows, columns):
-            padded_bins = padded_bins.cpu().numpy()
-
-            out = np.zeros([rows * columns])
-            start = 0
-            for i in range(padded_bins.shape[0]):
-                end = padded_bins[i] // blocking
-                while start < end:
-                    for j in range(columns):
-                        out[start * columns + j] = j + i * columns
-                    start += 1
-            return torch.from_numpy(out).cuda().short()
-
-        out = ops.topology(padded_bins,
-                           blocking,
-                           output_block_rows,
-                           output_block_columns)
-        expected_out = topology(padded_bins,
-                                blocking,
-                                output_block_rows,
-                                output_block_columns)
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..335140ce
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,97 @@
+import os
+from typing import List, Optional
+
+import pytest
+from composer.utils import reproducibility
+
+# Allowed options for pytest.mark.world_size()
+WORLD_SIZE_OPTIONS = (1, 2)
+
+# Enforce deterministic mode before any tests start.
+reproducibility.configure_deterministic_mode()
+
+# Add the path of any pytest fixture files you want to make global
+pytest_plugins = [
+    'tests.fixtures.autouse',
+    'tests.fixtures.fixtures'
+]
+
+
+def _get_world_size(item: pytest.Item):
+    """Returns the world_size of a test, defaults to 1."""
+    _default = pytest.mark.world_size(1).mark
+    return item.get_closest_marker('world_size', default=_default).args[0]
+
+
+
+def _get_option(config: pytest.Config, name: str, default: Optional[str] = None) -> str:  # type: ignore
+    val = config.getoption(name)
+    if val is not None:
+        assert isinstance(val, str)
+        return val
+    val = config.getini(name)
+    if val == []:
+        val = None
+    if val is None:
+        if default is None:
+            pytest.fail(f'Config option {name} is not specified but is required')
+        val = default
+    assert isinstance(val, str)
+    return val
+
+
+def _add_option(parser: pytest.Parser, name: str, help: str, choices: Optional[list[str]] = None):
+    parser.addoption(
+        f'--{name}',
+        default=None,
+        type=str,
+        choices=choices,
+        help=help,
+    )
+    parser.addini(
+        name=name,
+        help=help,
+        type='string',
+        default=None,
+    )
+
+
+def pytest_collection_modifyitems(
+    config: pytest.Config,
+    items: List[pytest.Item],
+) -> None:
+    """Filter tests by world_size (for multi-GPU tests)"""
+    world_size = int(os.environ.get('WORLD_SIZE', '1'))
+    print(f'world_size={world_size}')
+
+    conditions = [
+        lambda item: _get_world_size(item) == world_size,
+    ]
+
+    # keep items that satisfy all conditions
+    remaining = []
+    deselected = []
+    for item in items:
+        if all(condition(item) for condition in conditions):
+            remaining.append(item)
+        else:
+            deselected.append(item)
+
+    if deselected:
+        config.hook.pytest_deselected(items=deselected)
+        items[:] = remaining
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    _add_option(
+        parser,
+        'seed',
+        help="""\
+        Rank zero seed to use. `reproducibility.seed_all(seed + dist.get_global_rank())` will be invoked
+        before each test.""",
+    )
+
+
+def pytest_sessionfinish(session: pytest.Session, exitstatus: int):
+    if exitstatus == 5:
+        session.exitstatus = 0  # Ignore no-test-ran errors
diff --git a/tests/fixtures/autouse.py b/tests/fixtures/autouse.py
new file mode 100644
index 00000000..29fbdeb8
--- /dev/null
+++ b/tests/fixtures/autouse.py
@@ -0,0 +1,97 @@
+import gc
+import logging
+import os
+
+import composer
+import pytest
+import torch
+from composer.devices import DeviceCPU, DeviceGPU
+from composer.utils import dist, reproducibility
+
+
+@pytest.fixture(autouse=True)
+def clear_cuda_cache(request: pytest.FixtureRequest):
+    """Clear memory between GPU tests."""
+    marker = request.node.get_closest_marker('gpu')
+    if marker is not None and torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()  # Only gc on GPU tests as it 2x slows down CPU tests
+
+
+@pytest.fixture(autouse=True)
+def reset_mlflow_tracking_dir():
+    """Reset MLFlow tracking dir so it doesn't persist across tests"""
+    try:
+        import mlflow
+        mlflow.set_tracking_uri(None)  # type: ignore
+    except ModuleNotFoundError:
+        # MLFlow not installed
+        pass
+
+
+@pytest.fixture(scope='session')
+def cleanup_dist():
+    """Ensure all dist tests clean up resources properly."""
+    yield
+    # Avoid race condition where a test is still writing to a file on one rank
+    # while the file system is being torn down on another rank.
+    dist.barrier()
+
+
+@pytest.fixture(autouse=True, scope='session')
+def configure_dist(request: pytest.FixtureRequest):
+    # Configure dist globally when the world size is greater than 1,
+    # so individual tests that do not use the trainer
+    # do not need to worry about manually configuring dist.
+
+    if dist.get_world_size() == 1:
+        return
+
+    device = None
+
+    for item in request.session.items:
+        device = DeviceCPU() if item.get_closest_marker('gpu') is None else DeviceGPU()
+        break
+
+    assert device is not None
+
+    if not dist.is_initialized():
+        dist.initialize_dist(device, timeout=300.0)
+    # Hold PyTest until all ranks have reached this barrier. Ensure that no rank starts
+    # any test before other ranks are ready to start it, which could be a cause of random timeouts
+    # (e.g. rank 1 starts the next test while rank 0 is finishing up the previous test).
+    dist.barrier()
+
+
+@pytest.fixture(autouse=True)
+def set_log_levels():
+    """Ensures all log levels are set to DEBUG."""
+    logging.basicConfig()
+    logging.getLogger(composer.__name__).setLevel(logging.DEBUG)
+
+
+@pytest.fixture(autouse=True)
+def seed_all(rank_zero_seed: int, monkeypatch: pytest.MonkeyPatch):
+    """Monkeypatch reproducibility get_random_seed to always return the rank zero seed, and set the random seed before
+    each test to the rank local seed."""
+    monkeypatch.setattr(reproducibility, 'get_random_seed', lambda: rank_zero_seed)
+    reproducibility.seed_all(rank_zero_seed + dist.get_global_rank())
+
+
+@pytest.fixture(autouse=True)
+def remove_run_name_env_var():
+    # Remove environment variables for run names in unit tests
+    composer_run_name = os.environ.get('COMPOSER_RUN_NAME')
+    run_name = os.environ.get('RUN_NAME')
+
+    if 'COMPOSER_RUN_NAME' in os.environ:
+        del os.environ['COMPOSER_RUN_NAME']
+    if 'RUN_NAME' in os.environ:
+        del os.environ['RUN_NAME']
+
+    yield
+
+    if composer_run_name is not None:
+        os.environ['COMPOSER_RUN_NAME'] = composer_run_name
+    if run_name is not None:
+        os.environ['RUN_NAME'] = run_name
diff --git a/tests/fixtures/fixtures.py b/tests/fixtures/fixtures.py
new file mode 100644
index 00000000..48645a89
--- /dev/null
+++ b/tests/fixtures/fixtures.py
@@ -0,0 +1,9 @@
+import pytest
+from tests.conftest import _get_option
+
+
+@pytest.fixture
+def rank_zero_seed(pytestconfig: pytest.Config) -> int:
+    """Read the rank_zero_seed from the CLI option."""
+    seed = _get_option(pytestconfig, 'seed', default='0')
+    return int(seed)
diff --git a/tests/layers/dmoe_test.py b/tests/layers/dmoe_test.py
new file mode 100644
index 00000000..3ead862a
--- /dev/null
+++ b/tests/layers/dmoe_test.py
@@ -0,0 +1,180 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import partial
+
+import pytest
+import torch
+
+from megablocks import grouped_gemm_util as gg
+from megablocks.layers import dmoe, moe, testing
+from megablocks.layers.arguments import Arguments
+
+# min size: (1, 2, 128, 2, 1)
+_FORWARD_TESTS_DEFAULT = (
+    (16, 1024, 512, 1, 1),
+    (16, 1024, 512, 2, 1),
+    (16, 1024, 512, 4, 1),
+    (16, 1024, 512, 8, 1),
+    (8, 2048, 512, 1, 1),
+    (8, 2048, 512, 2, 1),
+    (8, 2048, 512, 4, 1),
+    (16, 1024, 512, 2, 2),
+    (16, 1024, 512, 4, 2),
+    (16, 1024, 512, 4, 4),
+    (16, 1024, 512, 8, 2),
+    (16, 1024, 512, 8, 4),
+    (16, 1024, 512, 8, 8),
+    (16, 1024, 128, 1, 1),
+)
+
+_FORWARD_TESTS_GROUPED_MLP = tuple([
+    p + ('grouped',) for p in _FORWARD_TESTS_DEFAULT
+]) if gg.grouped_gemm_is_available() else ()
+
+_FORWARD_TESTS_SPARSE_MLP = tuple([
+    p + ('sparse',) for p in _FORWARD_TESTS_DEFAULT
+])
+
+_FORWARD_TESTS = (_FORWARD_TESTS_SPARSE_MLP + _FORWARD_TESTS_GROUPED_MLP)
+
+_DENSE_TESTS = (
+    (16, 1024, 512),
+    (8, 2048, 512),
+)
+
+
+def construct_moes(hidden_size: int,
+                   ffn_hidden_size: int,
+                   moe_num_experts: int = 1,
+                   moe_capacity_factor: int = 1,
+                   moe_top_k: int = 1,
+                   mlp_impl: str = 'sparse'):
+    init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
+    args = Arguments(hidden_size=hidden_size,
+                     ffn_hidden_size=ffn_hidden_size,
+                     moe_num_experts=moe_num_experts,
+                     moe_capacity_factor=moe_capacity_factor,
+                     moe_top_k=moe_top_k,
+                     init_method=init_method,
+                     memory_optimized_mlp=True,
+                     mlp_type='mlp',
+                     mlp_impl=mlp_impl,
+                     fp16=False,
+                     bf16=True)
+
+    mlp = testing.FFN(args)
+    moe_mlp = moe.MoE(args)
+    dmoe_mlp = dmoe.dMoE(args)
+
+    mlp.cuda(torch.cuda.current_device()).to(torch.bfloat16)
+    moe_mlp.cuda(torch.cuda.current_device()).to(torch.bfloat16)
+    dmoe_mlp.cuda(torch.cuda.current_device()).to(torch.bfloat16)
+
+    # Set the baseline parameters to match exactly.
+    with torch.no_grad():
+        ne, hs, fhs = moe_mlp.experts.mlp.w1.size()
+        w1 = dmoe_mlp.experts.mlp.w1.view([ne, fhs, hs])
+        moe_mlp.experts.mlp.w1.copy_(torch.transpose(w1, 1, 2).contiguous())
+        moe_mlp.experts.mlp.w2.copy_(dmoe_mlp.experts.mlp.w2.view([ne, fhs,
+                                                                   hs]))
+        moe_mlp.router.layer.weight.copy_(dmoe_mlp.router.layer.weight)
+        if moe_num_experts == 1:
+            mlp.w1.copy_(moe_mlp.experts.mlp.w1.squeeze())
+            mlp.w2.copy_(moe_mlp.experts.mlp.w2.squeeze())
+    return args, mlp, moe_mlp, dmoe_mlp
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs', 'num_experts', 'top_k', 'mlp_impl'),
+                         _FORWARD_TESTS)
+def test_dmoe_forward(bs: int,
+                      sl: int,
+                      hs: int,
+                      num_experts: int,
+                      top_k: int,
+                      mlp_impl: str):
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+    _, _, _, layer = construct_moes(hidden_size=hs,
+                                    ffn_hidden_size=hs * 2,
+                                    moe_num_experts=num_experts,
+                                    moe_top_k=top_k,
+                                    mlp_impl=mlp_impl)
+
+    out, _ = layer(x)
+    assert out.shape == x.shape
+    moe.clear_load_balancing_loss()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs', 'num_experts', 'top_k', 'mlp_impl'),
+                         _FORWARD_TESTS)
+def test_dmoe_forward_backward(bs: int,
+                               sl: int,
+                               hs: int,
+                               num_experts: int,
+                               top_k: int,
+                               mlp_impl: str):
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+    x.requires_grad_(True)
+
+    args, _, _, layer = construct_moes(hidden_size=hs,
+                                       ffn_hidden_size=hs * 2,
+                                       moe_num_experts=num_experts,
+                                       moe_top_k=top_k,
+                                       mlp_impl=mlp_impl)
+
+    out, _ = layer(x)
+    assert out.shape == x.shape
+    loss = out.sum() + moe.batched_load_balancing_loss(args)
+    loss.backward()
+    assert x.grad is not None
+    layer.zero_grad(set_to_none=True)
+    x.grad = None
+    moe.clear_load_balancing_loss()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs'), _DENSE_TESTS)
+def test_dmoe_forward_vs_baseline(bs: int,
+                                  sl: int,
+                                  hs: int,
+                                  mlp_impl: str = 'sparse'):
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+
+    _, mlp, _, dmoe_mlp = construct_moes(hidden_size=hs,
+                                         ffn_hidden_size=hs * 2,
+                                         moe_num_experts=1,
+                                         moe_capacity_factor=1,
+                                         moe_top_k=1,
+                                         mlp_impl=mlp_impl)
+
+    expected_out = mlp(x)
+    out, _ = dmoe_mlp(x)
+    assert out.shape == x.shape == expected_out.shape
+    assert torch.allclose(out, expected_out)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs', 'num_experts', 'top_k', 'mlp_impl'),
+                         _FORWARD_TESTS)
+def test_dmoe_forward_vs_moe(bs: int,
+                             sl: int,
+                             hs: int,
+                             num_experts: int,
+                             top_k: int,
+                             mlp_impl: str):
+    torch.manual_seed(42)
+
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+
+    _, _, moe_mlp, dmoe_mlp = construct_moes(hidden_size=hs,
+                                             ffn_hidden_size=hs,
+                                             moe_num_experts=num_experts,
+                                             moe_capacity_factor=0,
+                                             mlp_impl=mlp_impl)
+
+    expected_out, _ = moe_mlp(x)
+    out, _ = dmoe_mlp(x)
+    assert out.shape == x.shape == expected_out.shape
+    assert torch.allclose(out, expected_out)
diff --git a/tests/layers/glu_test.py b/tests/layers/glu_test.py
new file mode 100644
index 00000000..0487ec88
--- /dev/null
+++ b/tests/layers/glu_test.py
@@ -0,0 +1,106 @@
+from functools import partial
+
+import pytest
+import stk
+import torch
+
+from megablocks.layers import dmlp_registry, testing
+from megablocks.layers.arguments import Arguments
+from megablocks.layers.glu import GroupedGLU, SparseGLU
+
+_DENSE_TESTS = (
+    (16, 1024, 512),
+    (8, 2048, 512),
+)
+
+
+def construct_dmoe_glu(
+        hidden_size: int,
+        ffn_hidden_size: int,
+        mlp_impl: str ='sparse',
+        memory_optimized_mlp: bool =False):
+    init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
+    args = Arguments(
+        hidden_size=hidden_size,
+        ffn_hidden_size=ffn_hidden_size,
+        moe_num_experts=1,
+        moe_top_k=1,
+        init_method=init_method,
+        memory_optimized_mlp=memory_optimized_mlp,
+        mlp_type='glu',
+        mlp_impl=mlp_impl,
+        fp16=False,
+        bf16=True)
+
+    glu = testing.GLU(args)
+    dmoe_glu = dmlp_registry.get(args)
+
+    dmoe_glu.cuda(torch.cuda.current_device()).to(torch.bfloat16)
+    glu.cuda(torch.cuda.current_device()).to(torch.bfloat16)
+
+    with torch.no_grad():
+        glu.w1.copy_(dmoe_glu.w1.T)
+        glu.v1.copy_(dmoe_glu.v1.T)
+        glu.w2.copy_(dmoe_glu.w2)
+
+    return args, glu, dmoe_glu
+
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs'), _DENSE_TESTS)
+def test_glu_forward_grouped_mlp(bs: int, sl: int, hs: int):
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+
+    _, glu, dmoe_glu = construct_dmoe_glu(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2,
+        mlp_impl='grouped')
+
+    expected_out = glu(x)
+    tokens_per_expert = torch.tensor([bs * sl]).cuda()
+    out = dmoe_glu(x.view(bs * sl, hs), tokens_per_expert)
+    out = out.view(sl, bs, hs)
+
+    assert out.shape == x.shape == expected_out.shape
+    assert torch.allclose(out, expected_out)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs'), _DENSE_TESTS)
+def test_glu_forward_grouped_mlp_mem_opt(bs: int, sl: int, hs: int):
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+
+    _, glu, dmoe_glu = construct_dmoe_glu(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2,
+        mlp_impl='grouped',
+        memory_optimized_mlp=True)
+
+    expected_out = glu(x)
+    tokens_per_expert = torch.tensor([bs * sl]).cuda()
+    out = dmoe_glu(x.view(bs * sl, hs), tokens_per_expert)
+    out = out.view(sl, bs, hs)
+
+    assert out.shape == x.shape == expected_out.shape
+    assert torch.allclose(out, expected_out)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs'), _DENSE_TESTS)
+def test_glu_forward_sparse_mlp(bs: int, sl: int, hs: int):
+    x = torch.randn(sl, bs, hs).to(torch.bfloat16).cuda()
+
+    _, glu, dmoe_glu = construct_dmoe_glu(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2,
+        mlp_impl='sparse')
+
+    expected_out = glu(x)
+    with torch.no_grad():
+        topo = stk.random.mask(bs * sl, hs * 2, 0, blocking=128).cuda()
+    out = dmoe_glu(x.view(bs * sl, hs), topo)
+    out = out.view(sl, bs, hs)
+
+    assert out.shape == x.shape == expected_out.shape
+    assert torch.allclose(out, expected_out)
diff --git a/tests/layers/moe_test.py b/tests/layers/moe_test.py
new file mode 100644
index 00000000..75ea196f
--- /dev/null
+++ b/tests/layers/moe_test.py
@@ -0,0 +1,148 @@
+from functools import partial
+
+import pytest
+import torch
+
+from megablocks.layers import moe, testing
+from megablocks.layers.arguments import Arguments
+
+_FORWARD_TESTS = (
+    (16, 1024, 512, 1, 1),
+    (16, 1024, 512, 2, 1),
+    (16, 1024, 512, 4, 1),
+    (16, 1024, 512, 8, 1),
+    (8, 2048, 512, 1, 1),
+    (8, 2048, 512, 2, 1),
+    (8, 2048, 512, 4, 1),
+    (16, 1024, 512, 2, 2),
+    (16, 1024, 512, 4, 2),
+    (16, 1024, 512, 4, 4),
+    (16, 1024, 512, 8, 2),
+    (16, 1024, 512, 8, 4),
+    (16, 1024, 512, 8, 8),
+)
+
+
+_DENSE_TESTS = (
+    (16, 1024, 512),
+    (8, 2048, 512),
+)
+
+
+def construct_moe(
+        hidden_size,
+        ffn_hidden_size,
+        moe_num_experts=1,
+        moe_capacity_factor=1,
+        moe_top_k=1):
+    init_method = partial(torch.nn.init.normal_, mean=0.0, std=0.1)
+    args = Arguments(
+        hidden_size=hidden_size,
+        ffn_hidden_size=ffn_hidden_size,
+        moe_num_experts=moe_num_experts,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_top_k=moe_top_k,
+        init_method=init_method)
+
+    mlp = testing.FFN(args)
+    moe_mlp = moe.MoE(args)
+
+    mlp.cuda(torch.cuda.current_device()).half()
+    moe_mlp.cuda(torch.cuda.current_device()).half()
+
+    # Set the baseline parameters to match exactly.
+    if moe_num_experts == 1:
+        with torch.no_grad():
+            mlp.w1.copy_(moe_mlp.experts.mlp.w1.squeeze())
+            mlp.w2.copy_(moe_mlp.experts.mlp.w2.squeeze())
+    return args, mlp, moe_mlp
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs', 'num_experts', 'top_k'),
+                         _FORWARD_TESTS)
+def test_moe_forward(bs: int, sl: int, hs: int, num_experts: int, top_k: int):
+    x = torch.randn(sl, bs, hs).half().cuda()
+
+    _, _, layer = construct_moe(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2,
+        moe_num_experts=num_experts,
+        moe_top_k=top_k)
+
+    out, _ = layer(x)
+    assert out.shape == x.shape
+    moe.clear_load_balancing_loss()
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs', 'num_experts', 'top_k'),
+                         _FORWARD_TESTS)
+def test_moe_forward_backward(bs: int, sl: int, hs: int, num_experts: int, top_k: int):
+    x = torch.randn(sl, bs, hs).half().cuda()
+    x.requires_grad_(True)
+
+    args, _, layer = construct_moe(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2,
+        moe_num_experts=num_experts,
+        moe_top_k=top_k)
+
+    out, _ = layer(x)
+    assert out.shape == x.shape
+
+    loss = out.sum() + moe.batched_load_balancing_loss(args)
+    loss.backward()
+    layer.zero_grad(set_to_none=True)
+    x.grad = None
+    moe.clear_load_balancing_loss()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs'), _DENSE_TESTS)
+def test_moe_forward_vs_dense(bs: int, sl: int, hs: int):
+    x = torch.randn(sl, bs, hs).half().cuda()
+
+    _, mlp, moe_mlp = construct_moe(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2)
+
+    expected_out = mlp(x)
+    out, _ = moe_mlp(x)
+    assert out.shape == x.shape == expected_out.shape
+    assert torch.allclose(out, expected_out)
+    moe.clear_load_balancing_loss()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('bs', 'sl', 'hs'), _DENSE_TESTS)
+def test_moe_forward_backward_vs_dense(bs: int, sl: int, hs: int):
+    x = torch.randn(sl, bs, hs).half().cuda()
+    x.requires_grad_(True)
+
+    _, mlp, moe_mlp = construct_moe(
+        hidden_size=hs,
+        ffn_hidden_size=hs * 2)
+
+    out, _ = moe_mlp(x)
+    loss = out.sum()
+    loss.backward()
+    w1_grad = moe_mlp.experts.mlp.w1.grad.detach().squeeze()
+    w2_grad = moe_mlp.experts.mlp.w2.grad.detach().squeeze()
+    moe_mlp.zero_grad(set_to_none=True)
+    x.grad = None
+    moe.clear_load_balancing_loss()
+
+    expected_out = mlp(x)
+    expected_loss = expected_out.sum()
+    expected_loss.backward()
+    expected_w1_grad = mlp.w1.grad.detach()
+    expected_w2_grad = mlp.w2.grad.detach()
+    mlp.zero_grad(set_to_none=True)
+    x.grad = None
+
+    # Verify the gradients match.
+    assert w1_grad.shape ==  expected_w1_grad.shape
+    assert w2_grad.shape == expected_w2_grad.shape
+    assert torch.allclose(w1_grad, expected_w1_grad)
+    assert torch.allclose(w2_grad, expected_w2_grad)
+    moe.clear_load_balancing_loss()
diff --git a/megablocks/layers/parallelism_test.py b/tests/layers/parallelism_test.py
similarity index 79%
rename from megablocks/layers/parallelism_test.py
rename to tests/layers/parallelism_test.py
index aa3d6584..0aa42694 100644
--- a/megablocks/layers/parallelism_test.py
+++ b/tests/layers/parallelism_test.py
@@ -1,11 +1,12 @@
 import functools
 
-from megablocks.layers import dmoe, arguments, mpu
-from megablocks import benchmark_util
 import numpy as np
+import pytest
 import torch
 
-_TESTS = (
+from megablocks.layers import arguments, dmoe, mpu
+
+_PARALLELISM_TESTS = (
     (64, 1024, 512, 2048, 64, 1, False),
     (64, 1024, 512, 2048, 64, 1, True),
     # Test with fewer experts than ranks to verify tensor
@@ -14,15 +15,26 @@
     (4, 1, 512, 2048, 4, 1, True),
 )
 
+# Todo: Fix this long term
+@pytest.fixture
+def group():
+    return None
+
+
+@pytest.mark.world_size(2)
+@pytest.mark.gpu
+@pytest.mark.parametrize(('batch_size', 'sequence_length', 'hidden_size', 'ffn_hidden_size', 'num_experts', 'top_k', 'memory_optimized'),
+                         _PARALLELISM_TESTS)
 def test_expert_parallel_versus_weight_parallel(
         group,
-        batch_size,
-        sequence_length,
-        hidden_size,
-        ffn_hidden_size,
-        num_experts,
-        top_k,
-        memory_optimized):
+        batch_size: int,
+        sequence_length: int,
+        hidden_size: int,
+        ffn_hidden_size: int,
+        num_experts: int,
+        top_k: int,
+        memory_optimized: bool):
+
     init_fn = functools.partial(torch.nn.init.normal_, mean=0.0, std=0.1)
     ep_args = arguments.Arguments(
         hidden_size=hidden_size,
@@ -71,10 +83,10 @@ def test_expert_parallel_versus_weight_parallel(
     for i in range(torch.distributed.get_world_size(group)):
         torch.distributed.barrier(group)
         if i == rank:
-            np.testing.assert_allclose(
+            assert np.testing.assert_allclose(
                 out.detach().float().cpu(),
                 expected_out.detach().float().cpu(),
-                rtol=1e-4, atol=1e-4)
+                rtol=1e-4, atol=1e-4) is None
 
     # Test backward.
     out.mean().backward()
@@ -99,36 +111,24 @@ def permute(x):
     wp_w2_grad = gather(wp.experts.mlp.w2.grad)
     ep_w2_grad = permute(gather(ep.experts.mlp.w2.grad))
     if rank == 0:
-        np.testing.assert_allclose(
+        assert np.testing.assert_allclose(
             wp_w2_grad.float().cpu(),
             ep_w2_grad.float().cpu(),
-            rtol=1e-5, atol=1e-5)
+            rtol=1e-5, atol=1e-5) is None
 
     wp_w1_grad = gather(wp.experts.mlp.w1.grad)
     ep_w1_grad = permute(gather(ep.experts.mlp.w1.grad))
     if rank == 0:
-        np.testing.assert_allclose(
+        assert np.testing.assert_allclose(
             wp_w1_grad.float().cpu(),
             ep_w1_grad.float().cpu(),
-            rtol=1e-5, atol=1e-5)
+            rtol=1e-5, atol=1e-5) is None
 
     # Verify the router weight gradient, which is not sharded.
     for i in range(torch.distributed.get_world_size(group)):
         torch.distributed.barrier(group)
         if i == rank:
-            np.testing.assert_allclose(
+            assert np.testing.assert_allclose(
                 wp.router.layer.weight.grad.float().cpu(),
                 ep.router.layer.weight.grad.float().cpu(),
-                rtol=1e-5, atol=1e-5)
-
-
-if __name__ == '__main__':
-    assert torch.distributed.is_available()
-    group = torch.distributed.init_process_group(backend='nccl')
-    local_rank = torch.distributed.get_rank(group)
-    torch.cuda.set_device(local_rank)
-
-    for args in _TESTS:
-        if local_rank == 0:
-            print(f"TEST: {args}")
-        test_expert_parallel_versus_weight_parallel(group, *args)
+                rtol=1e-5, atol=1e-5) is None
diff --git a/tests/ops/binned_gather_test.py b/tests/ops/binned_gather_test.py
new file mode 100644
index 00000000..cc59ae3c
--- /dev/null
+++ b/tests/ops/binned_gather_test.py
@@ -0,0 +1,66 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+from megablocks import ops
+
+BINNED_GATHER_TESTS = (
+    (4, 2, 2, 1),
+    (4, 2, 2, 2),
+    (4, 2, 2, 4),
+    (1024, 1536, 4, 1),
+    (1024, 1536, 4, 2),
+    (1024, 1536, 4, 4),
+    (1024, 1536, 64, 1),
+    (1024, 1536, 64, 2),
+    (1024, 1536, 64, 4),
+    (1024, 1536, 128, 1),
+    (1024, 1536, 128, 2),
+    (1024, 1536, 128, 4),
+    (16384, 768, 4, 1),
+    (16384, 768, 4, 2),
+    (16384, 768, 4, 4),
+    (16384, 768, 64, 1),
+    (16384, 768, 64, 2),
+    (16384, 768, 64, 4),
+    (16384, 768, 128, 1),
+    (16384, 768, 128, 2),
+    (16384, 768, 128, 4),
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('sl', 'hs', 'ne', 'top_k'), BINNED_GATHER_TESTS)
+def test_binned_gather(sl: int, hs: int, ne: int, top_k: int):
+    # NOTE: Capacity factor == 1.
+    ec = (sl * top_k) // ne
+
+    # Create the data and indices.
+    x = torch.randn((sl, hs)).cuda().half()
+
+    # Randomly assign tokens to experts.
+    top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+    _, indices = ops.sort(top_expert)
+    bins = ops.inclusive_cumsum(ops.histogram(top_expert, ne), 0)
+
+    def binned_gather(x: torch.Tensor, indices: torch.Tensor,
+                      bins: torch.Tensor, ec: int, top_k: int):
+        x = x.cpu().numpy()
+        indices = indices.cpu().numpy()
+        bins = bins.cpu().numpy()
+        start = 0
+        out = np.zeros((ne, ec, hs))
+        for i in range(ne):
+            end = bins[i]
+            for j in range(min(ec, end - start)):
+                index = indices[start + j] // top_k
+                out[i, j, :] = x[index, :]
+            start = end
+        return torch.from_numpy(out).cuda().half()
+
+    out = ops.binned_gather(x, indices, bins, ec, top_k)
+    expected_out = binned_gather(x, indices, bins, ec, top_k)
+    assert torch.all(torch.eq(out, expected_out))
diff --git a/tests/ops/binned_scatter_test.py b/tests/ops/binned_scatter_test.py
new file mode 100644
index 00000000..2d1c5856
--- /dev/null
+++ b/tests/ops/binned_scatter_test.py
@@ -0,0 +1,75 @@
+import numpy as np
+import pytest
+import torch
+
+from megablocks import ops
+
+_BINNED_SCATTER_TESTS = (
+    (4, 2, 2, 1),
+    (4, 2, 2, 2),
+    (4, 2, 2, 4),
+    (1024, 1536, 4, 1),
+    (1024, 1536, 4, 2),
+    (1024, 1536, 4, 4),
+    (1024, 1536, 64, 1),
+    (1024, 1536, 64, 2),
+    (1024, 1536, 64, 4),
+    (1024, 1536, 128, 1),
+    (1024, 1536, 128, 2),
+    (1024, 1536, 128, 4),
+    (16384, 768, 4, 1),
+    (16384, 768, 4, 2),
+    (16384, 768, 4, 4),
+    (16384, 768, 64, 1),
+    (16384, 768, 64, 2),
+    (16384, 768, 64, 4),
+    (16384, 768, 128, 1),
+    (16384, 768, 128, 2),
+    (16384, 768, 128, 4),
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('sl', 'hs', 'ne', 'top_k'), _BINNED_SCATTER_TESTS)
+def testBinnedScatter(sl: int, hs: int, ne: int, top_k: int):
+    # NOTE: Capacity factor == 1.
+    ec = (sl * top_k) // ne
+
+    # Create the data and indices.
+    x = torch.randn((sl, hs)).cuda().half()
+
+    # Randomly assign tokens to experts.
+    top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+    _, indices = ops.sort(top_expert)
+    bins = ops.inclusive_cumsum(ops.histogram(top_expert, ne), 0)
+
+    # Sample weights for the scatter reduce.
+    weights = torch.rand((sl * top_k,)).cuda().half()
+
+    x = ops.binned_gather(x, indices, bins, ec, top_k)
+
+    def binned_scatter(x: torch.Tensor, indices: torch.Tensor,
+                      weights: torch.Tensor, bins: torch.Tensor, top_k: int):
+        x = x.cpu().numpy()
+        indices = indices.cpu().numpy()
+        weights = weights.cpu().numpy()
+        bins = bins.cpu().numpy()
+        start = 0
+        out = np.zeros((sl, hs))
+        for i in range(ne):
+            end = bins[i]
+            for j in range(min(ec, end - start)):
+                index = indices[start + j]
+                scale = weights[index]
+                index //= top_k
+
+                out[index, :] += scale * x[i, j, :]
+            start = end
+        return torch.from_numpy(out).cuda().half()
+    out = ops.binned_scatter(x, indices, weights, bins, top_k)
+    expected_out = binned_scatter(x, indices, weights, bins, top_k)
+
+    # NOTE: We need to check approximate equality because the
+    # scatter reduce uses atomics.
+    assert np.testing.assert_allclose(
+        out.cpu(), expected_out.cpu(), rtol=5e-3) is None
diff --git a/tests/ops/cumsum_test.py b/tests/ops/cumsum_test.py
new file mode 100644
index 00000000..a1b71605
--- /dev/null
+++ b/tests/ops/cumsum_test.py
@@ -0,0 +1,44 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from megablocks import ops
+
+CUMSUM_TESTS = (
+    (1, 32),
+    (2, 32),
+    (2, 1024),
+    (4, 1024),
+    (8, 1024),
+    (16, 1024),
+    (32, 1024),
+    (64, 1024),
+    (128, 1024),
+    (2, 16384),
+    (4, 16384),
+    (8, 16384),
+    (16, 16384),
+    (32, 16384),
+    (64, 16384),
+    (128, 16384),
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('n', 'm'), CUMSUM_TESTS)
+def test_exclusive_cumsum(n: int, m: int):
+    x = torch.randint(0, 2, (n, m)).long().cuda()
+    out = ops.exclusive_cumsum(x, 1) * x
+    expected_out = (torch.cumsum(x, dim=1) - 1) * x
+    assert torch.all(torch.eq(out, expected_out))
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('n', 'm'), CUMSUM_TESTS)
+def test_inclusive_cumsum(n: int, m: int):
+    x = torch.randint(0, 2, (n, m)).long().cuda()
+    out = ops.inclusive_cumsum(x, 1)
+    expected_out = torch.cumsum(x, dim=1)
+    assert torch.all(torch.eq(out, expected_out))
diff --git a/megablocks/ops/histogram_test.py b/tests/ops/histogram_test.py
similarity index 72%
rename from megablocks/ops/histogram_test.py
rename to tests/ops/histogram_test.py
index b188cc7e..25b30cb9 100644
--- a/megablocks/ops/histogram_test.py
+++ b/tests/ops/histogram_test.py
@@ -1,9 +1,10 @@
-import unittest
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
 
-from absl.testing import parameterized
-from megablocks import ops
+import pytest
 import torch
 
+from megablocks import ops
 
 _HISTOGRAM_TESTS = (
     (1, 32, torch.int16, 128),
@@ -63,19 +64,20 @@
 )
 
 
-class HistogramTest(parameterized.TestCase):
-
-    @parameterized.parameters(*_HISTOGRAM_TESTS)
-    def testHistogram(self, m, n, dtype, max_val):
-        x = torch.randint(0, max_val, (m, n)).cuda().to(dtype)
+# Override the seed_all fixture in autouse.py because
+# _histc_cuda does not have a deterministic implementation
+@pytest.fixture()
+def seed_all():
+    torch.use_deterministic_algorithms(False)
+    return
 
-        out = ops.histogram(x, max_val)
-        expected_out = torch.stack(
-            [torch.histc(y, max_val, 0, max_val - 1)
-             for y in torch.split(x, 1)]
-        )
-        self.assertTrue(torch.all(torch.eq(out, expected_out)))
 
+@pytest.mark.gpu
+@pytest.mark.parametrize(('m', 'n', 'dtype', 'max_val'), _HISTOGRAM_TESTS)
+def test_histogram(m: int, n: int, dtype: torch.dtype, max_val: int):
+    x = torch.randint(0, max_val, (m, n)).cuda().to(dtype)
 
-if __name__ == '__main__':
-    unittest.main()
+    out = ops.histogram(x, max_val)
+    expected_out = torch.stack(
+        [torch.histc(y, max_val, 0, max_val - 1) for y in torch.split(x, 1)])
+    assert torch.all(torch.eq(out, expected_out))
diff --git a/tests/ops/padded_gather_test.py b/tests/ops/padded_gather_test.py
new file mode 100644
index 00000000..e6eb7f75
--- /dev/null
+++ b/tests/ops/padded_gather_test.py
@@ -0,0 +1,89 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+from megablocks import ops
+
+PADDED_GATHER_TESTS = (
+    (4, 2, 2, 1),
+    (4, 2, 2, 2),
+    (1024, 1, 4, 1),
+    (1024, 1, 4, 2),
+    (1024, 1, 4, 4),
+    (1024, 1, 64, 1),
+    (1024, 1, 64, 2),
+    (1024, 1, 64, 4),
+    (1024, 1, 128, 1),
+    (1024, 1, 128, 2),
+    (1024, 1, 128, 4),
+    (1024, 1536, 4, 1),
+    (1024, 1536, 4, 2),
+    (1024, 1536, 4, 4),
+    (1024, 1536, 64, 1),
+    (1024, 1536, 64, 2),
+    (1024, 1536, 64, 4),
+    (1024, 1536, 128, 1),
+    (1024, 1536, 128, 2),
+    (1024, 1536, 128, 4),
+    (16384, 768, 4, 1),
+    (16384, 768, 4, 2),
+    (16384, 768, 4, 4),
+    (16384, 768, 64, 1),
+    (16384, 768, 64, 2),
+    (16384, 768, 64, 4),
+    (16384, 768, 128, 1),
+    (16384, 768, 128, 2),
+    (16384, 768, 128, 4),
+    (16384, 1, 4, 1),
+    (16384, 1, 4, 2),
+    (16384, 1, 4, 4),
+    (16384, 1, 64, 1),
+    (16384, 1, 64, 2),
+    (16384, 1, 64, 4),
+    (16384, 1, 128, 1),
+    (16384, 1, 128, 2),
+    (16384, 1, 128, 4),
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('sl', 'hs', 'ne', 'top_k'), PADDED_GATHER_TESTS)
+def testPaddedGather(sl: int, hs: int, ne: int, top_k: int):
+    # Create the data and indices.
+    x = torch.randn((sl, hs)).cuda().half()
+
+    # Randomly assign tokens to experts.
+    top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+    bin_ids, indices = ops.sort(top_expert)
+    tokens_per_expert = ops.histogram(top_expert, ne)
+    padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+    padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+    bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+
+    def padded_gather(x: torch.Tensor, indices: torch.Tensor,
+                      bin_ids: torch.Tensor, bins: torch.Tensor,
+                      padded_bins: torch.Tensor, top_k: int):
+        x = x.cpu().numpy()
+        indices = indices.cpu().numpy()
+        bin_ids = bin_ids.cpu().numpy()
+        bins = bins.cpu().numpy()
+        padded_bins = padded_bins.cpu().numpy()
+
+        out = np.zeros((padded_bins[-1], hs))
+        in_idx = 0
+        for i, end in enumerate(bins):
+            out_idx = 0 if i == 0 else padded_bins[i - 1]
+            end = bins[i]
+            while in_idx < end:
+                load_idx = indices[in_idx] // top_k
+                out[out_idx, :] = x[load_idx, :]
+                in_idx += 1
+                out_idx += 1
+        return torch.from_numpy(out).cuda().half()
+
+    out = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+    expected_out = padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+    assert torch.all(torch.eq(out, expected_out))
diff --git a/tests/ops/padded_scatter_test.py b/tests/ops/padded_scatter_test.py
new file mode 100644
index 00000000..ebd04a86
--- /dev/null
+++ b/tests/ops/padded_scatter_test.py
@@ -0,0 +1,136 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+from megablocks import ops
+
+PADDED_SCATTER_TESTS = [
+    (4, 2, 2, 2),
+    (4, 2, 2, 1),
+    (4, 2, 2, 1),
+    (4, 2, 2, 1),
+    (4, 2, 2, 2),
+    (4, 2, 2, 2),
+    (1024, 1, 4, 1),
+    (1024, 1, 4, 2),
+    (1024, 1, 4, 4),
+    (1024, 1, 4, 1),
+    (1024, 1, 4, 2),
+    (1024, 1, 4, 4),
+    (1024, 1, 4, 1),
+    (1024, 1, 4, 2),
+    (1024, 1, 4, 4),
+    (1024, 1, 64, 1),
+    (1024, 1, 64, 2),
+    (1024, 1, 64, 4),
+    (1024, 1, 128, 1),
+    (1024, 1, 128, 2),
+    (1024, 1, 128, 4),
+    (1024, 1536, 4, 1),
+    (1024, 1536, 4, 2),
+    (1024, 1536, 4, 4),
+    (1024, 1536, 4, 4),
+    (1024, 1536, 4, 4),
+    (1024, 1536, 64, 1),
+    (1024, 1536, 64, 2),
+    (1024, 1536, 64, 4),
+    (1024, 1536, 128, 1),
+    (1024, 1536, 128, 2),
+    (1024, 1536, 128, 4),
+    (1024, 1536, 128, 1),
+    (1024, 1536, 128, 1),
+    (16384, 768, 4, 1),
+    (16384, 768, 4, 2),
+    (16384, 768, 4, 4),
+    (16384, 768, 64, 1),
+    (16384, 768, 64, 2),
+    (16384, 768, 64, 4),
+    (16384, 768, 128, 1),
+    (16384, 768, 128, 2),
+    (16384, 768, 128, 4),
+    (16384, 1, 4, 1),
+    (16384, 1, 4, 2),
+    (16384, 1, 4, 4),
+    (16384, 1, 64, 1),
+    (16384, 1, 64, 2),
+    (16384, 1, 64, 4),
+    (16384, 1, 128, 1),
+    (16384, 1, 128, 2),
+    (16384, 1, 128, 4),
+    (16384, 1, 128, 2),
+    (16384, 1, 128, 2),
+]
+
+
+def _to_numpy(x: torch.Tensor) -> np.ndarray:
+    return x.detach().cpu().numpy()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize((
+    'sl',
+    'hs',
+    'ne',
+    'top_k',
+), PADDED_SCATTER_TESTS)
+def testPaddedScatter(sl: int, hs: int, ne: int, top_k: int):
+    # Create the data and indices.
+    x = torch.randn((sl, hs), requires_grad=True).cuda().half()
+
+    # Randomly assign tokens to experts.
+    top_expert = torch.randint(0, ne, (sl * top_k,)).cuda().int()
+    bin_ids, indices = ops.sort(top_expert)
+    tokens_per_expert = ops.histogram(top_expert, ne)
+    padded_tokens_per_expert = ops.round_up(tokens_per_expert, 128)
+    padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+    bins = ops.inclusive_cumsum(tokens_per_expert, 0)
+
+    # Sample weights for the scatter reduce.
+    weights = torch.rand((sl * top_k,), requires_grad=True).cuda().half()
+
+    # Gather the data to prepare for backwards.
+    x = ops.padded_gather(x, indices, bin_ids, bins, padded_bins, top_k)
+
+    def padded_scatter(x: torch.Tensor, indices: torch.Tensor,
+                       bin_ids: torch.Tensor, weights: torch.Tensor,
+                       bins: torch.Tensor, padded_bins: torch.Tensor,
+                       top_k: int):
+        x = x.detach().cpu().numpy()
+        indices: np.ndarray = _to_numpy(indices)
+        bin_ids: np.ndarray = _to_numpy(bin_ids)
+        weights: np.ndarray = _to_numpy(weights)
+        bins: np.ndarray = _to_numpy(bins)
+        padded_bins: np.ndarray = _to_numpy(padded_bins)
+
+        out = np.zeros((indices.shape[0] // top_k, hs))
+        out_idx = 0
+        for i in range(len(bins)):
+            in_idx = 0 if i == 0 else padded_bins[i - 1]
+            end = bins[i]
+            while out_idx < end:
+                store_idx = indices[out_idx]
+                scale = weights[store_idx]
+                store_idx //= top_k
+
+                out[store_idx, :] += scale * x[in_idx, :]
+                out_idx += 1
+                in_idx += 1
+        return torch.from_numpy(out).cuda().half()
+
+    out = ops.padded_scatter(x, indices, bin_ids, weights, bins, padded_bins,
+                             top_k)
+    expected_out = padded_scatter(x, indices, bin_ids, weights, bins,
+                                  padded_bins, top_k)
+
+    out.backward(torch.randn_like(out))  # sanity check backward pass
+
+    # NOTE: We need to check approximate equality because the scatter reduce uses atomics.
+    # np.testing.assert_allclose returns `None` if no error and raises an AssertionError if an error exists
+    assert np.testing.assert_allclose(
+        _to_numpy(out),
+        _to_numpy(expected_out),
+        rtol=5e-3,
+    ) is None
diff --git a/tests/ops/replicate_test.py b/tests/ops/replicate_test.py
new file mode 100644
index 00000000..94aeb67a
--- /dev/null
+++ b/tests/ops/replicate_test.py
@@ -0,0 +1,108 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+try:
+    import megablocks_ops as backend  # type: ignore
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError("No module named 'megablocks_ops'.") from e
+
+from megablocks import ops
+
+
+def promote_scalar(x: torch.Tensor) -> torch.Tensor:
+    return x.view(1) if not len(x.size()) else x
+
+
+REPLICATE_TESTS = [
+    (8, 1, 1),
+    (8, 2, 1),
+    (8, 4, 1),
+    (8, 8, 1),
+    (8, 2, 2),
+    (8, 4, 2),
+    (8, 8, 2),
+    (8, 2, 4),
+    (8, 4, 4),
+    (8, 8, 4),
+    (8, 2, 8),
+    (8, 4, 8),
+    (8, 8, 8),
+    (16384, 2, 1),
+    (16384, 4, 1),
+    (16384, 8, 1),
+    (16384, 16, 1),
+    (16384, 32, 1),
+    (16384, 64, 1),
+    (16384, 128, 1),
+    (16384, 2, 2),
+    (16384, 4, 2),
+    (16384, 8, 2),
+    (16384, 16, 2),
+    (16384, 32, 2),
+    (16384, 64, 2),
+    (16384, 128, 2),
+    (16384, 2, 4),
+    (16384, 4, 4),
+    (16384, 8, 4),
+    (16384, 16, 4),
+    (16384, 32, 4),
+    (16384, 64, 4),
+    (16384, 128, 4),
+    (16384, 2, 8),
+    (16384, 4, 8),
+    (16384, 8, 8),
+    (16384, 16, 8),
+    (16384, 32, 8),
+    (16384, 64, 8),
+    (16384, 128, 8),
+]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('tokens', 'num_centers', 'top_k'), REPLICATE_TESTS)
+def test_replicate(tokens: int, num_centers: int, top_k: int):
+    tokens_to_centers = torch.randint(0, num_centers, (tokens,)).cuda().int()
+    tokens_per_center = ops.histogram(tokens_to_centers, num_centers)
+    bins = ops.inclusive_cumsum(tokens_per_center, 0)
+    bins = promote_scalar(bins)
+    center_weights = torch.randn(top_k, num_centers).cuda().half()
+
+    def replicate(x: torch.Tensor, bins: torch.Tensor, num_outputs: int):
+        x = x.cpu().numpy()
+        bins = bins.cpu().numpy()
+        out = np.zeros((x.shape[0], num_outputs))
+        for batch_idx in range(x.shape[0]):
+            start = 0
+            for i, end in enumerate(bins):
+                value = x[batch_idx, i]
+                while start < end:
+                    out[batch_idx, start] = value
+                    start += 1
+        return torch.from_numpy(out).cuda().half()
+
+    out = ops.replicate(center_weights, bins, tokens)
+    expected_out = replicate(center_weights, bins, tokens)
+    assert torch.all(torch.eq(out, expected_out))
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('tokens', 'num_centers', 'top_k'), REPLICATE_TESTS)
+def test_replicate_backward(tokens: int, num_centers: int, top_k: int):
+    tokens_to_centers = torch.randint(0, num_centers, (tokens,)).cuda().int()
+    tokens_per_center = ops.histogram(tokens_to_centers, num_centers)
+    bins = ops.inclusive_cumsum(tokens_per_center, 0)
+    bins = promote_scalar(bins)
+    center_weights = torch.randn(top_k, num_centers).cuda().half()
+
+    grad = ops.replicate(center_weights, bins, tokens)
+
+    out = torch.empty_like(center_weights)
+    backend.replicate_backward(grad, bins, out)
+    expected_out = center_weights * tokens_per_center.view([1, num_centers])
+
+    # NOTE: This floating-point reduction could be a problem for training stability and accuracy.
+    assert torch.allclose(out, expected_out, rtol=1e-2)
diff --git a/tests/ops/sort_test.py b/tests/ops/sort_test.py
new file mode 100644
index 00000000..e07f2e1a
--- /dev/null
+++ b/tests/ops/sort_test.py
@@ -0,0 +1,66 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Dict, Optional, Union
+
+import numpy as np
+import pytest
+import torch
+
+from megablocks import ops
+
+SORT_TESTS = [
+    (32, torch.int16, None),
+    (1024, torch.int16, None),
+    (16384, torch.int16, None),
+    (32, torch.int32, None),
+    (1024, torch.int32, None),
+    (16384, torch.int32, None),
+    (32, torch.int64, None),
+    (1024, torch.int64, None),
+    (16384, torch.int64, None),
+    (32, torch.int16, 128),
+    (1024, torch.int16, 128),
+    (16384, torch.int16, 128),
+    (32, torch.int32, 128),
+    (1024, torch.int32, 128),
+    (16384, torch.int32, 128),
+    (32, torch.int64, 128),
+    (1024, torch.int64, 128),
+    (16384, torch.int64, 128),
+]
+
+
+def torch_to_numpy_dtype(
+        dtype: torch.dtype) -> Union[np.int16, np.int32, np.int64]:
+    types: Dict[torch.dtype, Union[np.int16, np.int32, np.int64]] = {
+        torch.int16: np.int16,
+        torch.int32: np.int32,
+        torch.int64: np.int64
+    }
+    return types[dtype]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    ('n', 'dtype', 'max_val'),
+    SORT_TESTS,
+)
+def test_sort(n: int, dtype: torch.dtype, max_val: Optional[int]):
+    if max_val is None:
+        max_val = np.iinfo(torch_to_numpy_dtype(dtype)).max
+    end_bit = int(np.ceil(np.log2(max_val)))
+    x = torch.randint(0, max_val, (n,)).cuda().to(dtype)
+
+    out, indices = ops.sort(x, end_bit)
+    expected_out, expected_indices = torch.sort(x)
+    assert torch.all(torch.eq(out, expected_out))
+
+    # NOTE: The indices can be in different order depending
+    # on sort stability if multiple values in the array are
+    # equal.
+    data = torch.empty_like(x)
+    data.scatter_(0, indices.long(), out)
+    expected_data = torch.empty_like(x)
+    expected_data.scatter_(0, expected_indices, expected_out)
+    assert torch.all(torch.eq(data, expected_data))
diff --git a/tests/ops/topology_test.py b/tests/ops/topology_test.py
new file mode 100644
index 00000000..a7135be0
--- /dev/null
+++ b/tests/ops/topology_test.py
@@ -0,0 +1,69 @@
+# Copyright 2024 MosaicML MegaBlocks authors
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+from megablocks import ops
+
+TOPOLOGY_TESTS = (
+    (1024, 1536, 2),
+    (1024, 1536, 4),
+    (1024, 1536, 8),
+    (1024, 1536, 16),
+    (1024, 1536, 32),
+    (1024, 1536, 64),
+    (1024, 1536, 128),
+    (1024, 1536, 256),
+    (1024, 1536, 512),
+    (16384, 768, 2),
+    (16384, 768, 4),
+    (16384, 768, 8),
+    (16384, 768, 16),
+    (16384, 768, 32),
+    (16384, 768, 64),
+    (16384, 768, 128),
+    (16384, 768, 256),
+    (16384, 768, 512),
+    (16384, 768, 1024),
+    (8, 14336, 8),
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(('sl', 'hs', 'ne'), TOPOLOGY_TESTS)
+def test_topology(sl: int, hs: int, ne: int):
+    # Create the data and indices.
+    blocking = 128
+    assert hs % blocking == 0
+
+    # Randomly assign tokens to experts.
+    top_expert = torch.randint(0, ne, (sl,)).cuda().int()
+    tokens_per_expert = ops.histogram(top_expert, ne)
+    padded_tokens_per_expert = ops.round_up(tokens_per_expert, blocking)
+    padded_bins = ops.inclusive_cumsum(padded_tokens_per_expert, 0)
+
+    # Dimensions for the output indices.
+    output_block_rows = int(padded_bins[-1]) // blocking
+    output_block_columns = hs // blocking
+
+    def topology(padded_bins: torch.Tensor, blocking: torch.Tensor, rows: int,
+                 columns: int):
+        padded_bins = padded_bins.cpu().numpy()
+
+        out = np.zeros([rows * columns])
+        start = 0
+        for i in range(padded_bins.shape[0]):
+            end = padded_bins[i] // blocking
+            while start < end:
+                for j in range(columns):
+                    out[start * columns + j] = j + i * columns
+                start += 1
+        return torch.from_numpy(out).cuda().short()
+
+    out = ops.topology(padded_bins, blocking, output_block_rows,
+                       output_block_columns)
+    expected_out = topology(padded_bins, blocking, output_block_rows,
+                            output_block_columns)
+    assert torch.all(torch.eq(out, expected_out))