Add dynamic dependency management for CUDA and ROCm (#460)

## 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]>

.github/workflows/amd-ci.yml

@@ -39,7 +39,8 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install flake8 isort black
+        pip install --no-deps .[fmt]
+
     - name: Run checkstyle
       run: make checkstyle
 
@@ -58,7 +59,8 @@ jobs:
 
     - name: Setup Dependencies
       run: |
-        python -m pip install -e .[dev,amd] --extra-index-url https://download.pytorch.org/whl/nightly/rocm6.2
+        python -m pip install --upgrade pip
+        pip install -e .[dev] --extra-index-url https://download.pytorch.org/whl/nightly/rocm6.2
     
     - name: List Python Environments
       run: python -m pip list

.github/workflows/nvi-ci.yml

@@ -38,7 +38,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install flake8 isort black
+        pip install --no-deps .[fmt]
 
     - name: Run checkstyle
       run: make checkstyle

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 site/
 .cache/
 .venv/
+venv/
 .ipynb_checkpoints/
 
 # Misc
@@ -16,4 +17,4 @@ dist/
 uv.lock
 
 # Benchmark images
-benchmark/visualizations
+benchmark/visualizations

README.md

@@ -146,11 +146,13 @@ To install from source:
 ```bash
 git clone https://github.com/linkedin/Liger-Kernel.git
 cd Liger-Kernel
+
+# Install Default Dependencies
+# Setup.py will detect whether you are using AMD or NVIDIA
 pip install -e .
-# or if installing on amd platform
-pip install -e .[amd] --extra-index-url https://download.pytorch.org/whl/nightly/rocm6.2 # rocm6.2
-# or if using transformers
-pip install -e .[transformers]
+
+# Setup Development Dependencies
+pip install -e ".[dev]"
 ```
 
 

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=42", "wheel"]
+requires = ["setuptools>=42", "wheel", "setuptools-scm"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -9,53 +9,18 @@ description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }
-dependencies = [
-    "torch>=2.1.2",
-    "triton>=2.3.1",
-]
-
-
-[project.optional-dependencies]
-transformers = [
-    "transformers~=4.0"
-]
-
-trl = [
-    "trl>=0.11.0",
-]
-
-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",
-    "torchvision>=0.16.2",
-    "seaborn",
-]
-
-amd = [
-    "torch>=2.6.0.dev",
-    "setuptools-scm>=8",
-    "torchvision>=0.20.0.dev",
-    "triton>=3.0.0",
-]
-
+dynamic = ["dependencies", "optional-dependencies"]
 
+[tool.setuptools]
+package-dir = {"" = "src"}
 
 [tool.setuptools.packages.find]
 where = ["src"]
-include = ["liger_kernel", "liger_kernel.*"]
+include = ["liger_kernel*"]
+namespaces = false
 
 [tool.pytest.ini_options]
-pythonpath = [
-  "src",
-  "."
-]
+pythonpath = ["src", "."]
 asyncio_mode = "auto"
 log_cli = true
-log_cli_level = "INFO"
+log_cli_level = "INFO"

setup.py

@@ -0,0 +1,71 @@
+# 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(),
+)