Skip to content

Commit

Permalink
update setup.py
Browse files Browse the repository at this point in the history
  • Loading branch information
XiaohanZhangCMU committed Aug 14, 2024
1 parent 26a9386 commit ae8037f
Showing 1 changed file with 15 additions and 21 deletions.
36 changes: 15 additions & 21 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,29 +52,27 @@
]

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.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.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 = {}
Expand All @@ -92,29 +90,25 @@
]

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'] = [
Expand Down

0 comments on commit ae8037f

Please sign in to comment.