Removing ao from pyproject.toml

.github/workflows/build_docs.yaml

@@ -39,7 +39,7 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install dependencies
         run: |
-          python -m pip install torch torchvision
+          python -m pip install torch torchvision torchao
           python -m pip install -e .
           cd docs
           python -m pip install -r requirements.txt

.github/workflows/build_linux_wheels.yaml

@@ -41,3 +41,4 @@ jobs:
       trigger-event: ${{ github.event_name }}
       build-platform: 'python-build-package'
       pip-install-torch-extra-args: torchvision
+      pip-install-torch-extra-args: torchao

.github/workflows/gpu_test.yaml

@@ -39,10 +39,10 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install torch nightly
         if: ${{ matrix.torch-version == 'nightly' }}
-        run: python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu118
+        run: python -m pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu118
       - name: Install torch stable
         if: ${{ matrix.torch-version == 'stable' }}
-        run: python -m pip install torch torchvision
+        run: python -m pip install torch torchvision torchao
       - name: Install remaining dependencies
         run: |
           python -m pip install -e ".[dev]"

.github/workflows/recipe_test.yaml

@@ -39,7 +39,7 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install dependencies
         run: |
-          python -m pip install torch torchvision
+          python -m pip install torch torchvision torchao
           python -m pip install -e ".[dev]"
           python -m pip install lm-eval==0.4.*
       - name: Run recipe tests with coverage

.github/workflows/recipe_test_nightly.yaml

@@ -39,17 +39,14 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install torch nightly
         if: ${{ matrix.torch-version == 'nightly' }}
-        run: python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
+        run: python -m pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu121
       - name: Install torch stable
         if: ${{ matrix.torch-version == 'stable' }}
-        run: python -m pip install torch torchvision
+        run: python -m pip install torch torchvision torchao
       - name: Install remaining dependencies
         run: |
           python -m pip install -e ".[dev]"
           python -m pip install lm-eval==0.4.*
-      - name: Install torchao nightly
-        if: ${{ matrix.torch-version == 'nightly' }}
-        run: pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121
       - name: Run recipe tests with coverage
         run: pytest tests -m integration_test --cov=. --cov-report=xml --durations=20 -vv
       - name: Upload Coverage to Codecov

.github/workflows/regression_test.yaml

@@ -48,17 +48,14 @@ jobs:
           python3 -m pip install awscli==1.32.6
       - name: Install torch nightly
         if: ${{ matrix.torch-version == 'nightly' }}
-        run: python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu118
+        run: python -m pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu118
       - name: Install torch stable
         if: ${{ matrix.torch-version == 'stable' }}
-        run: python -m pip install torch torchvision
+        run: python -m pip install torch torchvision torchao
       - name: Install remaining dependencies
         run: |
           python -m pip install -e ".[dev]"
           python -m pip install lm-eval==0.4.*
-      - name: Install torchao nightly
-        if: ${{ matrix.torch-version == 'nightly' }}
-        run: pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121
       - name: Run regression tests with coverage
         run: pytest tests -m slow_integration_test --silence-s3-logs --cov=. --cov-report=xml --durations=20 -vv
       - name: Upload Coverage to Codecov

.github/workflows/unit_test.yaml

@@ -33,7 +33,7 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install dependencies
         run: |
-          python -m pip install torch torchvision
+          python -m pip install torch torchvision torchao
           python -m pip install -e ".[dev]"
       - name: Run unit tests with coverage
         run: pytest tests --cov=. --cov-report=xml --durations=20 -vv

README.md

@@ -156,14 +156,15 @@ You can find a full list of all our Llama3 configs [here](recipes/configs/llama3
 
 ## Installation
 
-**Step 1:** [Install PyTorch](https://pytorch.org/get-started/locally/). torchtune is tested with the latest stable PyTorch release as well as the preview nightly version. For fine-tuning the multimodal LLMs available in the repo, you'll need to install torchvision as well.
+**Step 1:** [Install PyTorch](https://pytorch.org/get-started/locally/). torchtune is tested with the latest stable PyTorch release as well as the preview nightly version. torchtune leverages
+torchvision for fine-tuning multimodal LLMs and torchao for the latest in quantization techniques, you should install these as well.
 
 ```
-# Install stable version of PyTorch using pip
-pip install torch torchvision
+# Install stable version of PyTorch libraries using pip
+pip install torch torchvision torchao
 
 # Nightly install for latest features
-pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
+pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu121
 ```
 
 **Step 2:** The latest stable version of torchtune is hosted on PyPI and can be downloaded with the following command:

docs/source/install.rst

@@ -5,7 +5,18 @@ Install Instructions
 ====================
 
 **Pre-requisites**: torchtune requires PyTorch, so please install for your proper host and environment
-using the `Start Locally <https://pytorch.org/get-started/locally/>`_ page.
+using the `Start Locally <https://pytorch.org/get-started/locally/>`_ page. You should also install
+torchvision (for multimodal LLMs) and torchao (for quantization APIs). You can install either stable or
+nightly versions with the following commands:
+
+.. code-block:: bash
+
+    # Install stable version of PyTorch libraries using pip
+    pip install torch torchvision torchao
+
+    # Nightly install for latest features
+    pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu121
+
 
 Install via PyPI
 ----------------

pyproject.toml

@@ -26,8 +26,6 @@ dependencies = [
     "tqdm",
     "omegaconf",
 
-    # Quantization
-    "torchao==0.3.1",
 ]
 dynamic = ["version"]
 

tests/__init__.py

@@ -3,3 +3,20 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+# Check at the top-level that torchao is installed.
+# This is better than doing it at every import site.
+# We have to do this because it is not currently possible to
+# properly support both nightly and stable installs of PyTorch + torchao
+# in pyproject.toml.
+try:
+    import torchao  # noqa
+except ImportError as e:
+    raise ImportError(
+        """
+    torchao not installed.
+    If you are using stable PyTorch please install via pip install torchao.
+    If you are using nightly PyTorch please install via
+    pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121."
+    """
+    ) from e

torchtune/__init__.py

@@ -6,6 +6,24 @@
 
 __version__ = ""
 
+
+# Check at the top-level that torchao is installed.
+# This is better than doing it at every import site.
+# We have to do this because it is not currently possible to
+# properly support both nightly and stable installs of PyTorch + torchao
+# in pyproject.toml.
+try:
+    import torchao  # noqa
+except ImportError as e:
+    raise ImportError(
+        """
+    torchao not installed.
+    If you are using stable PyTorch please install via pip install torchao.
+    If you are using nightly PyTorch please install via
+    pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121."
+    """
+    ) from e
+
 from torchtune import datasets, models, modules, utils
 
 __all__ = [datasets, models, modules, utils]