diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
index 9d227cc25..3aaef6516 100644
--- a/.github/actions/setup-venv/action.yml
+++ b/.github/actions/setup-venv/action.yml
@@ -8,6 +8,10 @@ inputs:
     description: Update this to invalidate the cache
     required: true
     default: v0
+  torch-version:
+    description: The PyTorch version to install
+    required: false
+    default: '<2.2'
 runs:
   using: composite
   steps:
@@ -30,7 +34,7 @@ runs:
       id: virtualenv-cache
       with:
         path: .venv
-        key: ${{ inputs.cache-prefix }}-${{ runner.os }}-${{ env.PYTHON_VERSION }}-${{ hashFiles('requirements.txt', 'dev-requirements.txt', 'hf_olmo/requirements.txt') }}
+        key: ${{ inputs.cache-prefix }}-${{ runner.os }}-${{ env.PYTHON_VERSION }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}
 
     - if: steps.virtualenv-cache.outputs.cache-hit != 'true'
       shell: bash
@@ -38,8 +42,8 @@ runs:
         # Set up virtual environment without cache hit.
         test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv
         . .venv/bin/activate
-        pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
-        pip install -e .[dev]
+        pip install 'torch${{ inputs.torch-version }}' --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install -e .[all]
         pip install -e hf_olmo
 
     - if: steps.virtualenv-cache.outputs.cache-hit == 'true'
@@ -47,7 +51,7 @@ runs:
       run: |
         # Set up virtual environment from cache hit.
         . .venv/bin/activate
-        pip install --no-deps -e .[dev]
+        pip install --no-deps -e .[all]
         pip install --no-deps -e hf_olmo
 
     - shell: bash
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 087b9926b..754a222f1 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -16,7 +16,7 @@ on:
 
 env:
   # Change this to invalidate existing cache.
-  CACHE_PREFIX: v1
+  CACHE_PREFIX: v2
   PYTHONPATH: ./
   TOKENIZERS_PARALLELISM: 'false'
 
@@ -78,10 +78,10 @@ jobs:
         uses: actions/cache@v3
         with:
           path: .mypy_cache
-          key: mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}-${{ github.ref }}-${{ github.sha }}
+          key: mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}-${{ github.ref }}-${{ github.sha }}
           restore-keys: |
-            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}-${{ github.ref }}
-            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}
+            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}-${{ github.ref }}
+            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}
 
       - name: ${{ matrix.task.name }}
         run: |
@@ -172,10 +172,11 @@ jobs:
         with:
           python-version: '3.10'
 
-      - name: Install requirements
-        run: |
-          pip install --upgrade pip setuptools wheel build
-          pip install -r dev-requirements.txt
+      - name: Setup Python environment
+        uses: ./.github/actions/setup-venv
+        with:
+          python-version: '3.10'
+          cache-prefix: ${{ env.CACHE_PREFIX }}
 
       - name: Prepare environment
         run: |
diff --git a/README.md b/README.md
index d84f5a660..0575491db 100644
--- a/README.md
+++ b/README.md
@@ -25,7 +25,7 @@ To install from source (recommended for training/fine-tuning) run:
 ```bash
 git clone https://github.com/allenai/OLMo.git
 cd OLMo
-pip install -e .
+pip install -e .[all]
 ```
 
 Otherwise you can install the model code by itself directly from PyPI with:
diff --git a/dev-requirements.txt b/dev-requirements.txt
deleted file mode 100644
index 3c17627e5..000000000
--- a/dev-requirements.txt
+++ /dev/null
@@ -1,10 +0,0 @@
-ruff
-mypy>=1.0,<1.4
-black>=23.1,<24.0
-isort>=5.12,<5.13
-pytest
-pytest-sphinx
-twine>=1.11.0
-setuptools
-wheel
-build
diff --git a/docker/Dockerfile.gantry b/docker/Dockerfile.gantry
index 18fd894b0..fad65043f 100644
--- a/docker/Dockerfile.gantry
+++ b/docker/Dockerfile.gantry
@@ -8,7 +8,10 @@ FROM olmo-torch2-base
 
 WORKDIR /stage
 
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+COPY pyproject.toml .
+RUN mkdir olmo && touch olmo/__init__.py && \
+    pip install --no-cache-dir .[all] && \
+    pip uninstall -y ai2-olmo && \
+    rm -rf olmo/
 
 WORKDIR /app/olmo
diff --git a/docker/Dockerfile.lumi b/docker/Dockerfile.lumi
index 74d036d00..9bcf84a62 100644
--- a/docker/Dockerfile.lumi
+++ b/docker/Dockerfile.lumi
@@ -88,8 +88,12 @@ RUN cd /opt && \
     sed -i 's/hostname -I/hostname -s/g' /usr/local/lib/python3.10/dist-packages/deepspeed/comm/comm.py
 
 # Install more dependencies
-COPY requirements.txt requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+COPY pyproject.toml .
+RUN mkdir olmo && touch olmo/__init__.py && \
+    pip install --no-cache-dir .[all] && \
+    pip uninstall -y ai2-olmo && \
+    rm -rf olmo/
+
 RUN pip install --no-cache-dir py-spy
 RUN pip install --no-cache-dir wandb --upgrade
 
diff --git a/pyproject.toml b/pyproject.toml
index a1592d8ee..db9af8201 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,14 +4,54 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ai2-olmo"
-dynamic = ["version", "dependencies", "optional-dependencies"]
+dynamic = ["version"]
 readme = "README.md"
 description = "Open Language Model (OLMo)"
 authors = [
     { name = "Allen Institute for Artificial Intelligence", email = "olmo@allenai.org" }
 ]
 requires-python = ">=3.8"
-license = {file = "LICENSE"}
+license = { file = "LICENSE" }
+dependencies = [
+    "numpy",
+    "torch>=2.0,<2.2",
+    "omegaconf",
+    "rich",
+    "boto3",
+    "google-cloud-storage",
+    "tokenizers",
+    "packaging",
+    "cached_path",
+    "transformers",
+]
+
+[project.optional-dependencies]
+dev = [
+    "ruff",
+    "mypy>=1.0,<1.4",
+    "black>=23.1,<24.0",
+    "isort>=5.12,<5.13",
+    "pytest",
+    "pytest-sphinx",
+    "twine>=1.11.0",
+    "setuptools",
+    "wheel",
+    "build",
+]
+train = [
+    "wandb",
+    "beaker-gantry",
+    "click",
+    "torchmetrics",
+    "smashed[remote]>=0.21.1",
+    "safetensors",
+    "datasets",
+    "scikit-learn",
+    "msgspec>=0.14.0",
+]
+all = [
+    "ai2-olmo[dev,train]",
+]
 
 [project.urls]
 Homepage = "https://github.com/allenai/OLMo"
@@ -25,8 +65,6 @@ olmo = ["py.typed"]
 
 [tool.setuptools.dynamic]
 version = { attr = "olmo.version.VERSION" }
-dependencies = { file = ["requirements.txt", "hf_olmo/requirements.txt"] }
-optional-dependencies = { dev = { file = ["dev-requirements.txt"] } }
 
 [tool.setuptools.packages.find]
 include = ["olmo*", "hf_olmo*"]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 9256a2663..000000000
--- a/requirements.txt
+++ /dev/null
@@ -1,19 +0,0 @@
-numpy
-torch>=2.0
-torchmetrics
-tokenizers
-click
-rich
-cached-path
-beaker-gantry
-omegaconf
-wandb
-boto3
-google-cloud-storage
-packaging
-datasets
-scikit-learn
-smashed[remote]>=0.21.1
-msgspec>=0.14.0
-cached_path
-safetensors
\ No newline at end of file
diff --git a/scripts/beaker/beaker_interactive.sh b/scripts/beaker/beaker_interactive.sh
index 0a2f40593..c6456dfe8 100644
--- a/scripts/beaker/beaker_interactive.sh
+++ b/scripts/beaker/beaker_interactive.sh
@@ -49,4 +49,4 @@ gh repo clone allenai/LLM
 cd LLM
 
 # Install other dependencies.
-pip install -e '.[dev]'
+pip install -e '.[all]'
diff --git a/scripts/test_entrypoint.sh b/scripts/test_entrypoint.sh
index 4f289eba1..5dfcaea52 100644
--- a/scripts/test_entrypoint.sh
+++ b/scripts/test_entrypoint.sh
@@ -29,7 +29,7 @@ git checkout --quiet "$COMMIT_SHA"
 
 # Install dependencies.
 pip install --upgrade pip
-pip install --no-cache-dir '.[dev]'
+pip install --no-cache-dir '.[all]'
 
 # Create directory for results.
 mkdir -p /results