diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index e89e27642..181fd9dc9 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -6,7 +6,7 @@ on: types: [opened, synchronize, reopened, ready_for_review] paths: - '**.py' - - 'pyproject.toml' + - 'requirements.txt' - '.github/workflows/*.yml' - "*.[q]md" - "examples/**/*.y[a]?ml" diff --git a/.github/workflows/multi-gpu-e2e.yml b/.github/workflows/multi-gpu-e2e.yml index d2d3123e3..2bb499ded 100644 --- a/.github/workflows/multi-gpu-e2e.yml +++ b/.github/workflows/multi-gpu-e2e.yml @@ -4,7 +4,8 @@ on: pull_request: paths: - 'tests/e2e/multigpu/**.py' - - 'pyproject.toml' + - 'requirements.txt' + - 'setup.py' - 'pyproject.toml' - '.github/workflows/multi-gpu-e2e.yml' - 'scripts/cutcrossentropy_install.py' diff --git a/.github/workflows/tests-nightly.yml b/.github/workflows/tests-nightly.yml index 8ccebfdc1..663b0476e 100644 --- a/.github/workflows/tests-nightly.yml +++ b/.github/workflows/tests-nightly.yml @@ -72,6 +72,14 @@ jobs: run: | pip3 install torch==${{ matrix.pytorch_version }} torchvision + - name: Update requirements.txt + run: | + sed -i 's#^transformers.*#transformers @ git+https://github.com/huggingface/transformers.git@main#' requirements.txt + sed -i 's#^peft.*#peft @ git+https://github.com/huggingface/peft.git@main#' requirements.txt + sed -i 's#^accelerate.*#accelerate @ git+https://github.com/huggingface/accelerate.git@main#' requirements.txt + sed -i 's#^trl.*#trl @ git+https://github.com/huggingface/trl.git@main#' requirements.txt + sed -i 's#^datasets.*#datasets @ git+https://github.com/huggingface/datasets.git@main#' requirements.txt + - name: Install dependencies run: | pip3 show torch @@ -80,15 +88,6 @@ jobs: python scripts/cutcrossentropy_install.py | sh pip3 install -r requirements-dev.txt -r requirements-tests.txt - - name: Override with nightly HF packages - run: | - pip3 install --no-deps \ - "transformers @ git+https://github.com/huggingface/transformers.git@main" \ - "peft @ git+https://github.com/huggingface/peft.git@main" \ - "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \ - "trl @ git+https://github.com/huggingface/trl.git@main" \ - "datasets @ git+https://github.com/huggingface/datasets.git@main" - - name: Make sure PyTorch version wasn't clobbered run: | python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__" diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index ad9f551a0..5099e447c 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -7,7 +7,7 @@ on: - "main" paths: - '**.py' - - 'pyproject.toml' + - 'requirements.txt' - '.github/workflows/*.yml' - 'requirements-tests.txt' - 'cicd/cicd.sh' @@ -16,7 +16,7 @@ on: types: [opened, synchronize, reopened, ready_for_review] paths: - '**.py' - - 'pyproject.toml' + - 'requirements.txt' - '.github/workflows/*.yml' - 'requirements-tests.txt' - 'cicd/cicd.sh' diff --git a/MANIFEST.in b/MANIFEST.in index 729f7e7a3..3fbb0edca 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,6 @@ +include requirements.txt include README.md include LICENSE -include VERSION +include src/setuptools_axolotl_dynamic_dependencies.py include src/axolotl/utils/chat_templates/templates/*.jinja +recursive-include axolotl *.py diff --git a/pyproject.toml b/pyproject.toml index a4e9f32a6..9cee4a520 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,143 +1,15 @@ [build-system] -requires = ["setuptools>=64", "wheel", "setuptools_scm>=8"] +requires = ["setuptools>=64", "wheel", "setuptools_scm>=8", "packaging==26.0"] build-backend = "setuptools.build_meta" [project] name = "axolotl" -dynamic = ["version"] +dynamic = ["version", "dependencies", "optional-dependencies"] description = "LLM Trainer" readme = "README.md" requires-python = ">=3.10" # license = "Apache-2.0" -dependencies = [ - # Core ML stack - "torch>=2.6.0", - "packaging==26.0", - "huggingface_hub>=1.1.7", - "peft>=0.18.1", - "tokenizers>=0.22.1", - "transformers==5.3.0", - "accelerate==1.13.0", - "datasets==4.5.0", - "trl==0.29.0", - "hf_xet==1.3.2", - "kernels==0.12.2", - "trackio>=0.16.1", - "typing-extensions>=4.15.0", - "optimum==1.16.2", - "hf_transfer", - "sentencepiece", - "gradio>=6.2.0,<7.0", - "modal==1.3.0.post1", - "pydantic>=2.10.6", - "addict", - "fire", - "PyYAML>=6.0", - "requests", - "wandb", - "einops", - "colorama", - "numba>=0.61.2", - "numpy>=2.2.6", - - # Evaluation & metrics - "evaluate==0.4.1", - "scipy", - "nvidia-ml-py==12.560.30", - "art", - "tensorboard", - "python-dotenv==1.0.1", - - # Remote filesystems - "s3fs>=2024.5.0", - "gcsfs>=2025.3.0", - "adlfs>=2024.5.0", - "ocifs==1.3.2", - - "zstandard==0.22.0", - "fastcore", - - # lm eval harness - "lm_eval==0.4.7", - "langdetect==1.0.9", - "immutabledict==4.2.0", - "antlr4-python3-runtime==4.13.2", - - "schedulefree==1.4.1", - "openenv-core==0.1.0", - - # Axolotl contribs - "axolotl-contribs-lgpl==0.0.7", - "axolotl-contribs-mit==0.0.6", - - # Telemetry - "posthog==6.7.11", - - "mistral-common==1.10.0", - - # Platform-specific (Linux only) - "bitsandbytes==0.49.1 ; sys_platform != 'darwin'", - "triton>=3.4.0 ; sys_platform != 'darwin'", - "xformers>=0.0.23.post1 ; sys_platform != 'darwin'", - "liger-kernel==0.7.0 ; sys_platform != 'darwin'", - "torchao==0.16.0 ; sys_platform != 'darwin' and platform_machine != 'aarch64'", - - # Architecture-specific - "fla-core==0.4.1 ; platform_machine != 'aarch64'", - "flash-linear-attention==0.4.1 ; platform_machine != 'aarch64'", -] - -[project.optional-dependencies] -flash-attn = ["flash-attn==2.8.3"] -ring-flash-attn = [ - "flash-attn==2.8.3", - "ring-flash-attn>=0.1.7", -] -deepspeed = [ - "deepspeed>=0.18.6,<0.19.0", - "deepspeed-kernels", -] -mamba-ssm = [ - "mamba-ssm==1.2.0.post1", - "causal_conv1d", -] -auto-gptq = [ - "auto-gptq==0.5.1", -] -mlflow = [ - "mlflow", -] -galore = [ - "galore_torch", -] -apollo = [ - "apollo-torch", -] -optimizers = [ - "galore_torch", - "apollo-torch", - "lomo-optim==0.1.1", - "torch-optimi==0.2.1", - "came_pytorch==0.1.3", -] -ray = [ - "ray[train]>=2.52.1", -] -vllm = [ - "vllm>=0.10.0", -] -llmcompressor = [ - "llmcompressor>=0.10.0", -] -fbgemm-gpu = ["fbgemm-gpu-genai>=1.3.0"] -opentelemetry = [ - "opentelemetry-api", - "opentelemetry-sdk", - "opentelemetry-exporter-prometheus", - "prometheus-client", -] - [project.scripts] axolotl = "axolotl.cli.main:main" @@ -146,15 +18,18 @@ Homepage = "https://axolotl.ai/" Documentation = "https://docs.axolotl.ai/" Repository = "https://github.com/axolotl-ai-cloud/axolotl.git" -[tool.setuptools] -include-package-data = true +[tool.setuptools_scm] -[tool.setuptools.packages.find] -where = ["src"] +[tool.setuptools] +py-modules = ["setuptools_axolotl_dynamic_dependencies"] +include-package-data = true [tool.setuptools.dynamic] version = { file = "VERSION" } +[tool.setuptools.cmdclass] +build_py = "setuptools_axolotl_dynamic_dependencies.BuildPyCommand" + [tool.ruff] line-length = 88 target-version = "py310" @@ -192,40 +67,5 @@ markers = [ "slow: marks tests as slow", ] -# UV specific configuration -[tool.uv] -prerelease = "allow" -conflicts = [ - [ - { package = "axolotl" }, - { extra = "vllm" }, - ], - [ - { package = "axolotl" }, - { extra = "flash-attn" }, - ], - [ - { package = "axolotl" }, - { extra = "ring-flash-attn" }, - ], - [ - { package = "axolotl" }, - { extra = "mamba-ssm" }, - ], - [ - { package = "axolotl" }, - { extra = "auto-gptq" }, - ], - [ - { package = "axolotl" }, - { extra = "fbgemm-gpu" }, - ], -] - [tool.uv.extra-build-dependencies] axolotl = ["huggingface_hub"] -mamba-ssm = ["torch"] -causal-conv1d = ["torch"] -flash-attn = ["torch"] -deepspeed = ["torch"] -auto-gptq = ["torch"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 000000000..3fd75c3fa --- /dev/null +++ b/requirements.txt @@ -0,0 +1,78 @@ +--extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/ + +# START section of dependencies that don't install on Darwin/MacOS +bitsandbytes==0.49.1 +triton>=3.4.0 +mamba-ssm==1.2.0.post1 +xformers>=0.0.23.post1 +liger-kernel==0.7.0 +# END section + +packaging==26.0 +huggingface_hub>=1.1.7 +peft>=0.18.1 +tokenizers>=0.22.1 +transformers==5.3.0 +accelerate==1.13.0 +datasets==4.5.0 +deepspeed>=0.18.6,<0.19.0 +trl==0.29.0 +hf_xet==1.3.2 +kernels==0.12.2 + +fla-core==0.4.1 +flash-linear-attention==0.4.1 + +trackio>=0.16.1 +typing-extensions>=4.15.0 + +optimum==1.16.2 +hf_transfer +sentencepiece +gradio>=6.2.0,<7.0 + +modal==1.3.0.post1 +pydantic>=2.10.6 +addict +fire +PyYAML>=6.0 +requests +wandb +einops +colorama +numba>=0.61.2 +numpy>=2.2.6 + +# qlora things +evaluate==0.4.1 +scipy +nvidia-ml-py==12.560.30 +art +tensorboard +python-dotenv==1.0.1 + +# remote filesystems +s3fs>=2024.5.0 +gcsfs>=2025.3.0 +adlfs>=2024.5.0 +ocifs==1.3.2 + +zstandard==0.22.0 +fastcore + +# lm eval harness +lm_eval==0.4.7 +langdetect==1.0.9 +immutabledict==4.2.0 +antlr4-python3-runtime==4.13.2 + +torchao==0.16.0 +openenv-core==0.1.0 +schedulefree==1.4.1 + +axolotl-contribs-lgpl==0.0.7 +axolotl-contribs-mit==0.0.6 +# telemetry +posthog==6.7.11 + +mistral-common==1.10.0 diff --git a/setup.py b/setup.py new file mode 100644 index 000000000..d4e1894f7 --- /dev/null +++ b/setup.py @@ -0,0 +1,230 @@ +"""setup.py for axolotl""" + +import os +import platform +import re +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path + +from setuptools import find_packages, setup + + +def parse_requirements(extras_require_map): + _install_requires = [] + _dependency_links = [] + with open("./requirements.txt", encoding="utf-8") as requirements_file: + lines = [r.strip() for r in requirements_file.readlines()] + for line in lines: + is_extras = "deepspeed" in line or "mamba-ssm" in line + if line.startswith("--extra-index-url"): + # Handle custom index URLs + _, url = line.split() + _dependency_links.append(url) + elif not is_extras and line and line[0] != "#": + # Handle standard packages + _install_requires.append(line) + try: + xformers_version = [req for req in _install_requires if "xformers" in req][0] + install_xformers = platform.machine() != "aarch64" + if platform.machine() == "aarch64": + # skip on ARM64 + skip_packages = [ + "torchao", + "fla-core", + "flash-linear-attention", + ] + _install_requires = [ + req + for req in _install_requires + if re.split(r"[>=<]", req)[0].strip() not in skip_packages + ] + if "Darwin" in platform.system(): + # skip packages not compatible with OSX + skip_packages = [ + "bitsandbytes", + "triton", + "mamba-ssm", + "xformers", + "liger-kernel", + ] + _install_requires = [ + req + for req in _install_requires + if re.split(r"[>=<]", req)[0].strip() not in skip_packages + ] + print( + _install_requires, [req in skip_packages for req in _install_requires] + ) + else: + # detect the version of torch already installed + # and set it so dependencies don't clobber the torch version + try: + torch_version = version("torch") + except PackageNotFoundError: + torch_version = "2.8.0" # default to torch 2.8.0 + _install_requires.append(f"torch=={torch_version}") + + version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version) + if version_match: + major, minor, patch = version_match.groups() + major, minor = int(major), int(minor) + patch = ( + int(patch) if patch is not None else 0 + ) # Default patch to 0 if not present + else: + raise ValueError("Invalid version format") + + torch_parts = torch_version.split("+") + if len(torch_parts) == 2: + torch_cuda_version = torch_parts[1] + _dependency_links.append( + f"https://download.pytorch.org/whl/{torch_cuda_version}" + ) + + if (major, minor) >= (2, 10): + extras_require_map.pop("fbgemm-gpu") + extras_require_map["fbgemm-gpu"] = [ + "fbgemm-gpu==1.5.0", + "fbgemm-gpu-genai==1.5.0", + ] + if not install_xformers: + _install_requires.pop(_install_requires.index(xformers_version)) + extras_require_map["vllm"] = ["vllm>=0.17.1"] + elif (major, minor) >= (2, 9): + extras_require_map.pop("fbgemm-gpu") + extras_require_map["fbgemm-gpu"] = [ + "fbgemm-gpu==1.4.0", + "fbgemm-gpu-genai==1.4.2", + ] + if not install_xformers: + _install_requires.pop(_install_requires.index(xformers_version)) + if patch == 0: + extras_require_map["vllm"] = ["vllm==0.13.0"] + else: + extras_require_map["vllm"] = ["vllm==0.14.0"] + elif (major, minor) >= (2, 8): + extras_require_map.pop("fbgemm-gpu") + extras_require_map["fbgemm-gpu"] = ["fbgemm-gpu-genai==1.3.0"] + extras_require_map["vllm"] = ["vllm==0.11.0"] + if not install_xformers: + _install_requires.pop(_install_requires.index(xformers_version)) + elif (major, minor) >= (2, 7): + _install_requires.pop(_install_requires.index(xformers_version)) + if patch == 0: + if install_xformers: + _install_requires.append("xformers==0.0.30") + # vllm 0.9.x is incompatible with latest transformers + extras_require_map.pop("vllm") + else: + if install_xformers: + _install_requires.append("xformers==0.0.31") + extras_require_map["vllm"] = ["vllm==0.10.1"] + elif (major, minor) >= (2, 6): + _install_requires.pop(_install_requires.index(xformers_version)) + if install_xformers: + _install_requires.append("xformers==0.0.29.post3") + # since we only support 2.6.0+cu126 + _dependency_links.append("https://download.pytorch.org/whl/cu126") + extras_require_map.pop("vllm") + elif (major, minor) >= (2, 5): + _install_requires.pop(_install_requires.index(xformers_version)) + if install_xformers: + if patch == 0: + _install_requires.append("xformers==0.0.28.post2") + else: + _install_requires.append("xformers>=0.0.28.post3") + extras_require_map.pop("vllm") + elif (major, minor) >= (2, 4): + extras_require_map.pop("vllm") + if install_xformers: + if patch == 0: + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers>=0.0.27") + else: + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers==0.0.28.post1") + else: + raise ValueError("axolotl requires torch>=2.4") + + except PackageNotFoundError: + pass + return _install_requires, _dependency_links, extras_require_map + + +def get_package_version(): + with open( + Path(os.path.dirname(os.path.abspath(__file__))) / "VERSION", + "r", + encoding="utf-8", + ) as fin: + version_ = fin.read().strip() + return version_ + + +extras_require = { + "flash-attn": ["flash-attn==2.8.3"], + "ring-flash-attn": [ + "flash-attn==2.8.3", + "ring-flash-attn>=0.1.7", + ], + "deepspeed": [ + "deepspeed==0.18.2", + "deepspeed-kernels", + ], + "mamba-ssm": [ + "mamba-ssm==1.2.0.post1", + "causal_conv1d", + ], + "auto-gptq": [ + "auto-gptq==0.5.1", + ], + "mlflow": [ + "mlflow", + ], + "galore": [ + "galore_torch", + ], + "apollo": [ + "apollo-torch", + ], + "optimizers": [ + "galore_torch", + "apollo-torch", + "lomo-optim==0.1.1", + "torch-optimi==0.2.1", + "came_pytorch==0.1.3", + ], + "ray": [ + "ray[train]>=2.52.1", + ], + "vllm": [ + "vllm==0.10.0", + ], + "llmcompressor": [ + "llmcompressor==0.5.1", + ], + "fbgemm-gpu": ["fbgemm-gpu-genai==1.3.0"], + "opentelemetry": [ + "opentelemetry-api", + "opentelemetry-sdk", + "opentelemetry-exporter-prometheus", + "prometheus-client", + ], +} +install_requires, dependency_links, extras_require_build = parse_requirements( + extras_require +) + +setup( + version=get_package_version(), + package_dir={"": "src"}, + packages=find_packages("src"), + install_requires=install_requires, + dependency_links=dependency_links, + entry_points={ + "console_scripts": [ + "axolotl=axolotl.cli.main:main", + ], + }, + extras_require=extras_require_build, +) diff --git a/src/setuptools_axolotl_dynamic_dependencies.py b/src/setuptools_axolotl_dynamic_dependencies.py new file mode 100644 index 000000000..3bb54cda8 --- /dev/null +++ b/src/setuptools_axolotl_dynamic_dependencies.py @@ -0,0 +1,102 @@ +""" +dynamic requirements for axolotl +""" + +import platform +import re +from importlib.metadata import PackageNotFoundError, version + +from setuptools.command.build_py import build_py as _build_py + + +def parse_requirements(): + _install_requires = [] + _dependency_links = [] + with open("./requirements.txt", encoding="utf-8") as requirements_file: + lines = [r.strip() for r in requirements_file.readlines()] + for line in lines: + is_extras = ( + "flash-attn" in line + or "flash-attention" in line + or "deepspeed" in line + or "mamba-ssm" in line + or "lion-pytorch" in line + ) + if line.startswith("--extra-index-url"): + # Handle custom index URLs + _, url = line.split() + _dependency_links.append(url) + elif not is_extras and line and line[0] != "#": + # Handle standard packages + _install_requires.append(line) + + try: + xformers_version = [req for req in _install_requires if "xformers" in req][0] + torchao_version = [req for req in _install_requires if "torchao" in req][0] + + if "Darwin" in platform.system(): + # don't install xformers on MacOS + _install_requires.pop(_install_requires.index(xformers_version)) + else: + # detect the version of torch already installed + # and set it so dependencies don't clobber the torch version + try: + torch_version = version("torch") + except PackageNotFoundError: + torch_version = "2.5.1" + _install_requires.append(f"torch=={torch_version}") + + version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version) + if version_match: + major, minor, patch = version_match.groups() + major, minor = int(major), int(minor) + patch = ( + int(patch) if patch is not None else 0 + ) # Default patch to 0 if not present + else: + raise ValueError("Invalid version format") + + if (major, minor) >= (2, 5): + _install_requires.pop(_install_requires.index(xformers_version)) + if patch == 0: + _install_requires.append("xformers==0.0.28.post2") + else: + _install_requires.append("xformers==0.0.28.post3") + elif (major, minor) >= (2, 4): + if patch == 0: + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers>=0.0.27") + else: + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers==0.0.28.post1") + elif (major, minor) >= (2, 3): + _install_requires.pop(_install_requires.index(torchao_version)) + if patch == 0: + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers>=0.0.26.post1") + else: + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers>=0.0.27") + elif (major, minor) >= (2, 2): + _install_requires.pop(_install_requires.index(torchao_version)) + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers>=0.0.25.post1") + else: + _install_requires.pop(_install_requires.index(torchao_version)) + _install_requires.pop(_install_requires.index(xformers_version)) + _install_requires.append("xformers>=0.0.23.post1") + + except PackageNotFoundError: + pass + return _install_requires, _dependency_links + + +class BuildPyCommand(_build_py): + """ + custom build_py command to parse dynamic requirements + """ + + def finalize_options(self): + super().finalize_options() + install_requires, _ = parse_requirements() + self.distribution.install_requires = install_requires