mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Stefan HeandByronHsu
b735584899
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Closes https://github.com/linkedin/Liger-Kernel/issues/451 This PR implements dynamic dependency management that automatically detects the user's hardware platform (CUDA, ROCm, or CPU) and installs the appropriate dependencies. This ensures a smoother installation experience across different hardware configurations. ### Key Changes - Added `_get_platform()` function to detect hardware platform using `nvidia-smi` or `rocm-smi` - Implemented `get_default_dependencies()` to return platform-specific core dependencies - Implemented `get_optional_dependencies()` to handle dev and transformers extras - Updated `pyproject.toml` to use dynamic dependencies from setup.py ### Platform-specific Dependencies - **CUDA/CPU**: - Core: torch>=2.1.2, triton>=2.3.1 - Full dev environment including transformers, testing tools, and visualization libraries - **ROCm**: - Core: torch>=2.6.0.dev, triton>=3.0.0, plus ROCm-specific requirements - Streamlined dev dependencies focused on essential tools <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ### setup.py vs pyproject.toml #### setup.py - Traditional Python package configuration - Allows dynamic configuration via Python code - Needed for complex/conditional dependencies - Supports custom build steps #### pyproject.toml - Modern, static configuration (PEP 517/518) - Cleaner, declarative format - Works well with modern tools (poetry, hatch) - Preferred for simple packages #### Using Both - `pyproject.toml`: Build system specs, basic metadata - `setup.py`: Dynamic dependencies, custom build logic - Common in projects needing both static config and dynamic install - ## Testing Done (In Progress) <!--- This is a required section; please describe how this change was tested. ---> 1. Clone the repository 2. Test installation on different platforms: ```bash # Basic installation pip install -e . # With dev tools pip install -e .[dev] # With transformers pip install -e .[dev,transformers] ``` - [x] CPU - [x] CUDA - NVIDIA Tested with following config - NVIDIA-SMI 535.129.03 - Driver Version: 535.129.03 - CUDA Version: 12.2 - [ ] ROCm - AMD - Need help from the community to test on ROCm platform <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> ## Hardware Type: - NVIDIA-SMI 535.129.03 - CUDA Version: 12.2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: ByronHsu <[email protected]>
72 lines
1.8 KiB
Python
72 lines
1.8 KiB
Python
# setup.py
|
|
|
|
import subprocess
|
|
from typing import Literal
|
|
|
|
from setuptools import setup
|
|
|
|
|
|
def get_default_dependencies():
|
|
"""Determine the appropriate dependencies based on detected hardware."""
|
|
platform = get_platform()
|
|
|
|
if platform in ["cuda", "cpu"]:
|
|
return [
|
|
"torch>=2.1.2",
|
|
"triton>=2.3.1",
|
|
]
|
|
elif platform == "rocm":
|
|
return [
|
|
"torch>=2.6.0.dev",
|
|
"triton>=3.0.0",
|
|
]
|
|
|
|
|
|
def get_optional_dependencies():
|
|
"""Get optional dependency groups."""
|
|
return {
|
|
"dev": [
|
|
"transformers>=4.44.2",
|
|
"matplotlib>=3.7.2",
|
|
"flake8>=4.0.1.1",
|
|
"black>=24.4.2",
|
|
"isort>=5.13.2",
|
|
"pytest>=7.1.2",
|
|
"pytest-xdist",
|
|
"pytest-rerunfailures",
|
|
"datasets>=2.19.2",
|
|
"seaborn",
|
|
],
|
|
"fmt": ["flake8", "isort", "black"],
|
|
}
|
|
|
|
|
|
# TODO: add intel XPU
|
|
def get_platform() -> Literal["cuda", "rocm", "cpu"]:
|
|
"""
|
|
Detect whether the system has NVIDIA or AMD GPU without torch dependency.
|
|
"""
|
|
# Try nvidia-smi first
|
|
try:
|
|
subprocess.run(["nvidia-smi"], check=True)
|
|
print("NVIDIA GPU detected")
|
|
return "cuda"
|
|
except (subprocess.SubprocessError, FileNotFoundError):
|
|
# If nvidia-smi fails, check for ROCm
|
|
try:
|
|
subprocess.run(["rocm-smi"], check=True)
|
|
print("ROCm GPU detected")
|
|
return "rocm"
|
|
except (subprocess.SubprocessError, FileNotFoundError):
|
|
print("No GPU detected")
|
|
return "cpu"
|
|
|
|
|
|
setup(
|
|
name="liger_kernel",
|
|
package_dir={"": "src"},
|
|
packages=["liger_kernel"],
|
|
install_requires=get_default_dependencies(),
|
|
extras_require=get_optional_dependencies(),
|
|
)
|