Remove all dependencies except torch

.github/workflows/doc_build.yml

@@ -42,8 +42,7 @@ jobs:
           python -m pip install torch
           python -m pip install -e .
           pip install -r dev-requirements.txt
-          cd docs
-          python -m pip install -r requirements.txt
+          python -m pip install -r docs/requirements.txt
       - name: Build docs
         env:
           TORCHAO_VERSION_DOCS: ${{ github.ref }}

.github/workflows/nightly_smoke_test.yml

@@ -2,7 +2,7 @@ name: PyTorch CUDA Nightly Smoke Test
 
 on:
   schedule:
-    # 6 am PST every day 
+    # 6 am PST every day
     - cron: "0 14 * * *"
   workflow_dispatch:
 
@@ -34,7 +34,6 @@ jobs:
       script: |
         python -m pip install --upgrade pip
         pip install ${{ matrix.torch-spec }}
-        pip install -r requirements.txt
         pip install -r dev-requirements.txt
         python setup.py install
         pytest test --verbose -s

.github/workflows/regression_test.yml

@@ -31,7 +31,7 @@ jobs:
             torch-spec: 'torch==2.3.0'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.1"
-          - name: CUDA Nightly 
+          - name: CUDA Nightly
             runs-on: linux.g5.12xlarge.nvidia.gpu
             torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu121'
             gpu-arch-type: "cuda"
@@ -65,7 +65,6 @@ jobs:
         export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
         python -m pip install --upgrade pip
         pip install ${{ matrix.torch-spec }}
-        pip install -r requirements.txt
         pip install -r dev-requirements.txt
         pip install .
         pytest test --verbose -s

README.md

@@ -31,21 +31,15 @@ From source
 ```Shell
 git clone https://github.com/pytorch/ao
 cd ao
-pip install -r requirements.txt
-pip install -r dev-requirements.txt
+python setup.py install 
 ```
 
-There are two options;
--If you plan to be developing the library run:
+If you plan to be developing the library run:
 ```Shell
+pip install -r dev-requirements.txt
 python setup.py develop
 ```
 
-If you want to install from source run
-```Shell
-python setup.py install 
-```
-
 ** Note:
 If you are running into any issues while building `ao` cpp extensions you can instead build using
 

dev-requirements.txt

@@ -4,13 +4,17 @@ unittest-xml-reporting
 parameterized
 packaging
 transformers
+hypothesis # Avoid test derandomization warning
+sentencepiece # for gpt-fast tokenizer
+expecttest
 
 # For prototype features and benchmarks
 bitsandbytes #needed for testing triton quant / dequant ops for 8-bit optimizers
 matplotlib
 pandas
 fire  # QOL for commandline scripts
 tabulate  # QOL for printing tables to stdout
+numpy # mx benchmarks leverage benchmark timer which uses numpy
 
 # Custom CUDA Extensions
 ninja

setup.py

@@ -94,7 +94,7 @@ def get_extensions():
     pytorch_dep += "==" + os.getenv("PYTORCH_VERSION")
 
 requirements = [
-    "numpy",
+    # "numpy",
     pytorch_dep,
 ]
 

torchao/__init__.py

@@ -1,6 +1,11 @@
 import torch
 import logging
 
+# torch/nested/_internal/nested_tensor.py:417: UserWarning: Failed to initialize NumPy: No module named 'numpy'
+import warnings
+warnings.filterwarnings("ignore", message="Failed to initialize NumPy: No module named 'numpy'")
+
+
 # We use this "hack" to set torchao.__version__ correctly
 # the version of ao is dependent on environment variables for multiple architectures
 # For local development this will default to whatever is version.txt

torchao/utils.py

@@ -1,9 +1,8 @@
 import torch
-import torch.utils.benchmark as benchmark
 from typing import Tuple
 from functools import reduce
+from importlib.metadata import version
 from math import gcd
-from packaging import version
 import torch.nn.utils.parametrize as parametrize
 import itertools
 
@@ -19,6 +18,7 @@
     "TORCH_VERSION_AFTER_2_2",
     "TORCH_VERSION_AFTER_2_3",
     "TORCH_VERSION_AFTER_2_4",
+    "TORCH_VERSION_AFTER_2_5",
 ]
 
 
@@ -64,8 +64,9 @@ def wrapper(*args, **kwargs):
 
 
 def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
+    import torch.utils.benchmark as benchmark # this avoids importing numpy when torchao module is loaded
+
     # Manual warmup
-
     f(*args, **kwargs)
     f(*args, **kwargs)
 
@@ -163,20 +164,14 @@ def unwrap_tensor_subclass(model, filter_fn=None):
         unwrap_tensor_subclass(child)
     return model
 
-if version.parse(torch.__version__) >= version.parse("2.4.0.dev"):
-    TORCH_VERSION_AFTER_2_4 = True
-else:
-    TORCH_VERSION_AFTER_2_4 = False
 
-if version.parse(torch.__version__) >= version.parse("2.3.0.dev"):
-    TORCH_VERSION_AFTER_2_3 = True
-else:
-    TORCH_VERSION_AFTER_2_3 = False
+def torch_version_at_least(min_version):
+    return version("torch") >= min_version
 
-if version.parse(torch.__version__) >= version.parse("2.2.0.dev"):
-    TORCH_VERSION_AFTER_2_2 = True
-else:
-    TORCH_VERSION_AFTER_2_2 = False
+TORCH_VERSION_AFTER_2_5 = torch_version_at_least("2.5.0.dev")
+TORCH_VERSION_AFTER_2_4 = torch_version_at_least("2.4.0.dev")
+TORCH_VERSION_AFTER_2_3 = torch_version_at_least("2.3.0.dev")
+TORCH_VERSION_AFTER_2_2 = torch_version_at_least("2.2.0.dev")
 
 def is_fbcode():
     return not hasattr(torch.version, "git_version")