* change python lint in workflow

* change module install script
* move test dir

.github/workflows/python-app.yml

@@ -30,10 +30,7 @@ jobs:
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
     - name: Lint with flake8
       run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+        flake8 . --count --show-source --statistics --max-line-length=127 --max-complexity 15
     - name: Test with pytest
       run: |
         pytest

aspen/__init__.py

@@ -4,7 +4,7 @@
 from aspen.modelargs import TokenizerArgs, LlamaModelArgs, MultiLoraBatchData, LoraBatchDataConfig
 from aspen.dataset import DataSet
 from aspen.modelloader import load_llama_7b_weight, load_llama_tf_weight
-from aspen.modelloader import load_alpaca_lora_7b_weight, load_random_lora_7b_weight
+from aspen.modelloader import load_random_lora_7b_weight
 from aspen.modelloader import save_lora_model
 
 __all__ = [
@@ -20,7 +20,6 @@
     "convert_hf_to_pth",
     "load_llama_7b_weight",
     "load_llama_tf_weight",
-    "load_alpaca_lora_7b_weight",
     "load_random_lora_7b_weight",
     "save_lora_model"
 ]

aspen/dataset.py

@@ -1,7 +1,9 @@
+from aspen.modelargs import MultiLoraBatchData, LoraBatchDataConfig
+from aspen.tokenizer import Tokenizer
+
 import math
 import json
 import random
-from aspen import MultiLoraBatchData, LoraBatchDataConfig, Tokenizer
 from typing import Dict, List, Tuple
 
 
@@ -37,7 +39,9 @@ def __get_lora_text_data(self) -> Dict[str, List[str]]:
                             "{output}", raw_data_output).replace("{instruction}", raw_data_instruction)
                     else:
                         text_data = lora_config["prompt_input"].replace(
-                            "{output}", raw_data_output).replace("{instruction}", raw_data_instruction).replace("{input}", raw_data_input)
+                            "{output}", raw_data_output).replace(
+                            "{instruction}", raw_data_instruction).replace(
+                            "{input}", raw_data_input)
                     lora_text_data[lora_name].append(text_data)
 
         return lora_text_data
@@ -122,8 +126,6 @@ def get_batch_data(self) -> MultiLoraBatchData:
                 f"    epoch: {self.lora_cnt_epochs_[lora_name] + 1} / {self.lora_num_epochs_[lora_name]}")
             print(
                 f"    step : {self.lora_start_idx_[lora_name]} / {len(self.lora_token_data_[lora_name])}")
-            print(
-                f"         : {(self.lora_cnt_epochs_[lora_name] * len(self.lora_token_data_[lora_name]) + self.lora_start_idx_[lora_name]) * 100 / (self.lora_num_epochs_[lora_name] * len(self.lora_token_data_[lora_name]))}%")
 
         # align batch data
         max_token_len = math.ceil(max_token_len / 8) * 8

aspen/model.py

@@ -1,4 +1,4 @@
-from aspen import LlamaModelArgs, MultiLoraBatchData
+from aspen.modelargs import LlamaModelArgs, MultiLoraBatchData
 
 import torch
 import torch.nn.functional as F
@@ -45,7 +45,8 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor:
     return einops.rearrange(x, "... d r -> ... (d r)")
 
 
-def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, angle: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor,
+                     angle: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
     # data shape is: batch_size * max_seq_len * n_head * n_dim
     _, max_seq_len, _, dim_head = xq.shape
 
@@ -177,7 +178,11 @@ def set_lora_parameter(self, adapter_name: str, r: int, lora_alpha: int, lora_dr
                 adapter_name, r, lora_alpha, lora_dropout)
 
     # @torch.compile
-    def forward(self, data: torch.Tensor, mask: torch.Tensor, rope_angle: Tuple[torch.Tensor, torch.Tensor], input_args: MultiLoraBatchData):
+    def forward(self,
+                data: torch.Tensor,
+                mask: torch.Tensor,
+                rope_angle: Tuple[torch.Tensor, torch.Tensor],
+                input_args: MultiLoraBatchData):
         batch_size, max_seq_len, _ = data.shape
 
         attention_norm_data = self.attention_norm_.forward(data)

aspen/modelargs.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List, Dict
+from typing import List
 
 
 @dataclass

aspen/modelloader.py

@@ -1,7 +1,8 @@
+from aspen.model import LlamaModel, Linear, RMSNorm
+
 import sys
 import torch
 
-from aspen import LlamaModel, Linear, RMSNorm
 from transformers import LlamaForCausalLM
 
 
@@ -91,7 +92,12 @@ def load_llama_tf_weight(model: LlamaModel, llama_model_path: str, dev: str):
             print(f"Not use layer {layer_name}.", file=sys.stderr)
 
 
-def load_random_lora_7b_weight(model: LlamaModel, adapter_name: str, r: int, dim: int, target_module: str, device: str) -> None:
+def load_random_lora_7b_weight(model: LlamaModel,
+                               adapter_name: str,
+                               r: int,
+                               dim: int,
+                               target_module: str,
+                               device: str) -> None:
     norm_mean = 0
     norm_std = 1e-3
     target_module_name_list = ["q_proj", "k_proj", "v_proj", "o_proj", "w1_proj", "w2_proj", "w3_proj"]

aspen/utils.py

@@ -1,6 +1,7 @@
 import torch
 from transformers import LlamaForCausalLM
 
+
 # convert huggingface model to pytorch model
 def convert_hf_to_pth(source: str, dest: str):
     src_model = LlamaForCausalLM.from_pretrained(source)

config/lora.json

@@ -1,6 +1,6 @@
 {
-    "base_model": "",
-    "token_model": "",
+    "base_model": "/modules/llama-7b/7B/consolidated.00.pth",
+    "token_model": "/modules/llama-7b/tokenizer.model",
     "cutoff_len": 512,
     "group_by_length": false,
     "expand_right": true,

legacy.py

@@ -1,7 +1,4 @@
-from aspen import LlamaModel, Tokenizer, DataSet
-from aspen import LlamaModelArgs, MultiLoraBatchData
-from aspen import load_llama_7b_weight, load_random_lora_7b_weight
-from aspen import save_lora_model
+import aspen
 
 import json
 import torch
@@ -10,25 +7,25 @@
 with open('config/lora.json', 'r', encoding='utf8') as fp:
     config = json.load(fp)
 
-args = LlamaModelArgs()
-tokenizer = Tokenizer(config["token_model"])
+args = aspen.LlamaModelArgs()
+tokenizer = aspen.Tokenizer(config["token_model"])
 tokenizer.pad_id_ = 0
 args.max_seq_len_ = 4096
 args.device = config["device"]
 args.vocab_size_ = tokenizer.n_words_
 args.pad_id_ = tokenizer.pad_id_
 args.n_heads_ = 32
-llama_model = LlamaModel(args)
+llama_model = aspen.LlamaModel(args)
 
 
 def setup_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
 
 
-def init_lora_model(llama_model: LlamaModel):
+def init_lora_model(llama_model: aspen.LlamaModel):
     for lora_config in config["lora"]:
-        load_random_lora_7b_weight(
+        aspen.load_random_lora_7b_weight(
             llama_model,
             lora_config["name"],
             lora_config["r"],
@@ -42,8 +39,8 @@ def init_lora_model(llama_model: LlamaModel):
 if __name__ == "__main__":
     setup_seed(42)
 
-    data_set = DataSet(config, tokenizer)
-    load_llama_7b_weight(llama_model, config["base_model"], config["device"])
+    data_set = aspen.DataSet(config, tokenizer)
+    aspen.load_llama_7b_weight(llama_model, config["base_model"], config["device"])
     init_lora_model(llama_model)
 
     torch.cuda.empty_cache()
@@ -54,7 +51,7 @@ def init_lora_model(llama_model: LlamaModel):
     while not data_set.check_done():
         optimizer.zero_grad()
         loss_fn = torch.nn.CrossEntropyLoss()
-        input: MultiLoraBatchData = data_set.get_batch_data()
+        input: aspen.MultiLoraBatchData = data_set.get_batch_data()
 
         step_cnt += 1
 
@@ -83,9 +80,7 @@ def init_lora_model(llama_model: LlamaModel):
 
         if step_cnt % config["save_step"] == 0:
             for lora_config in config["lora"]:
-                save_lora_model(
-                    llama_model, lora_config["output"] + f".bin{step_cnt}", lora_config["name"])
+                aspen.save_lora_model(llama_model, lora_config["output"] + f".bin{step_cnt}", lora_config["name"])
 
     for lora_config in config["lora"]:
-        save_lora_model(
-            llama_model, lora_config["output"], lora_config["name"])
+        aspen.save_lora_model(llama_model, lora_config["output"], lora_config["name"])

mlora.py

@@ -14,9 +14,7 @@
 #
 # Copyright (C) 2023 All Rights Reserved.
 #
-# Email:   
 # Github:  https://github.com/TUDB-Labs/multi-lora-fine-tune
-# Website: 
 
 import datetime
 import argparse
@@ -31,17 +29,20 @@
 
 args = parser.parse_args()
 
-def log(msg:str):
+
+def log(msg: str):
     if args.log:
         print('[%s] ASPEN: %s' % (datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), msg))
 
+
 if torch.cuda.is_available():
     log('NVIDIA CUDA initialized successfully.')
     log('Total %i GPU(s) detected.' % torch.cuda.device_count())
 else:
     print('ASPEN requires NVIDIA CUDA computing capacity. Please check your PyTorch installation.')
     exit(-1)
 
+
 def prep_llm():
     args = aspen.LlamaModelArgs()
     tokenizer = aspen.Tokenizer(args.model_name_or_path + os.sep + 'tokenizer.model')
@@ -55,5 +56,6 @@ def prep_llm():
     aspen.load_llama_tf_weight(model, args.model_name_or_path, args.device)
     return tokenizer, model
 
+
 if __name__ == "__main__":
     tokenizer, model = prep_llm()

pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "aspen"
+version = "0.0.1"
+description = "A tool for fine-tuning large language models (LLMs) using the LoRA or QLoRA methods more efficiently."
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "torch==2.0.1",
+    "einops",
+    "xformers",
+    "transformers",
+    "bitsandbytes",
+    "sentencepiece",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/TUDB-Labs/multi-lora-fine-tune"
+"Bug Tracker" = "https://github.com/TUDB-Labs/multi-lora-fine-tune/issues"
+
+[tool.setuptools.packages.find]
+include = ["aspen"]

requirements.txt

@@ -1,5 +1,6 @@
 torch==2.0.1
-sentencepiece
-bitsandbytes
-xformers
 einops
+xformers
+transformers
+bitsandbytes
+sentencepiece

aspen/tests/loader_test.py → tests/loader_test.py

@@ -18,4 +18,4 @@ def test_load_llam2(self):
 
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()