Add oumi[quantization] optional dependency (#1902)

Author: wizeng23

configs/examples/quantization/README.md

@@ -4,6 +4,8 @@
 
 This directory contains example configurations for model quantization using Oumi's AWQ and BitsAndBytes quantization methods.
 
+> **NOTE**: Quantization requires a GPU to run.
+
 ## Configuration Files
 
 - **`awq_quantization_config.yaml`** - AWQ 4-bit quantization with calibration
@@ -15,7 +17,7 @@ This directory contains example configurations for model quantization using Oumi
 # Simplest command-line usage
 oumi quantize --method awq_q4_0 --model "TinyLlama/TinyLlama-1.1B-Chat-v1.0" --output quantized_model
 
-# Using configuration file (requires GPU)
+# Using configuration file
 oumi quantize --config configs/examples/quantization/awq_quantization_config.yaml
 ```
 
@@ -40,10 +42,12 @@ oumi quantize --config configs/examples/quantization/awq_quantization_config.yam
 ## Requirements
 
 ```bash
-# For AWQ quantization
+pip install oumi[quantization]
+
+# Alternatively, for AWQ quantization only
 pip install autoawq
 
-# For BitsAndBytes quantization
+# Alternatively, for BitsAndBytes quantization only
 pip install bitsandbytes
 ```
 

docs/user_guides/quantization.md

@@ -4,6 +4,8 @@
 
 This guide covers the `oumi quantize` command for reducing model size while maintaining performance.
 
+> **NOTE**: Quantization requires a GPU to run.
+
 ## Quick Start
 
 ```bash
@@ -93,10 +95,12 @@ Currently supported output formats:
 ## Installation
 
 ```bash
-# For AWQ quantization
+pip install oumi[quantization]
+
+# Alternatively, for AWQ quantization only
 pip install autoawq
 
-# For BitsAndBytes quantization
+# Alternatively, for BitsAndBytes quantization only
 pip install bitsandbytes
 ```
 

notebooks/Oumi - A Tour.ipynb

@@ -493,7 +493,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.11.13"
   }
  },
  "nbformat": 4,

notebooks/Oumi - Quantization Tutorial.ipynb

@@ -31,13 +31,17 @@
     "\n",
     "⚠️ **DEVELOPMENT STATUS**: The quantization feature is currently under active development. Some features may change in future releases.\n",
     "\n",
-    "First, let's install Oumi with GPU support and the required quantization libraries:\n",
-    "\n",
-    "```bash\n",
-    "pip install oumi[gpu]\n",
-    "pip install autoawq\n",
-    "pip install triton==3.0.0  # Required for AWQ inference compatibility\n",
-    "```"
+    "First, let's install Oumi with GPU support and the required quantization libraries:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install oumi[gpu,quantization]\n",
+    "%pip install triton==3.0.0  # Required for AWQ inference compatibility"
    ]
   },
   {

pyproject.toml

@@ -154,6 +154,8 @@ evaluation = [
     "sentencepiece>=0.1.98",
 ]
 
+quantization = ["autoawq>=0.2.0,<0.3", "bitsandbytes>=0.45.0,<0.46"]
+
 bitnet = ["onebitllms>=0.0.3"]
 
 cambrian = [

src/oumi/quantize/awq_quantizer.py

@@ -60,7 +60,7 @@ def raise_if_requirements_not_met(self):
         if self._awq is None:
             raise RuntimeError(
                 "AWQ quantization requires autoawq library.\n"
-                "Install with: `pip install autoawq`\n"
+                "Install with: `pip install oumi[quantization]`\n"
             )
 
         if not torch.cuda.is_available():