From 9548f57af0da88549619302062844b53ca929d77 Mon Sep 17 00:00:00 2001 From: xiaohanzhangcmu Date: Wed, 14 Aug 2024 11:31:09 -0700 Subject: [PATCH] update setup.py --- setup.py | 100 ++++++++++++++++++++++++++++++++----------------------- 1 file changed, 58 insertions(+), 42 deletions(-) diff --git a/setup.py b/setup.py index e56b85f9f7..dda2ccb491 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ import copy import os -import re +from typing import Any, Dict, Mapping import setuptools from setuptools import setup @@ -15,14 +15,15 @@ _REPO_REAL_PATH = os.path.dirname(os.path.realpath(__file__)) _PACKAGE_REAL_PATH = os.path.join(_REPO_REAL_PATH, _PACKAGE_DIR) -# Read the repo version +# Read the llm-foundry version # We can't use `.__version__` from the library since it's not installed yet -with open(os.path.join(_PACKAGE_REAL_PATH, '__init__.py')) as f: +version_path = os.path.join(_PACKAGE_REAL_PATH, '_version.py') +with open(version_path, encoding='utf-8') as f: + version_globals: Dict[str, Any] = {} + version_locals: Mapping[str, object] = {} content = f.read() -# regex: '__version__', whitespace?, '=', whitespace, quote, version, quote -# we put parens around the version so that it becomes elem 1 of the match -expr = re.compile(r"""^__version__\W+=\W+['"]([0-9\.]*)['"]""", re.MULTILINE) -repo_version = expr.findall(content)[0] + exec(content, version_globals, version_locals) + repo_version = str(version_locals['__version__']) # Use repo README for PyPi description with open('README.md', 'r', encoding='utf-8') as fh: @@ -45,44 +46,43 @@ classifiers = [ 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', + 'Programming Language :: Python :: 3.11', ] install_requires = [ - 'mosaicml[libcloud,wandb,oci,gcs]>=0.20.1,<0.21', - 'mlflow>=2.10,<3', - 'accelerate>=0.25,<0.26', # for HF inference `device_map` - 'transformers>=4.38.2,<4.39', - 'mosaicml-streaming>=0.7.4,<0.8', - 'torch>=2.2.1,<2.3', - 'datasets>=2.16,<2.17', + 'mosaicml[libcloud,wandb,oci,gcs,mlflow]>=0.23.4,<0.24', + 'mlflow>=2.14.1,<2.16', + 'accelerate>=0.25,<0.34', # for HF inference `device_map` + 'transformers>=4.43.2,<4.44', + 'mosaicml-streaming>=0.8.0,<0.9', + 'torch>=2.3.0,<2.4', + 'datasets>=2.19,<2.20', 'fsspec==2023.6.0', # newer version results in a bug in datasets that duplicates data - 'sentencepiece==0.1.97', - 'einops==0.7.0', + 'sentencepiece==0.2.0', + 'einops==0.8.0', 'omegaconf>=2.2.3,<3', 'slack-sdk<4', 'mosaicml-cli>=0.6.10,<1', - 'onnx==1.14.0', - 'onnxruntime==1.15.1', - 'cmake>=3.25.0,<=3.26.3', # required for triton-pre-mlir below - # PyPI does not support direct dependencies, so we remove this line before uploading from PyPI - 'triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir_sm90#subdirectory=python', + 'onnx==1.16.2', + 'onnxruntime==1.18.1', 'boto3>=1.21.45,<2', - 'huggingface-hub>=0.17.0,<1.0', + 'huggingface-hub>=0.19.0,<0.25', 'beautifulsoup4>=4.12.2,<5', # required for model download utils 'tenacity>=8.2.3,<9', - 'dask[distributed]>=2023.11.0', + 'catalogue>=2,<3', + 'typer<1', ] extra_deps = {} extra_deps['dev'] = [ + 'coverage[toml]==7.6.1', 'pre-commit>=3.4.0,<4', 'pytest>=7.2.1,<8', - 'pytest_codeblocks>=0.16.1,<0.17', - 'pytest-cov>=4,<5', + 'pytest_codeblocks>=0.16.1,<0.18', + 'pytest-cov>=4,<6', 'pyright==1.1.256', 'toml>=0.10.2,<0.11', 'packaging>=21,<23', @@ -90,42 +90,54 @@ ] extra_deps['databricks'] = [ - 'mosaicml[databricks]>=0.20.1,<0.21', + 'mosaicml[databricks]>=0.23.4,<0.24', 'databricks-sql-connector>=3,<4', 'databricks-connect==14.1.0', 'lz4>=4,<5', ] extra_deps['tensorboard'] = [ - 'mosaicml[tensorboard]>=0.20.1,<0.21', + 'mosaicml[tensorboard]>=0.23.4,<0.24', ] # Flash 2 group kept for backwards compatibility extra_deps['gpu-flash2'] = [ - 'flash-attn==2.5.0', + 'flash-attn>=2.5.8,<3', ] extra_deps['gpu'] = copy.deepcopy(extra_deps['gpu-flash2']) extra_deps['peft'] = [ - 'mosaicml[peft]>=0.20.1,<0.21', -] - -extra_deps['olmo'] = [ - 'ai2-olmo>0.2.4', + 'mosaicml[peft]>=0.23.4,<0.24', ] extra_deps['openai'] = [ 'openai==1.3.8', 'tiktoken==0.4.0', ] -extra_deps['all-cpu'] = set( - dep for key, deps in extra_deps.items() for dep in deps if 'gpu' not in key) -extra_deps['all'] = set(dep for key, deps in extra_deps.items() for dep in deps - if key not in {'gpu-flash2', 'all-cpu'}) -extra_deps['all-flash2'] = set(dep for key, deps in extra_deps.items() - for dep in deps - if key not in {'gpu', 'all', 'all-cpu'}) + +extra_deps['megablocks'] = [ + 'megablocks==0.5.1', + 'grouped-gemm==0.1.4', +] + +extra_deps['databricks-serverless'] = { + dep for key, deps in extra_deps.items() for dep in deps + if 'gpu' not in key and 'megablocks' not in key and + 'databricks-connect' not in dep +} +extra_deps['all-cpu'] = { + dep for key, deps in extra_deps.items() for dep in deps + if 'gpu' not in key and 'megablocks' not in key +} +extra_deps['all'] = { + dep for key, deps in extra_deps.items() for dep in deps + if key not in {'gpu-flash2', 'all-cpu'} +} +extra_deps['all-flash2'] = { + dep for key, deps in extra_deps.items() for dep in deps + if key not in {'gpu', 'all', 'all-cpu'} +} setup( name=_PACKAGE_NAME, @@ -140,9 +152,13 @@ 'llmfoundry': ['py.typed'], }, packages=setuptools.find_packages( - exclude=['.github*', 'mcli*', 'scripts*', 'tests*']), + exclude=['.github*', 'mcli*', 'scripts*', 'tests*'], + ), classifiers=classifiers, install_requires=install_requires, extras_require=extra_deps, python_requires='>=3.9', + entry_points={ + 'console_scripts': ['llmfoundry = llmfoundry.cli.cli:app'], + }, )