Skip to content

Commit

Permalink
* change python lint in workflow
Browse files Browse the repository at this point in the history
* change module install script
* move test dir
  • Loading branch information
yezhengmao1 committed Aug 29, 2023
1 parent f5972ca commit 4b0fefb
Show file tree
Hide file tree
Showing 14 changed files with 79 additions and 65 deletions.
5 changes: 1 addition & 4 deletions .github/workflows/python-app.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions aspen/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__ = [
Expand 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"
]
10 changes: 6 additions & 4 deletions aspen/dataset.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
11 changes: 8 additions & 3 deletions aspen/model.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from aspen import LlamaModelArgs, MultiLoraBatchData
from aspen.modelargs import LlamaModelArgs, MultiLoraBatchData

import torch
import torch.nn.functional as F
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion aspen/modelargs.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from dataclasses import dataclass
from typing import List, Dict
from typing import List


@dataclass
Expand Down
10 changes: 8 additions & 2 deletions aspen/modelloader.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -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"]
Expand Down
1 change: 1 addition & 0 deletions aspen/utils.py
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
4 changes: 2 additions & 2 deletions config/lora.json
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
27 changes: 11 additions & 16 deletions legacy.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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"],
Expand All @@ -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()
Expand All @@ -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

Expand Down Expand Up @@ -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"])
8 changes: 5 additions & 3 deletions mlora.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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')
Expand All @@ -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()
30 changes: 30 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -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"]
7 changes: 4 additions & 3 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
torch==2.0.1
sentencepiece
bitsandbytes
xformers
einops
xformers
transformers
bitsandbytes
sentencepiece
24 changes: 0 additions & 24 deletions setup.py

This file was deleted.

2 changes: 1 addition & 1 deletion aspen/tests/loader_test.py → tests/loader_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,4 @@ def test_load_llam2(self):


if __name__ == '__main__':
unittest.main()
unittest.main()

0 comments on commit 4b0fefb

Please sign in to comment.