diff --git a/README.md b/README.md index f92e03a..3ba2e3e 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,9 @@ The following optional dependencies are available: - `fp8-infer`: `torchao` package for fp8 inference - `gptq`: `GPTQModel` package for W4A16 quantization - `mx`: `microxcaling` package for MX quantization -- `opt`: Shortcut for `fp8`, `gptq`, and `mx` installs +- `data`: `datasets` package for loading calibration data (required by `fms_mo.run_quant` and direct quantization) +- `launcher`: `accelerate` package for the multi-GPU launcher in `build/` +- `opt`: Shortcut for `fp8`, `gptq`, `mx`, `data`, and `launcher` installs - `aiu`: `ibm-fms` package for AIU model deployment - `torchvision`: `torch` package for image recognition training and inference - `triton`: `triton` package for matrix multiplication kernels diff --git a/fms_mo/dq.py b/fms_mo/dq.py index e7d644a..d4598c0 100644 --- a/fms_mo/dq.py +++ b/fms_mo/dq.py @@ -23,7 +23,6 @@ import logging # Third Party -from datasets import load_from_disk from torch.utils.data import DataLoader from tqdm import tqdm from transformers import ( @@ -51,6 +50,15 @@ from fms_mo.utils.eval_utils import Evaluator, eval_llm_1GPU from fms_mo.utils.utils import patch_torch_bmm, prepare_input +try: + # Third Party + from datasets import load_from_disk +except ImportError as e: + raise ImportError( + "The `datasets` package is required for data loading. Install it with " + "`pip install fms-model-optimizer[data]`." + ) from e + logger = logging.getLogger(__name__) diff --git a/fms_mo/run_quant.py b/fms_mo/run_quant.py index 4521dca..abc9184 100644 --- a/fms_mo/run_quant.py +++ b/fms_mo/run_quant.py @@ -33,7 +33,6 @@ import traceback # Third Party -from datasets import load_from_disk from torch.cuda import OutOfMemoryError from transformers import ( AutoModelForMaskedLM, @@ -62,6 +61,15 @@ from fms_mo.utils.import_utils import available_packages from fms_mo.utils.logging_utils import set_log_level +try: + # Third Party + from datasets import load_from_disk +except ImportError as e: + raise ImportError( + "The `datasets` package is required for data loading. Install it with " + "`pip install fms-model-optimizer[data]`." + ) from e + def quantize( model_args: ModelArguments, diff --git a/fms_mo/utils/calib_data.py b/fms_mo/utils/calib_data.py index a61d3a7..ccc0fe5 100755 --- a/fms_mo/utils/calib_data.py +++ b/fms_mo/utils/calib_data.py @@ -25,11 +25,19 @@ import random # Third Party -from datasets import load_dataset, load_from_disk from transformers import AutoTokenizer, BatchEncoding -import datasets import torch +try: + # Third Party + from datasets import load_dataset, load_from_disk + import datasets +except ImportError as e: + raise ImportError( + "The `datasets` package is required for data loading. Install it with " + "`pip install fms-model-optimizer[data]`." + ) from e + def return_tokenized_samples(nsamples, trainenc, seqlen, sequential=False): """Randomly crop nsamples sequence from trainenc, each with the length of seqlen. diff --git a/pyproject.toml b/pyproject.toml index 6190d2b..6f6f37a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,30 +23,32 @@ classifiers=[ dynamic = ["version"] dependencies = [ "numpy>=1.26.4,<2.3.0", -"accelerate>=0.20.3,!=0.34,<1.11", "transformers>4.45,<=5.12.1", "torch>=2.2.0,<2.12.0", "tqdm>=4.66.2,<5.0", -"datasets>=3.0.0,<5.0", "pandas", "safetensors", "pkginfo>1.10" ] [project.optional-dependencies] -examples = ["ninja>=1.11.1.1,<2.0", "evaluate", "huggingface_hub"] +# data loading for the quantization pipeline (run_quant, dq, calib_data) +data = ["datasets>=3.0.0,<5.0"] +# multi-GPU launcher (build/accelerate_launch.py) and examples +launcher = ["accelerate>=0.20.3,!=0.34,<1.11"] +examples = ["ninja>=1.11.1.1,<2.0", "evaluate", "huggingface_hub", "fms-model-optimizer[data, launcher]"] fp8 = ["llmcompressor", "torchao==0.11"] # FP8 matmul on CPU needs a fix before advancing torchao > 0.11 fp8-infer = ["torchao==0.11"] gptq = ["Cython", "gptqmodel>=1.7.3"] mx = ["microxcaling>=1.1"] -opt = ["fms-model-optimizer[fp8, gptq, mx]"] +opt = ["fms-model-optimizer[fp8, gptq, mx, data, launcher]"] aiu = ["ibm-fms>=0.0.8"] torchvision = ["torchvision>=0.17"] flash-attn = ["flash-attn>=2.5.3,<3.0"] triton = ["triton>=3.0,<3.5"] visualize = ["matplotlib", "graphviz", "pygraphviz", "tensorboard", "notebook"] dev = ["pre-commit>=3.0.4,<5.0"] -test = ["pytest", "pillow"] +test = ["pytest", "pillow", "fms-model-optimizer[data, launcher]"] [project.urls] homepage = "https://github.com/foundation-model-stack/fms-model-optimizer"