Merge branch 'main' into hanxian_8da4w_trial

.github/scripts/validate_binaries.sh

@@ -0,0 +1,2 @@
+pip install ${PYTORCH_PIP_PREFIX} torchao --index-url ${PYTORCH_PIP_DOWNLOAD_URL}
+python  ./test/smoke_tests/smoke_tests.py

.github/workflows/validate-binaries.yml

@@ -0,0 +1,49 @@
+name: Validate binaries
+
+on:
+  workflow_call:
+    inputs:
+      channel:
+        description: "Channel to use (nightly, test, release, all)"
+        required: false
+        type: string
+        default: release
+      ref:
+        description: "Reference to checkout, defaults to empty"
+        default: ""
+        required: false
+        type: string
+  workflow_dispatch:
+    inputs:
+      channel:
+        description: "Channel to use (nightly, test, release, all)"
+        required: true
+        type: choice
+        options:
+          - release
+          - nightly
+          - test
+          - all
+      ref:
+        description: "Reference to checkout, defaults to empty"
+        default: ""
+        required: false
+        type: string
+      pytorch_version:
+        description: "PyTorch version to validate (ie. 2.0, 2.2.2, etc.) - optional"
+        default: ""
+        required: false
+        type: string
+jobs:
+  validate-binaries:
+    uses: pytorch/test-infra/.github/workflows/validate-domain-library.yml@main
+    with:
+      package_type: "wheel"
+      version: ${{ inputs.version }}
+      os: "linux"
+      channel: ${{ inputs.channel }}
+      repository: "pytorch/ao"
+      with_cuda: "enable"
+      with_rocm: "disable"
+      smoke_test: "source ./.github/scripts/validate_binaries.sh"
+      install_torch: true

README.md

@@ -81,17 +81,34 @@ A key design principle for us is composability as in any new dtype or layout we
 
 
 ### Installation
+
 `torchao` makes liberal use of several new features in Pytorch, it's recommended to use it with the current nightly or latest stable version of PyTorch.
 
-Stable release from Pypi which will default to PyTorch 2.3.1 and CUDA 12.1
+#### Install torch
+
+Install torch stable
+
+```
+pip install torch
+```
+
+Or torch nightlies
+
+```
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
+```
+
+#### Install torchao
+
+Stable release from Pypi which will default to CUDA 12.1
 
 ```Shell
 pip install torchao
 ```
 
 Stable Release from the PyTorch index
 ```Shell
-pip install torchao --extra-index-url https://download.pytorch.org/whl/test/cu121 # full options are cpu/cu118/cu121/cu124
+pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124
 ```
 
 Nightly Release
@@ -116,6 +133,13 @@ python setup.py install
 * [vayuda](https://github.com/vayuda) with generic bitpacking kernels that were code generated using pure PyTorch [prototype/common](torchao/prototype/common)
 * [andreaskopf](https://github.com/andreaskoepf) and [melvinebenezer](https://github.com/melvinebenezer) with [1 bit LLMs](torchao/prototype/dtypes) Bitnet 1.58 bitpacked into uint2 and fully code-generated with torch.compile
 
+## Blogs and Videos
+* [Accelerating Neural Network Training with Semi-Structured (2:4) Sparsity](https://pytorch.org/blog/accelerating-neural-network-training/)
+* [https://mobiusml.github.io/whisper-static-cache-blog/](https://mobiusml.github.io/whisper-static-cache-blog/)
+* [Slaying OOMs at the Mastering LLM's course](https://x.com/HamelHusain/status/1800315287574847701)
+* [Advanced Quantization at CUDA MODE](https://youtu.be/1u9xUK3G4VM?si=4JcPlw2w8chPXW8J)
+* [Chip Huyen's GPU Optimization Workshop](https://www.youtube.com/live/v_q2JTIqE20?si=mf7HeZ63rS-uYpS6)
+
 ## How to contribute
 
 This repository is currently under heavy development

packaging/pre_build_script.sh

@@ -9,4 +9,12 @@ set -eux
 
 echo "This script is run before building torchao binaries"
 
+python -m pip install --upgrade pip
+if [ -z ${PYTORCH_VERSION:-} ]; then
+    PYTORCH_DEP="torch"
+else
+    PYTORCH_DEP="torch==$PYTORCH_VERSION"
+fi
+pip install $PYTORCH_DEP
+
 pip install setuptools wheel twine auditwheel

setup.py

@@ -110,15 +110,6 @@ def get_extensions():
 
     return ext_modules
 
-# Mimic code from torchvision https://github.com/pytorch/vision/blob/143d078b28f00471156a4e562dd3836370acc9ee/setup.py#L58
-pytorch_dep = "torch"
-if os.getenv("PYTORCH_VERSION"):
-    pytorch_dep += "==" + os.getenv("PYTORCH_VERSION")
-
-requirements = [
-    pytorch_dep,
-]
-
 setup(
     name=package_name,
     version=version+version_suffix,
@@ -128,7 +119,6 @@ def get_extensions():
         "torchao.kernel.configs": ["*.pkl"],
     },
     ext_modules=get_extensions() if use_cpp != "0" else None,
-    install_requires=requirements,
     extras_require={"dev": read_requirements("dev-requirements.txt")},
     description="Package for applying ao techniques to GPU models",
     long_description=open("README.md").read(),

test/smoke_tests/smoke_tests.py

@@ -0,0 +1,6 @@
+"""Run smoke tests"""
+
+import torchao
+
+
+print("torchao version is ", torchao.__version__)

torchao/__init__.py

@@ -39,6 +39,3 @@
     "autoquant",
     "quantize",
 ]
-
-# test ghstack flow
-# test ghtstack flow 2

torchao/prototype/mx_formats/custom_cast.py

@@ -15,7 +15,7 @@
 # TODO(future): if needed, make the below work on previous PyTorch versions,
 # just need to hunt down the previous location of `libdevice`. An assert
 # at the callsite prevents usage of this on unsupported versions.
-if TORCH_VERSION_AFTER_2_4:
+if TORCH_VERSION_AFTER_2_4 and has_triton():
     from torch._inductor.runtime.triton_helpers import libdevice
 
 from torchao.prototype.mx_formats.constants import (

version.txt

@@ -1 +1 @@
-0.3.0
+0.3.1