Merge pull request #283 from MeetKai/update-training

Update training dependencies installation

functionary/train/README.md

@@ -3,11 +3,14 @@
 # Create new virtual environment
 python3 -m venv venv && source venv/bin/activate
 
+# Install PyTorch 2.4
+pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124
+
 # Install Dependencies
-pip install -e . --index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.org/simple
+pip install -e .
 
 # Install Liger if using liger:
-pip install -e .[liger] --index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.org/simple
+pip install -e .[liger]
 ```
 ### Llama-2 models
 
@@ -158,10 +161,10 @@ For Lora fintuning, you need to install additional requirements:
 
 ```shell
 # To install dependencies for LoRA
-pip install -e .[lora] --index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.org/simple
+pip install -e .[lora]
 
 # To run LoRA finetuning with Liger
-pip install -e .[lora,liger] --index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.org/simple
+pip install -e .[lora,liger]
 ```
 Run script:
 

functionary/train/pyproject.toml

@@ -4,9 +4,6 @@ version = "0.0.1"
 description = "Chat language model that can use tools and interpret the results"
 requires-python = ">=3.9"
 dependencies = [
-    "torch==2.4.0+cu121",
-    "torchvision==0.19.0+cu121",
-    "torchaudio==2.4.0+cu121",
     "accelerate==0.34.0",
     "bitsandbytes==0.44.1",
     "scipy==1.11.3",
@@ -21,7 +18,7 @@ dependencies = [
     "tensorboard==2.15.1",
     "aenum==3.1.15",
     "transformers @ git+https://github.com/huggingface/transformers.git",
-    "flash-attn==v2.6.3",
+    "flash-attn==2.6.3",
     "json_source_map==1.0.5",
 ]