Setup package dependencies (#1068)

.github/azure-gpu-test.yml

@@ -41,7 +41,7 @@ jobs:
       displayName: "Image info & NVIDIA"
 
     - script: |
-        pip install -r requirements-all.txt pytest pytest-rerunfailures transformers>=4.38.0 einops protobuf
+        pip install '.[all,test]'
       displayName: 'Install dependencies'
 
     - bash: |

.github/workflows/cpu-tests.yml

@@ -49,10 +49,9 @@ jobs:
       run: |
         pip install -r requirements.txt
         pip list
-        # make sure all modules are importable
+        # make sure all modules are still importable with only the minimal dependencies available
         modules=$(
-          find * -type f -name "*.py" | \
-          grep -v tests | grep "/" | grep -v lm_eval | grep -v xla | \
+          find litgpt -type f -name "*.py" | \
           sed 's/\.py$//' | sed 's/\//./g' | \
           sed 's/.__init__//g' | xargs -I {} echo "import {};"
         )
@@ -61,15 +60,9 @@ jobs:
 
     - name: Install all dependencies
       run: |
-        pip install -r requirements-all.txt pytest pytest-rerunfailures pytest-timeout transformers>=4.38.0 einops protobuf
+        pip install '.[all,test]'
         pip list
 
-    - name: Run tests without the package installed
-      run: |
-        pytest -v --disable-pytest-warnings --strict-markers --color=yes --timeout 120
-
     - name: Run tests
       run: |
-        pip install . --no-deps
-
         pytest -v --disable-pytest-warnings --strict-markers --color=yes --timeout 120

README.md

@@ -86,17 +86,10 @@ Avoiding code duplication is **not** a goal. **Readability** and **hackability**
 
 ## Setup
 
-Clone the repo:
-
-```bash
-git clone https://github.com/Lightning-AI/litgpt
-cd litgpt
-```
-
 Install with all dependencies (including CLI, quantization, tokenizers for all models, etc.):
 
 ```bash
-pip install -r requirements-all.txt
+pip install 'litgpt[all]'
 ```
 
  

eval/lm_eval_harness.py

@@ -11,12 +11,8 @@
 from lm_eval import base, evaluator, tasks
 from lm_eval.base import BaseLM
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from litgpt.generate.base import generate
 from litgpt import GPT, Config, Tokenizer
+from litgpt.generate.base import generate
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, load_checkpoint
 
 

litgpt/__init__.py

@@ -14,7 +14,7 @@
 if not bool(_LIGHTNING_AVAILABLE):
     raise ImportError(
         "LitGPT requires lightning nightly. Please run:\n"
-        f" pip uninstall -y lightning; pip install -r requirements.txt\n{str(_LIGHTNING_AVAILABLE)}"
+        f" pip uninstall -y lightning; pip install litgpt\n{str(_LIGHTNING_AVAILABLE)}"
     )
 
 # Suppress excessive warnings, see https://github.com/pytorch/pytorch/issues/111632

litgpt/chat/base.py

@@ -9,10 +9,6 @@
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 from litgpt.generate.base import next_token
 from litgpt import GPT, Config, PromptStyle, Tokenizer
 from litgpt.prompts import load_prompt_style, has_prompt_style

litgpt/generate/adapter.py

@@ -9,10 +9,6 @@
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 from litgpt import Tokenizer, PromptStyle
 from litgpt.adapter import GPT, Config
 from litgpt.generate.base import generate

litgpt/generate/adapter_v2.py

@@ -9,10 +9,6 @@
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 from litgpt import Tokenizer, PromptStyle
 from litgpt.adapter_v2 import GPT, Config
 from litgpt.generate.base import generate

litgpt/generate/base.py

@@ -9,10 +9,7 @@
 import torch
 import torch._dynamo.config
 import torch._inductor.config
-
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
+from lightning.fabric.plugins import BitsandbytesPrecision
 
 from litgpt import GPT, Config, Tokenizer, PromptStyle
 from litgpt.prompts import load_prompt_style, has_prompt_style

litgpt/generate/full.py

@@ -9,10 +9,6 @@
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 from litgpt import GPT, Config, Tokenizer, PromptStyle
 from litgpt.generate.base import generate
 from litgpt.prompts import load_prompt_style, has_prompt_style

litgpt/generate/lora.py

@@ -9,10 +9,6 @@
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 from litgpt import Tokenizer, PromptStyle
 from litgpt.generate.base import generate
 from litgpt.lora import GPT, Config, merge_lora_weights

litgpt/generate/sequentially.py

@@ -17,10 +17,6 @@
 from lightning.fabric.utilities.init import _materialize_meta_tensors
 from typing_extensions import Type
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 import litgpt.generate.base as generate_base
 from litgpt import GPT, Config, Tokenizer
 from litgpt.model import Block, build_mask_cache

litgpt/generate/tp.py

@@ -15,10 +15,6 @@
 from lightning.fabric.utilities import rank_zero_only
 from torch.distributed._functional_collectives import all_reduce
 
-# support running without installing as a package
-wd = Path(__file__).parents[2].resolve()
-sys.path.append(str(wd))
-
 import litgpt.generate.base as generate_base
 from litgpt import GPT, Config, Tokenizer
 from litgpt.model import CausalSelfAttention, GptNeoxMLP, LLaMAMLP, LLaMAMoE

requirements-all.txt

@@ -12,5 +12,4 @@ pandas                    # litgpt.data.prepare_starcoder.py
 pyarrow                   # litgpt.data.prepare_starcoder.py
 tensorboard               # litgpt.pretrain
 torchmetrics              # litgpt.pretrain
-# eval
-git+https://github.com/EleutherAI/lm-evaluation-harness.git@115206dc89dad67b8beaa90051fb52db77f0a529
+lm_eval @ git+https://github.com/EleutherAI/lm-evaluation-harness.git@115206dc89dad67b8b

setup.py

@@ -1,24 +1,26 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
-import os
+from pathlib import Path
 
 from setuptools import find_packages, setup
 
-_PATH_ROOT = os.path.dirname(__file__)
-
-with open(os.path.join(_PATH_ROOT, "README.md"), encoding="utf-8") as fo:
-    readme = fo.read()
+root = Path(__file__).parent
+readme = (root / "README.md").read_text()
+requirements = (root / "requirements.txt").read_text().split("\n")
+requirements_all = (root / "requirements-all.txt").read_text().split("\n")
+requirements_all = [r for r in requirements_all if r and not r.strip().startswith("-r")]
 
 setup(
     name="litgpt",
     version="0.1.0",
     description="Open source large language model implementation",
     author="Lightning AI",
     url="https://github.com/lightning-AI/litgpt",
-    install_requires=[
-        "torch>=2.2.0",
-        "lightning @ git+https://github.com/Lightning-AI/lightning@f23b3b1e7fdab1d325f79f69a28706d33144f27e",
-    ],
+    install_requires=requirements,
+    extras_require={
+        "all": requirements_all,
+        "test": ["pytest", "pytest-rerunfailures", "pytest-timeout", "transformers>=4.38.0", "einops", "protobuf"],
+    },
     packages=find_packages(),
     long_description=readme,
     long_description_content_type="text/markdown",

tutorials/pretrain_tinyllama.md

@@ -49,7 +49,7 @@ In order to start pretraining litgpt on it, you need to read, tokenize, and writ
 First, install additional dependencies for preprocessing:
 
 ```bash
-pip install -r requirements-all.txt
+pip install '.[all]'
 ```
 
 You will need to have the tokenizer config available:

xla/README.md

@@ -50,7 +50,7 @@ Once inside the machine, clone the repository and install the dependencies:
 ```shell
 git clone https://github.com/Lightning-AI/litgpt
 cd litgpt
-pip install -r requirements.txt
+pip install .
 ```
 
 Install Optimized BLAS:

xla/finetune/adapter.py

@@ -14,13 +14,14 @@
 from lightning.fabric.strategies import XLAFSDPStrategy
 from lightning.fabric.utilities import ThroughputMonitor, measure_flops
 
+from litgpt.adapter import GPT, Block, Config, adapter_filter, mark_only_adapter_as_trainable
+from litgpt.tokenizer import Tokenizer
+from litgpt.utils import check_valid_checkpoint_dir, chunked_cross_entropy, estimate_flops, lazy_load, num_parameters
+
 # support running without installing as a package
 wd = Path(__file__).parent.parent.parent.resolve()
 sys.path.append(str(wd))
 
-from litgpt.adapter import GPT, Block, Config, adapter_filter, mark_only_adapter_as_trainable
-from litgpt.tokenizer import Tokenizer
-from litgpt.utils import check_valid_checkpoint_dir, chunked_cross_entropy, estimate_flops, lazy_load, num_parameters
 from xla.scripts.prepare_alpaca import generate_prompt
 from xla.generate.base import generate
 from xla.utils import rank_print, sequential_load_and_fsdp_wrap

xla/generate/adapter.py

@@ -8,16 +8,16 @@
 import lightning as L
 from lightning.fabric.accelerators import XLAAccelerator
 from lightning.fabric.strategies import XLAFSDPStrategy
+
+from litgpt import Tokenizer
+from litgpt.adapter import GPT, Block, Config
 from litgpt.prompts import Alpaca
+from litgpt.utils import check_valid_checkpoint_dir, lazy_load
 
 # support running without installing as a package
 wd = Path(__file__).parent.parent.parent.resolve()
 sys.path.append(str(wd))
 
-from litgpt import Tokenizer
-from litgpt.adapter import GPT, Block, Config
-from litgpt.utils import check_valid_checkpoint_dir, lazy_load
-
 from xla.generate.base import generate
 from xla.utils import rank_print
 

xla/generate/base.py

@@ -11,13 +11,14 @@
 from lightning.fabric.accelerators import XLAAccelerator
 from lightning.fabric.strategies import XLAFSDPStrategy
 
+from litgpt import GPT, Config, Tokenizer
+from litgpt.model import Block
+from litgpt.utils import check_valid_checkpoint_dir, lazy_load
+
 # support running without installing as a package
 wd = Path(__file__).parent.parent.parent.resolve()
 sys.path.append(str(wd))
 
-from litgpt import GPT, Config, Tokenizer
-from litgpt.model import Block
-from litgpt.utils import check_valid_checkpoint_dir, lazy_load
 from xla.utils import rank_print
 
 

xla/scripts/prepare_alpaca.py

@@ -3,21 +3,15 @@
 """Implementation derived from https://github.com/tloen/alpaca-lora"""
 
 import json
-import sys
 from pathlib import Path
 from typing import Optional
 
 import torch
 from lightning_utilities.core.imports import RequirementCache
-from torch.utils.data import random_split
-from tqdm import tqdm
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
 from lit_gpt.tokenizer import Tokenizer
 from lit_gpt.utils import CLI
+from torch.utils.data import random_split
+from tqdm import tqdm
 
 
 def prepare(