Upload 33 files

README.md

@@ -1,3 +1,101 @@
----
-license: cc-by-4.0
----
+# CDKA: Component Designed Kronecker Adapters
+
+Official PyTorch implementation of **Diving into Kronecker Adapters: Component Design Matters**.
+
+CDKA is a parameter-efficient fine-tuning framework for large-scale models. It studies how the component design of Kronecker adapters affects adapter capacity and downstream performance. In particular, CDKA exposes the component dimensions and the number of Kronecker components as controllable hyperparameters:
+
+$$
+\Delta W = \sum_{i=1}^{r} B^{(i)} \otimes A^{(i)},
+$$
+
+where `r1`, `r2`, and `r` determine the shape and number of Kronecker components.
+
+
+## Installation
+
+### 1. Clone the repository
+
+```bash
+git clone https://github.com/rainstonee/CDKA.git
+cd CDKA
+```
+
+### 2. Create a Python environment
+
+```bash
+conda create -n cdka python=3.10 -y
+conda activate cdka
+```
+
+### 3. Install dependencies
+
+```bash
+pip install --upgrade pip
+pip install -r requirements.txt
+```
+
+### 4. Install the modified PEFT package
+
+This repository includes a modified PEFT implementation in `peft.zip`. The custom PEFT version is required because CDKA extends the standard LoRA configuration with Kronecker-adapter-specific arguments such as `r1`, `r2`, and `r`.
+
+```bash
+unzip peft.zip
+pip install -e peft
+```
+
+## Training with CDKA
+
+Run the example script:
+
+```bash
+bash run.sh
+```
+
+The default example runs the Llama-2-7B model on `meta_math` dataset with PEFT enabled and Kronecker adapter hyperparameters specified through Hydra overrides.
+
+The script is equivalent to:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python run_exp.py \
+  +model=llama \
+  +peft=all \
+  +init=default \
+  +dataset_name=meta_math \
+  +seed=333 \
+  ++peft.lora_r1=2 \
+  ++peft.lora_r2=2 \
+  ++peft.lora_r=8 \
+  ++peft.lora_alpha=64
+```
+
+### Recommended default component configuration
+
+For most CDKA experiments, we recommend starting with the following default component configuration:
+
+```bash
+++peft.lora_r1=2 \
+++peft.lora_r2=8 \
+++peft.lora_r=4 \
+++peft.lora_alpha=64
+```
+
+This setting follows the main component-design principle of CDKA: use a small `r1`, a large `r2`, and a moderate number of Kronecker components `r`. In practice, this provides a balanced default configuration before task-specific tuning.
+
+
+## Citation
+
+If you find this repository useful, please cite:
+
+```bibtex
+@article{bai2026diving,
+  title={Diving into Kronecker Adapters: Component Design Matters},
+  author={Bai, Jiayu and Yu, Danchen and Liao, Zhenyu and Hou, TianQi and Zhou, Feng and Qiu, Robert C and Ling, Zenan},
+  journal={arXiv preprint arXiv:2602.01267},
+  year={2026}
+}
+```
+
+## Acknowledgement
+
+This codebase is built on [LoRA-GA](https://github.com/Outsider565/LoRA-GA). We thank the authors of LoRA-GA for releasing their implementation.
+

conf/config.yaml

@@ -0,0 +1,12 @@
+defaults:
+  - _self_
+  # - model: llama
+dry_run: False
+wandb:
+  project: "glue_new"
+  name: null
+hydra:
+  sweeper:
+    params:
+      ++dataset_name: mnli
+      ++seed: 0

conf/init/default.yaml

@@ -0,0 +1,7 @@
+mode: simple
+lora_A: kaiming
+lora_B: zeros
+# lora_A: zeros
+# lora_B: kaiming
+# This is the default way to initialize the model parameters.
+dtype: fp32

conf/init/gradient.yaml

@@ -0,0 +1,8 @@
+mode: gradient
+bsz: 1
+iters: 8
+direction: ArBr
+max_length: 1024
+dtype: fp32
+scale: stable
+stable_gamma: 16

conf/init/kaiming.yaml

@@ -0,0 +1,5 @@
+mode: simple
+lora_A: kaiming
+lora_B: kaiming
+dtype: fp32
+stable_gamma: 16

conf/init/qr.yaml

@@ -0,0 +1,2 @@
+mode: qr
+dtype: fp32

conf/init/svd.yaml

@@ -0,0 +1,5 @@
+mode: svd
+weight: default
+dtype: fp32
+stable_gamma: 16
+scale: default

conf/model/llama.yaml

@@ -0,0 +1,11 @@
+name: llama/llama-2-7b-hf
+type: CausalLM
+epochs: 1
+per_device_batch_size: 1
+real_batch_size: 32
+bf16: True
+eval_epochs: 1
+early_stopping_patience: 3
+max_length: 1024
+logging_steps: 1
+learning_rate: 2e-4

conf/model/llama3.yaml

@@ -0,0 +1,11 @@
+name: llama/llama-3.1-8b-hf
+type: CausalLM
+epochs: 1
+per_device_batch_size: 1
+real_batch_size: 64
+bf16: True
+eval_epochs: 1
+early_stopping_patience: 3
+max_length: 1024
+logging_steps: 1
+learning_rate: 2e-4

conf/model/t5base.yaml

@@ -0,0 +1,11 @@
+name: t5-base
+type: ConditionalGeneration
+epochs: 1
+per_device_batch_size: 32
+real_batch_size: 32
+bf16: False
+eval_epochs: 1
+early_stopping_patience: 5
+max_length: 128
+logging_steps: 1
+learning_rate: 2e-3

conf/model/t5large.yaml

@@ -0,0 +1,11 @@
+name: t5-large
+type: ConditionalGeneration
+epochs: 1
+per_device_batch_size: 32
+real_batch_size: 1
+bf16: True
+eval_epochs: 0.1
+early_stopping_patience: 3
+max_length: 128
+logging_steps: 1
+learning_rate: 5e-5

conf/model/t5small.yaml

@@ -0,0 +1,11 @@
+name: t5-small
+type: ConditionalGeneration
+epochs: 1
+per_device_batch_size: 32
+real_batch_size: 32
+bf16: True
+eval_epochs: 0.1
+early_stopping_patience: 3
+max_length: 128
+logging_steps: 1
+learning_rate: 5e-5

conf/model/t5xl.yaml

@@ -0,0 +1,11 @@
+name: google/flan-t5-xl
+type: ConditionalGeneration
+epochs: 5
+per_device_batch_size: 1
+real_batch_size: 32
+bf16: True
+eval_epochs: 1
+early_stopping_patience: 3
+max_length: 512
+logging_steps: 1
+learning_rate: 5e-5

conf/model/tinyllama.yaml

@@ -0,0 +1,11 @@
+name: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
+type: CausalLM
+epochs: 1
+per_device_batch_size: 1
+real_batch_size: 32
+bf16: True
+eval_epochs: 1
+early_stopping_patience: 3
+max_length: 512
+logging_steps: 1
+learning_rate: 1e-5

conf/peft/adalora.yaml

@@ -0,0 +1,10 @@
+use_peft: true
+lora_r: null
+lora_relative_r: null
+lora_target_modules: all
+train_embeddings: false
+use_rslora: false
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 16
+adalora: true

conf/peft/all.yaml

@@ -0,0 +1,11 @@
+use_peft: true
+lora_r1: null
+lora_r2: null
+lora_r: null
+lora_relative_r: null
+lora_target_modules: all
+train_embeddings: false
+use_rslora: True
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 64

conf/peft/all_w_emb.yaml

@@ -0,0 +1,9 @@
+use_peft: true
+lora_r: null
+lora_relative_r: null
+lora_target_modules: all
+train_embeddings: true
+use_rslora: false
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 16

conf/peft/att.yaml

@@ -0,0 +1,11 @@
+use_peft: true
+lora_r1: null
+lora_r2: null
+lora_r: null
+lora_relative_r: null
+lora_target_modules: [q_proj,k_proj,v_proj,o_proj]
+train_embeddings: false
+use_rslora: True
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 64

conf/peft/dora.yaml

@@ -0,0 +1,10 @@
+use_peft: true
+lora_r: null
+lora_relative_r: null
+lora_target_modules: all
+train_embeddings: false
+use_rslora: false
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 16
+dora: true

conf/peft/full_ft.yaml

@@ -0,0 +1,8 @@
+use_peft: false
+lora_r: null
+lora_relative_r: null
+lora_target_modules: null
+train_embeddings: false
+use_rslora: false
+use_loraplus: false
+loraplus_lr_ratio: 16

conf/peft/qv.yaml

@@ -0,0 +1,9 @@
+use_peft: true
+lora_r: null
+lora_relative_r: null
+lora_target_modules: [q,v]
+train_embeddings: false
+use_rslora: false
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 16

conf/peft/qv_w_emb.yaml

@@ -0,0 +1,9 @@
+use_peft: true
+lora_r: null
+lora_relative_r: null
+lora_target_modules: [q,v]
+train_embeddings: true
+use_rslora: false
+use_loraplus: false
+loraplus_lr_ratio: 16
+lora_alpha: 16

data.py

@@ -0,0 +1,488 @@
+from datasets import load_dataset, Dataset
+import typing as tp
+import functools
+import os
+import pickle
+import logging
+import datasets
+
+log = logging.getLogger(__name__)
+
+def cache_to_disk(root_datadir):
+    def decorator_cache(func):
+        @functools.wraps(func)
+        def wrapper_cache(*args, **kwargs):
+            if not os.path.exists(root_datadir):
+                os.makedirs(root_datadir)
+
+            func_name = func.__name__.replace("/", "")
+            cache_file = os.path.join(root_datadir, f"{func_name}.pkl")
+            if os.path.exists(cache_file):
+                with open(cache_file, "rb") as f:
+                    log.info(f"Loading cached data for {func.__name__}")
+                    return pickle.load(f)
+
+            result = func(*args, **kwargs)
+            with open(cache_file, "wb") as f:
+                pickle.dump(result, f)
+                log.info(f"Cached data for {func.__name__}")
+            return result
+
+        return wrapper_cache
+
+    return decorator_cache
+
+@cache_to_disk("data_cache")
+def load_emo():
+    dataset = load_dataset("emo")
+    label_map = {0: "others", 1: "happy", 2: "sad", 3: "angry"}
+    instruction = "classify the emotion of the text: "
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["text"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    test_set = dataset["test"]
+    return train_set, test_set, test_set
+
+@cache_to_disk("data_cache")
+def load_sst2():
+    dataset = load_dataset("glue", "sst2")
+    instruction = "classify the sentiment of the text: "
+    label_map = {0: "negative", 1: "positive", -1: "other"}
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["sentence"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_cola():
+    dataset = load_dataset("glue", "cola")
+    instruction = "classify the grammaticality of the text: "
+    label_map = {0: "unacceptable", 1: "acceptable", -1: "other"}
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["sentence"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_qqp():
+    dataset = load_dataset("glue", "qqp")
+    instruction = "classify the semantic similarity of the text: "
+    label_map = {0: "different", 1: "duplicate", -1: "other"}
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["question1"]}\n{e["question2"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_mrpc():
+    dataset = load_dataset("glue", "mrpc")
+    instruction = "classify the semantic similarity of the text: "
+    label_map = {0: "different", 1: "equivalent", -1: "other"}
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["sentence1"]}\n{e["sentence2"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_mnli():
+    dataset = load_dataset("glue", "mnli",download_mode="force_redownload")
+    instruction = "classify the semantic similarity of the text: "
+    label_map = {0: "entailment", 1: "neutral", 2: "contradiction", -1: "other"}
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["premise"]}\n{e["hypothesis"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation_matched"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_squad():
+    dataset = load_dataset("rajpurkar/squad")
+    instruction = "answer the question: "
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["question"]}\ncontext: {e["context"]}\nresult: ',
+            "y": ", ".join(e["answers"]["text"]),
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_qnli():
+    dataset = load_dataset("glue", "qnli")
+    instruction = "classify the semantic similarity of the question and the sentence: "
+    label_map = {0: "entailment", 1: "not_entailment", -1: "other"}
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'{instruction}{e["question"]}\n{e["sentence"]}\nresult: ',
+            "y": label_map[e["label"]],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["validation"]
+    test_set = dataset["test"]
+    return train_set, validation_set, test_set
+
+
+template_with_input = '''### Instruction:
+{instruction}
+
+### Input:
+{input}
+
+### Response:
+'''
+
+template_wo_input = '''Below is an instruction that describes a task. Write a response that appropriately completes the request.
+
+### Instruction:
+{instruction}
+
+### Response:
+'''
+
+@cache_to_disk("data_cache")
+def load_alpaca():
+    dataset = load_dataset("tatsu-lab/alpaca")
+    def alpaca_preprocess(instruction, input, output):
+        if input == "":
+            x = template_wo_input.format(instruction=instruction)
+        else:
+            x = template_with_input.format(instruction=instruction, input=input)
+        return {"x": x, "y": output}
+    dataset = dataset.map(
+        lambda e: alpaca_preprocess(e["instruction"], e["input"], e["output"])
+    )
+    # we sample 10% of the training set as validation set
+    train_set = dataset["train"].train_test_split(test_size=0.1)['train']
+    validation_set = dataset["train"].train_test_split(test_size=0.1)['test']
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_gsm8k():
+    dataset = load_dataset("gsm8k", "main")
+    #x = "Q: " + x[0] + "\n" + "A:"
+    dataset = dataset.map(
+        lambda e: {
+            "x": f'Q: {e["question"]}\nA: ',
+            "y": e["answer"],
+        }
+    )
+    train_set = dataset["train"]
+    validation_set = dataset["test"]
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_alpaca_gpt4():
+    dataset = load_dataset("tatsu-lab/alpaca")
+    def alpaca_preprocess(instruction, input, output):
+        if input == "":
+            x = template_wo_input.format(instruction=instruction)
+        else:
+            x = template_with_input.format(instruction=instruction, input=input)
+        return {"x": x, "y": output}
+    dataset = dataset.map(
+        lambda e: alpaca_preprocess(e["instruction"], e["input"], e["output"])
+    )
+    # we sample 10% of the training set as validation set
+    train_set = dataset["train"].train_test_split(test_size=0.1)['train']
+    validation_set = dataset["train"].train_test_split(test_size=0.1)['test']
+    return train_set, validation_set, validation_set
+
+@cache_to_disk("data_cache")
+def load_flan():
+    dataset = load_dataset("Muennighoff/flan", split='train', streaming=True)
+    def preprocess(data):
+        return {
+            "x": template_wo_input.format(instruction=data['inputs']),
+            "y": data['targets'],
+        }
+    train_samples = []
+    eval_samples = []
+    count = 0
+    dataset.shuffle(buffer_size=5000, seed=42)
+    from tqdm import tqdm
+    for sample in tqdm(dataset, total=110000):
+        processed_sample = preprocess(sample)
+        if count < 100000:  # First 100,000 samples for training
+            train_samples.append(processed_sample)
+        elif 100000 <= count < 110000:  # Next 10,000 samples for evaluation
+            eval_samples.append(processed_sample)
+        elif count >= 110000:  # Stop processing after collecting enough samples
+            break
+        count += 1
+    # convert to hf dataset
+    train_set = Dataset.from_list(train_samples)
+    eval_set = Dataset.from_list(eval_samples)
+    return train_set, eval_set, eval_set
+
+@cache_to_disk("data_cache")
+def load_meta_math(max_tokens=512):
+    dataset = load_dataset("meta-math/MetaMathQA", split='train')
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+    def preprocess(data):
+        return {
+            "x": f'Q: {data["query"]}\nA: ',
+            "y": data["response"].split("\nThe answer is:")[0]
+        }
+    train_samples = []
+    eval_samples = []
+    count = 0
+    dataset.shuffle(seed=42)
+    from tqdm import tqdm
+    bar = tqdm(dataset, total=110000)
+    total = 0
+    ok = 0
+    for sample in dataset:
+        total += 1
+        temp = preprocess(sample)
+        if len(tokenizer(temp['x']+' '+temp['y'])['input_ids']) >= max_tokens or "GSM" not in sample["type"]:
+            continue
+        bar.update(1)
+        bar.set_description(f"ok: {ok}/{total}")
+        ok += 1
+        processed_sample = preprocess(sample)
+        if count < 100000:  # First 100,000 samples for training
+            train_samples.append(processed_sample)
+        elif 100000 <= count < 110000:  # Next 10,000 samples for evaluation
+            eval_samples.append(processed_sample)
+        elif count >= 110000:  # Stop processing after collecting enough samples
+            break
+        count += 1
+    # convert to hf dataset
+    train_set = Dataset.from_list(train_samples)
+    eval_set = Dataset.from_list(eval_samples)
+    return train_set, eval_set, eval_set
+
+@cache_to_disk("data_cache")
+def load_flan_v2(max_tokens=512):
+    dataset = load_dataset("SirNeural/flan_v2", split='train', streaming=True)
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+    def preprocess(data):
+        return {
+            "x": data['inputs'],
+            "y": data['targets'],
+        }
+    train_samples = []
+    eval_samples = []
+    count = 0
+    dataset.shuffle(buffer_size=5000, seed=42)
+    from tqdm import tqdm
+    bar = tqdm(dataset, total=110000)
+    total = 0
+    ok = 0
+    for sample in dataset:
+        total += 1
+        temp = preprocess(sample)
+        if len(tokenizer(temp['x']+' '+temp['y'])['input_ids']) >= max_tokens:
+            continue
+        bar.update(1)
+        bar.set_description(f"ok: {ok}/{total}")
+        ok += 1
+        processed_sample = preprocess(sample)
+        if count < 100000:  # First 100,000 samples for training
+            train_samples.append(processed_sample)
+        elif 100000 <= count < 110000:  # Next 10,000 samples for evaluation
+            eval_samples.append(processed_sample)
+        elif count >= 110000:  # Stop processing after collecting enough samples
+            break
+        count += 1
+    # convert to hf dataset
+    train_set = Dataset.from_list(train_samples)
+    eval_set = Dataset.from_list(eval_samples)
+    return train_set, eval_set, eval_set
+
+@cache_to_disk("data_cache")
+def load_codefeedback(max_tokens=1024):
+    dataset = load_dataset("m-a-p/CodeFeedback-Filtered-Instruction", split='train')
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("llama/llama-2-7b-hf")
+    def preprocess(data):
+        y = data['answer']
+        y = "```".join(y.split("```")[:2]) + "```" # only keep the first code block
+        return {
+            "x": template_wo_input.format(
+                instruction=data['query']
+            ),
+            "y": y,
+        }
+    train_samples = []
+    eval_samples = []
+    count = 0
+    dataset.shuffle(seed=42)
+    from tqdm import tqdm
+    bar = tqdm(dataset, total=110000)
+    total = 0
+    ok = 0
+    for sample in dataset:
+        total += 1
+        temp = preprocess(sample)
+        if "```" not in sample['answer']:
+            continue
+        if len(tokenizer(temp['x']+' '+temp['y'])['input_ids']) >= max_tokens:
+            continue
+        bar.update(1)
+        bar.set_description(f"ok: {ok}/{total}")
+        ok += 1
+        processed_sample = preprocess(sample)
+        if count < 100000:
+            train_samples.append(processed_sample)
+        elif 100000 <= count < 110000:
+            eval_samples.append(processed_sample)
+        elif count >= 110000:  # Stop processing after collecting enough samples
+            break
+        count += 1
+    # convert to hf dataset
+    train_set = Dataset.from_list(train_samples)
+    eval_set = Dataset.from_list(eval_samples)
+    return train_set, eval_set, eval_set
+
+@cache_to_disk("data_cache")
+def load_wizardlm(max_tokens=1024):
+    dataset = load_dataset("silk-road/Wizard-LM-Chinese-instruct-evol", split='train')
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("llama/llama-2-7b-hf")
+    def preprocess(data):
+        y = data['output']
+        return {
+            "x": template_wo_input.format(
+                instruction=data['instruction']
+            ),
+            "y": y,
+        }
+    train_samples = []
+    eval_samples = []
+    count = 0
+    dataset.shuffle(seed=42)
+    from tqdm import tqdm
+    bar = tqdm(dataset, total=70000)
+    total = 0
+    ok = 0
+    for sample in dataset:
+        total += 1
+        temp = preprocess(sample)
+        if "sorry" in temp['y'].lower() or "as an ai" in temp['y'].lower():
+            continue
+        if len(tokenizer(temp['x']+' '+temp['y'])['input_ids']) >= max_tokens:
+            continue
+        bar.update(1)
+        bar.set_description(f"ok: {ok}/{total}")
+        ok += 1
+        processed_sample = temp
+        if count < 52000:
+            train_samples.append(processed_sample)
+        elif 52000 <= count < 70000:
+            eval_samples.append(processed_sample)
+        elif count >= 70000:  # Stop processing after collecting enough samples
+            break
+        count += 1
+    # convert to hf dataset
+    train_set = Dataset.from_list(train_samples)
+    eval_set = Dataset.from_list(eval_samples)
+    return train_set, eval_set, eval_set
+
+
+@cache_to_disk("data_cache")
+def load_common(max_tokens=1024):
+    # dataset = load_dataset("zwhe99/commonsense_170k", split='train')
+    dataset = load_dataset("json", data_files="commonsense_170k.json")['train']
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("llama/llama-2-7b-hf")
+    def preprocess(data):
+        y = data['output']
+        return {
+            "x": template_wo_input.format(
+                instruction=data['instruction']
+            ),
+            "y": y,
+        }
+    i = 0
+    train_samples = []
+    eval_samples = []
+    for sample in dataset:
+        i += 1
+        temp = preprocess(sample)
+        # print(temp)
+        if len(tokenizer(temp['x']+' '+temp['y'])['input_ids']) >= max_tokens:
+            continue
+        processed_sample = temp
+        train_samples.append(processed_sample)
+        if i == 1:
+            eval_samples.append(processed_sample)
+    # convert to hf dataset
+    train_set = Dataset.from_list(train_samples)
+    eval_set = Dataset.from_list(eval_samples)
+    return train_set, eval_set, eval_set
+
+DATASET_MAP = {
+    "sst2": load_sst2,
+    "cola": load_cola,
+    "qqp": load_qqp,
+    "mrpc": load_mrpc,
+    "mnli": load_mnli,
+    "emo": load_emo,
+    "squad": load_squad,
+    "alpaca": load_alpaca,
+    "qnli": load_qnli,
+    "gsm8k": load_gsm8k,
+    "alpaca_gpt4": load_alpaca_gpt4,
+    "flan": load_flan,
+    "flan_v2": load_flan_v2,
+    "meta_math": load_meta_math,
+    "codefeedback": load_codefeedback,
+    "wizard_lm": load_wizardlm,
+    "common": load_common,
+}
+
+
+if __name__ == "__main__":
+    # for dataset in [load_emo, load_sst2, load_cola, load_qqp, load_mrpc, load_mnli]:
+    #     train_set, val_set, test_set = dataset()
+    #     print(train_set[0])
+    #     print(val_set[0])
+    #     print(test_set[0])
+    #     print()
+    # print(load_alpaca())
+    # for name, dataset in DATASET_MAP.items():
+    #     train_set, val_set, test_set = dataset()
+    #     print(name)
+    #     print(train_set[0])
+    #     print(val_set[0])
+    #     print(test_set[0])
+    #     print()   
+    x, r, _ = load_common()
+    print(x[0]['x'])
+    print(x[0]['y'])
+    print(len(x))
+    print(len(r))

eval_gsm8k.py

@@ -0,0 +1,58 @@
+from data import load_gsm8k
+from utils import model_inference, initialize_text_to_text_model
+from fire import Fire
+import re
+import os
+from tqdm import tqdm
+from peft import PeftModel
+
+def extract_num(text):
+    # Regex pattern to find the number following '####'
+    pattern = r'####\s*(\d+)'
+    # Using re.search to find the first match
+    match = re.search(pattern, text)
+    if match:
+        result = match.group(1)
+        print(text)
+    else:
+        print(text)
+        result = ""
+    try:
+        return int(result.replace(",", ""))
+    except:
+        print(f"'{result}' can't be converted")
+        return 0
+
+def main(model_name = "llama/llama-2-7b-hf"):
+    _, _, test_set = load_gsm8k()
+    model_type = "CausalLM"
+    model, tokenizer = initialize_text_to_text_model(
+        model_name, model_type, True, flash_attention=True
+    )
+    model = PeftModel.from_pretrained(model, "checkpoint_dir")
+    model.generation_config.pad_token_id = tokenizer.pad_token_id
+    all = 0
+    correct = 0
+    t = tqdm(test_set)
+    for example in t:
+        # print(example['x'])
+        pred_text = model_inference(model, tokenizer, example['x'], model_type, max_target_length=512)
+        gt = extract_num(example["y"])
+        pred = extract_num(pred_text)
+        correct += int(gt == pred)
+        all += 1
+        t.set_description(f"Accuracy: {correct / all * 100:02f}%")
+
+    print("Acc:", correct / all)
+    # append to gsm8k_results.txt (create if not exists)
+    if not os.path.exists("gsm8k_results.txt"):
+        with open("gsm8k_results.txt", "w") as f:
+            f.write("Model Acc\n")
+    with open("gsm8k_results.txt", "a") as f:
+        f.write(f"{model_name} {correct / all}\n")
+
+if __name__ == "__main__":
+    Fire(main)
+
+
+

eval_humaneval.py

@@ -0,0 +1,88 @@
+from human_eval.data import write_jsonl, read_problems
+from fire import Fire
+from tqdm import trange, tqdm
+from utils import initialize_text_to_text_model, model_inference
+import re
+import os
+from human_eval.evaluation import evaluate_functional_correctness
+from peft import PeftModel
+
+ALPACA_PREFIX_TEMPLATE_MD = """Below is an instruction that describes a task.\n Write a response that appropriately completes the request.
+
+### Instruction:
+Complete the following Python code: 
+Notes: respond with the entire complete function definition
+do not add any comments, be as concise in your code as possible
+use only built-in libraries, assume no additional imports other than those provided (if any)
+use `    ` (4 spaces) for each level of indentation
+
+code:
+```python
+{PROMPT}
+```
+
+### Response:
+```python
+"""
+
+def post_process(text):
+    text = text.replace("```", "")
+    text = text.replace("\t", "    ")
+    text = re.sub(r'(""".*?"""|\'\'\'.*?\'\'\')', '', text, flags=re.DOTALL)
+    text = "\n".join([ll.rstrip() for ll in text.splitlines() if ll.strip()])
+    lines = text.split("\n")
+    spaces_for_each_line = []
+    for line in lines:
+        match = re.match(r'^( *)', line)
+        if match:
+            leading_spaces = len(match.group(1))
+            spaces_for_each_line.append(leading_spaces)
+    try:
+        def_line = [i for i, line in enumerate(lines) if "def" in line][0]
+        def_line_space = spaces_for_each_line[def_line]
+    except:
+        print("No def line found")
+        print(text)
+        def_line_space = 0
+    rank_unique_spaces = sorted(list(set(spaces_for_each_line)))
+    indentation_level = {}
+    i = 0
+    for space in rank_unique_spaces:
+        if space <= def_line_space:
+            indentation_level[space] = 0
+        else:
+            i += 1
+            indentation_level[space] = i
+    new_lines = []
+    for line, space in zip(lines, spaces_for_each_line):
+        new_lines.append("    " * indentation_level[space] + line.lstrip())
+    return "\n".join(new_lines)
+
+def generate_one_completion(model, tokenizer, model_type, prompt, template=True):
+    if template:
+        prompt_in = ALPACA_PREFIX_TEMPLATE_MD.format(PROMPT=prompt)
+    pred_text = model_inference(model, tokenizer, prompt_in, model_type, max_target_length=512)
+    post_pred = post_process(pred_text)
+    return post_pred
+
+
+
+
+def humaneval(model, tokenizer, save_dir, model_type = "CausalLM", model_name="llama/llama-2-7b-hf"):
+
+    import os
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    problems = read_problems()
+    num_samples_per_task = 1
+    samples = [
+        dict(task_id=task_id, completion=generate_one_completion(model, tokenizer, model_type, problems[task_id]["prompt"]))
+        for task_id in tqdm(problems, desc="Tasks")
+        for _ in range(num_samples_per_task)
+    ]
+
+    target_name = os.path.join(save_dir, f"{model_name.replace('/', '_')}_humaneval_samples.jsonl")
+    write_jsonl(target_name, samples)
+    results = evaluate_functional_correctness(target_name, [1])
+    print(results)
+

logTrainer.py

@@ -0,0 +1,238 @@
+from dataclasses import dataclass, field
+from functools import reduce
+from typing import Callable, Dict, List, Optional, Tuple, Union, Any
+
+import numpy as np
+import torch
+import wandb
+import torch.nn as nn
+from torch.utils.data import Dataset
+
+from transformers import Trainer, Seq2SeqTrainingArguments
+from transformers.data.data_collator import DataCollator
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.trainer import (
+    EvalPrediction,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    TrainerCallback,
+)
+from transformers.trainer_pt_utils import get_parameter_names
+from transformers.utils import is_sagemaker_mp_enabled, logging
+from peft.tuners.lora.layer import Linear as LoraLinear
+
+# include_keywords = ["block.0", "block.4"]
+include_keywords = ["encoder.block.2","encoder.block.3","encoder.block.4"]  # for T5
+# include_keywords = ["layers.27", "layers.6"]  # for Llama
+do_log = False
+
+
+def get_forward_hook(name):
+    def hook(module, input, output):
+        wandb.log(
+            {
+                f"{name}/input_mean": input[0].mean().item(),
+                f"{name}/input_std": input[0].std().item(),
+                f"{name}/output_mean": output.mean().item(),
+                f"{name}/output_std": output.std().item(),
+            },
+            commit=False,
+        )
+    return hook
+
+class LogTrainer(Trainer):
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: Seq2SeqTrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Optional[Callable[[], PreTrainedModel]] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (
+            None,
+            None,
+        ),
+        preprocess_logits_for_metrics: Optional[
+            Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+        ] = None,
+    ):
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+        self.is_peft = "PeftModel" in type(model).__name__
+        if self.is_peft:
+            for name, module in model.named_modules():
+                if isinstance(module, LoraLinear):
+                    self.scaling = module.scaling["default"]
+                    break
+        self.orig_A = None
+        self.orig_B = None
+        self.orig_W = None
+        self.gradient_accumulation_counter = 0
+
+    def training_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
+    ) -> torch.Tensor:
+        if not do_log:
+            return super().training_step(model, inputs)
+        if self.is_peft:
+            if self.orig_A is None:
+                self.orig_A = {}
+                self.orig_B = {}
+                for name, param in model.named_parameters():
+                    if param.requires_grad and any(
+                        [kw in name for kw in include_keywords]
+                    ):
+                        if "lora_A" in name:
+                            self.orig_A[name.split("lora_A.")[0]] = (
+                                param.detach().clone()
+                            )
+                        elif "lora_B" in name:
+                            self.orig_B[name.split("lora_B.")[0]] = (
+                                param.detach().clone()
+                            )
+                for name, module in model.named_modules():
+                    if any([kw in name for kw in include_keywords]) and isinstance(module, LoraLinear):
+                        breakpoint()
+                        hook = get_forward_hook(name)
+                        module.register_forward_hook(hook)
+        else:
+            if self.orig_W is None:
+                self.orig_W = {}
+                for name, param in model.named_parameters():
+                    if param.requires_grad and any(
+                        [kw in name for kw in include_keywords]
+                    ):
+                        self.orig_W[name] = param.detach().clone()
+
+        model.train()
+        inputs = self._prepare_inputs(inputs)
+
+        with self.compute_loss_context_manager():
+            loss = self.compute_loss(model, inputs)
+
+        if self.args.n_gpu > 1:
+            loss = loss.mean()  # mean() to average on multi-gpu parallel training
+
+        self.accelerator.backward(loss)
+        with torch.no_grad():
+            if (
+                self.gradient_accumulation_counter
+                % self.args.gradient_accumulation_steps
+                == self.args.gradient_accumulation_steps - 1
+            ):
+                if self.is_peft:
+                    A_dict = {}
+                    B_dict = {}
+                    for name, param in model.named_parameters():
+                        if param.requires_grad and any(
+                            [kw in name for kw in include_keywords]
+                        ):
+                            if "lora_A" in name:
+                                A_dict[name.split("lora_A.")[0]] = param
+                            elif "lora_B" in name:
+                                B_dict[name.split("lora_B.")[0]] = param
+                    assert (
+                        len(A_dict)
+                        == len(self.orig_A)
+                        == len(B_dict)
+                        == len(self.orig_B)
+                    ), (
+                        len(A_dict),
+                        len(self.orig_A),
+                        len(B_dict),
+                        len(self.orig_B),
+                    )
+                    for key in A_dict.keys():
+                        A = A_dict[key]
+                        B = B_dict[key]
+                        lora_r = A.shape[0]
+                        A_grad = A_dict[key].grad
+                        B_grad = B_dict[key].grad
+                        A_0 = self.orig_A[key]
+                        B_0 = self.orig_B[key]
+                        A_diff = A - A_0
+                        B_diff = B - B_0
+                        BA = torch.matmul(B, A)
+                        BA_0 = torch.matmul(B_0, A_0)
+                        BA_diff = BA - BA_0
+                        BA_diff_norm = torch.norm(BA_diff).item()
+                        A_diff_norm = torch.norm(A_diff).item()
+                        B_diff_norm = torch.norm(B_diff).item()
+                        A_norm = torch.norm(A).item()
+                        B_norm = torch.norm(B).item()
+                        A_grad_norm = torch.norm(A_grad).item()
+                        B_grad_norm = torch.norm(B_grad).item()
+                        # BA_singular_values = torch.svd(BA_diff.float(), compute_uv=False).S[:lora_r]
+                        BA_singular_values = torch.svd_lowrank(
+                            BA_diff.float(), q=2 * lora_r
+                        )[1][:lora_r]
+                        top_1_ratio = (
+                            BA_singular_values[0] / BA_singular_values.sum()
+                        ).item()
+                        top_4_ratio = (
+                            BA_singular_values[:4].sum() / BA_singular_values.sum()
+                        ).item()
+                        wandb.log(
+                            {
+                                f"A_norm/{key}": A_norm,
+                                f"B_norm/{key}": B_norm,
+                                f"A_grad_norm/{key}": A_grad_norm,
+                                f"B_grad_norm/{key}": B_grad_norm,
+                                f"A_diff_norm/{key}": A_diff_norm,
+                                f"B_diff_norm/{key}": B_diff_norm,
+                                f"BA_diff_norm/{key}": BA_diff_norm,
+                                f"scaled_BA_diff_norm/{key}": self.scaling
+                                * BA_diff_norm,
+                                f"BA_top_1_ratio/{key}": top_1_ratio,
+                                f"BA_top_4_ratio/{key}": top_4_ratio,
+                                "train/global_step": self.state.global_step,
+                            }
+                        )
+                else:
+                    W_dict = {}
+                    for name, param in model.named_parameters():
+                        if (
+                            param.requires_grad
+                            and any([kw in name for kw in include_keywords])
+                            and len(param.shape) == 2
+                        ):
+                            W_dict[name] = param
+                    for key in W_dict.keys():
+                        W = W_dict[key]
+                        W_grad = W.grad
+                        W_0 = self.orig_W[key]
+                        W_diff = W - W_0
+                        W_diff_norm = torch.norm(W_diff).item()
+                        W_norm = torch.norm(W).item()
+                        W_grad_norm = torch.norm(W_grad).item()
+                        U, S, V = torch.svd(W_diff.float())
+                        top_1_ratio = S[0] / S.sum()
+                        top_4_ratio = S[:4].sum() / S.sum()
+                        wandb.log(
+                            {
+                                f"W_norm/{key}": W_norm,
+                                f"W_grad_norm/{key}": W_grad_norm,
+                                f"W_diff_norm/{key}": W_diff_norm,
+                                "train/global_step": self.state.global_step,
+                                f"W_top_1_ratio/{key}": top_1_ratio.item(),
+                                f"W_top_4_ratio/{key}": top_4_ratio.item(),
+                            }
+                        )
+        self.gradient_accumulation_counter += 1
+
+        return loss.detach() / self.args.gradient_accumulation_steps

lora_plus.py

@@ -0,0 +1,210 @@
+from dataclasses import dataclass, field
+from functools import reduce
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+
+from peft.tuners import lora
+from transformers import Trainer, Seq2SeqTrainingArguments
+from transformers.data.data_collator import DataCollator
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.trainer import (EvalPrediction, PreTrainedModel,
+                                  PreTrainedTokenizerBase, TrainerCallback)
+from transformers.trainer_pt_utils import get_parameter_names
+from transformers.utils import is_sagemaker_mp_enabled, logging
+from logTrainer import LogTrainer
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class LoraPlusTrainingArguments(Seq2SeqTrainingArguments):
+    loraplus_lr_ratio: Optional[float] = field(
+        default=None, metadata={"help": "loraplus learning rate ratio lr_B / lr_A."}
+    )
+    loraplus_lr_embedding: Optional[float] = field(
+        default=1e-6,
+        metadata={"help": "loraplus learning rate for lora embedding layers."},
+    )
+
+
+def get_module(name, opt_model):
+    """
+    Retrieve a module from a model using its parameter name.
+    Args:
+        name (str): Full name of the parameter, typically including module path.
+        opt_model (torch.nn.Module): The model from which to retrieve the module.
+
+    Returns:
+        Module corresponding to the given name.
+    """
+    parent_idx = 2 if "lora" in name else 1
+    module_names = name.split(sep=".")[:-parent_idx]
+    module = reduce(getattr, module_names, opt_model)
+    return module
+
+
+def create_loraplus_optimizer(
+    opt_model,
+    optimizer_cls,
+    optimizer_kwargs,
+    loraplus_lr_ratio,
+    loraplus_lr_embedding=None,
+):
+    """
+    Creates an optimizer for the given model, applying LoRA-specific learning rate adjustments to different parameter groups.
+
+    Args:
+        opt_model (torch.nn.Module): The model for which the optimizer is being created.
+        optimizer_cls (class): The class of the optimizer to be used (e.g., torch.optim.Adam).
+        optimizer_kwargs (dict): A dictionary of keyword arguments for the optimizer's initialization.
+        loraplus_lr_ratio (float): The learning rate ratio to be applied to LoRA parameters.
+        loraplus_lr_embedding (float, optional): A specific learning rate for embedding parameters, with a default value if not provided.
+
+    Returns:
+        An instance of the specified optimizer class configured with the model's parameters organized into groups with custom learning rates.
+    """
+
+    assert loraplus_lr_ratio is not None, "loraplus_lr_ratio must be provided."
+
+    if loraplus_lr_embedding is None:
+        loraplus_lr_embedding = 1e-6
+
+    decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
+    decay_parameters = [name for name in decay_parameters if "bias" not in name]
+    param_groups = {
+        "groupA": {},
+        "groupB": {},
+        "groupB_no_decay": {},
+        "embedding": {},
+    }
+
+    for name, param in opt_model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        module = get_module(name, opt_model)
+        if isinstance(module, lora.Embedding):
+            param_groups["embedding"][name] = param
+        elif "lora_B" in name or param.ndim == 1:
+            if name in decay_parameters:
+                param_groups["groupB"][name] = param
+            else:
+                param_groups["groupB_no_decay"][name] = param
+        else:
+            param_groups["groupA"][name] = param
+
+    assigned_param_groups = ""
+    for group in param_groups:
+        assigned_param_groups += f"{group}\n {list(param_groups[group].keys())}\n\n"
+    logger.debug(assigned_param_groups)
+
+    lr = optimizer_kwargs["lr"]
+    weight_decay = optimizer_kwargs.get("weight_decay", 0.0)
+
+    optimizer_grouped_parameters = [
+        {
+            "params": list(param_groups["groupA"].values()),
+            "weight_decay": weight_decay,
+            "lr": lr,
+        },
+        {
+            "params": list(param_groups["embedding"].values()),
+            "weight_decay": weight_decay,
+            "lr": loraplus_lr_embedding,
+        },
+        {
+            "params": list(param_groups["groupB"].values()),
+            "weight_decay": weight_decay,
+            "lr": lr * loraplus_lr_ratio,
+        },
+        {
+            "params": list(param_groups["groupB_no_decay"].values()),
+            "weight_decay": 0.0,
+            "lr": lr * loraplus_lr_ratio,
+        },
+    ]
+
+    optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+    if optimizer_cls.__name__ == "Adam8bit":
+        import bitsandbytes
+
+        manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+
+        skipped = 0
+        for module in opt_model.modules():
+            if isinstance(module, nn.Embedding):
+                skipped += sum(
+                    {p.data_ptr(): p.numel() for p in module.parameters()}.values()
+                )
+                logger.info(f"skipped {module}: {skipped/2**20}M params")
+                manager.register_module_override(module, "weight", {"optim_bits": 32})
+                logger.debug(f"bitsandbytes: will optimize {module} in fp32")
+        logger.info(f"skipped: {skipped/2**20}M params")
+
+    return optimizer
+
+
+class LoraPlusTrainer(LogTrainer):
+    def __init__(
+        self,
+        model: Union[PreTrainedModel, nn.Module] = None,
+        args: LoraPlusTrainingArguments = None,
+        data_collator: Optional[DataCollator] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        model_init: Optional[Callable[[], PreTrainedModel]] = None,
+        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+        callbacks: Optional[List[TrainerCallback]] = None,
+        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (
+            None,
+            None,
+        ),
+        preprocess_logits_for_metrics: Optional[
+            Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+        ] = None,
+    ):
+        assert isinstance(
+            args, LoraPlusTrainingArguments
+        ), "args must be of type LoraPlusTrainingArguments"
+        super().__init__(
+            model,
+            args,
+            data_collator,
+            train_dataset,
+            eval_dataset,
+            tokenizer,
+            model_init,
+            compute_metrics,
+            callbacks,
+            optimizers,
+            preprocess_logits_for_metrics,
+        )
+
+    def create_optimizer(self):
+        """
+        Overrides the method to create an optimizer with LoRA+ specific adjustments.
+        """
+        if self.args.loraplus_lr_ratio is None:
+            return super().create_optimizer()
+
+        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
+        if self.optimizer is None:
+            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
+                self.args
+            )
+
+            loraplus_lr_ratio = getattr(self.args, "loraplus_lr_ratio", None)
+            loraplus_lr_embedding = getattr(self.args, "loraplus_lr_embedding", None)
+            self.optimizer = create_loraplus_optimizer(
+                opt_model,
+                optimizer_cls,
+                optimizer_kwargs,
+                loraplus_lr_ratio,
+                loraplus_lr_embedding,
+            )
+
+        return self.optimizer

peft.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a2008a24660d3a4e89498528727c16cad435e3549cae135ad3ca25afdcebefd
+size 6737683

requirements.txt

@@ -0,0 +1,149 @@
+accelerate==0.30.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+annotated-types==0.6.0
+anthropic==0.25.8
+antlr4-python3-runtime==4.9.3
+anyio==4.3.0
+asttokens==2.4.1
+attrs==23.2.0
+cachetools==5.3.3
+calflops==0.2.9
+certifi==2024.2.2
+charset-normalizer==3.3.2
+click==8.1.7
+comm==0.2.2
+datasets==2.19.1
+debugpy==1.8.1
+decorator==5.1.1
+dill==0.3.8
+distro==1.9.0
+dnspython==2.6.1
+docker-pycreds==0.4.0
+einops==0.8.0
+email_validator==2.1.1
+executing==2.0.1
+fastapi==0.111.0
+fastapi-cli==0.0.3
+filelock==3.14.0
+fire==0.6.0
+flash-attn==2.5.8
+frozenlist==1.4.1
+fsspec==2024.3.1
+gitdb==4.0.11
+GitPython==3.1.43
+h11==0.14.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.0
+hydra-core==1.3.2
+idna==3.7
+ipykernel==6.29.4
+ipython==8.24.0
+jedi==0.19.1
+Jinja2==3.1.4
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+jupyter_client==8.6.1
+jupyter_core==5.7.2
+markdown-it-py==3.0.0
+markdown2==2.4.13
+MarkupSafe==2.1.5
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.0.8
+multidict==6.0.5
+multiprocess==0.70.16
+nest-asyncio==1.6.0
+networkx==3.3
+nh3==0.2.17
+ninja==1.11.1.1
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-ml-py==12.535.161
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.1.105
+nvitop==1.3.2
+omegaconf==2.3.0
+openai==0.28.1
+orjson==3.10.3
+packaging==24.0
+pandas==2.2.2
+parso==0.8.4
+peft==0.10.0
+pexpect==4.9.0
+pillow==10.3.0
+platformdirs==4.2.1
+prompt-toolkit==3.0.43
+protobuf==4.25.3
+psutil==5.9.8
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyarrow==16.0.0
+pyarrow-hotfix==0.6
+pydantic==2.7.1
+pydantic_core==2.18.2
+Pygments==2.18.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.1
+pyzmq==26.0.3
+ray==2.21.0
+referencing==0.35.1
+regex==2024.5.10
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.1
+safetensors==0.4.3
+scipy==1.13.0
+sentencepiece==0.2.0
+sentry-sdk==2.1.1
+setproctitle==1.3.3
+shellingham==1.5.4
+shortuuid==1.0.13
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+stack-data==0.6.3
+starlette==0.37.2
+svgwrite==1.4.3
+sympy==1.12
+termcolor==2.4.0
+tiktoken==0.6.0
+tokenizers==0.19.1
+torch==2.3.0
+torchaudio==2.3.0
+torchprofile==0.0.4
+torchvision==0.18.0
+tornado==6.4
+tqdm==4.66.4
+traitlets==5.14.3
+transformers==4.44.0
+triton==2.3.0
+typer==0.12.3
+typing_extensions==4.11.0
+tzdata==2024.1
+ujson==5.9.0
+urllib3==2.2.1
+uvicorn==0.29.0
+uvloop==0.19.0
+wandb==0.17.0
+watchfiles==0.21.0
+wavedrom==2.0.3.post3
+wcwidth==0.2.13
+websockets==12.0
+xxhash==3.4.1
+yarl==1.9.4

run.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0 python run_exp.py +model=llama +peft=all ++peft.lora_r1=2 ++peft.lora_r2=2 ++peft.lora_r=8 ++peft.lora_alpha=64 +dataset_name=meta_math +init=default +seed=333

run_exp.py

@@ -0,0 +1,705 @@
+from peft import get_peft_model, LoraConfig, AdaLoraConfig, TaskType
+import os
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from utils import (
+    train_text_to_text_model,
+    model_inference,
+    initialize_text_to_text_model,
+    transform_dataset,
+    merge_llama,
+)
+import json
+import math
+from datasets import load_dataset
+import wandb
+from data import *
+from typing import List
+import torch
+from copy import deepcopy
+import logging
+from tqdm import tqdm, trange
+from typing import Tuple, List, Dict
+from peft.tuners.lora.layer import Linear as LoraLinear
+from split import rebuild
+import re
+import itertools
+import matplotlib.pyplot as plt
+from commonsense_evaluate import common_evaluate
+from eval_humaneval import humaneval
+# from eval_mtbench import evaluate_mtbench_from_model
+log = logging.getLogger(__name__)
+
+s = 0
+
+def kron(A, B):
+    return (A[:, None, :, None] * B[None, :, None, :]).reshape(A.shape[0] * B.shape[0], A.shape[1] * B.shape[1])
+
+def modified_gram_schmidt(W, eps=1e-12):
+    """
+    Modified Gram–Schmidt QR
+    W: (m, n)
+    Returns:
+        Q: (m, n)
+        R: (n, n)
+    """
+    m, n = W.shape
+    Q = W.clone()
+    R = torch.zeros(n, n, device=W.device, dtype=W.dtype)
+
+    for i in range(n):
+        R[i, i] = torch.norm(Q[:, i])
+        if R[i, i] < eps:
+            raise RuntimeError("Linearly dependent columns")
+
+        Q[:, i] = Q[:, i] / R[i, i]
+
+        for j in range(i + 1, n):
+            R[i, j] = torch.dot(Q[:, i], Q[:, j])
+            Q[:, j] = Q[:, j] - R[i, j] * Q[:, i]
+
+    return Q, R
+
+def seed_everything(seed: int):
+    import random, os
+    import numpy as np
+    import torch
+
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True
+
+
+def find_all_linear_modules(model) -> List[str]:
+    r"""
+    Finds all available modules to apply lora.
+    """
+    linear_cls = torch.nn.Linear
+
+    output_layer_names = ["lm_head", "embed_tokens"]
+
+    module_names = set()
+    for name, module in model.named_modules():
+        if isinstance(module, linear_cls) and not any(
+            [output_layer in name for output_layer in output_layer_names]
+        ):
+            module_names.add(name.split(".")[-1])
+    return list(module_names)
+
+
+def find_hidden_state_size(model):
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            return min(module.weight.shape)
+    return None
+
+
+def calculate_gain(
+    nonlinearity, param
+) -> float:
+    linear_fns = [
+        "linear",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+        "conv_transpose1d",
+        "conv_transpose2d",
+        "conv_transpose3d",
+    ]
+    if nonlinearity in linear_fns or nonlinearity == "sigmoid":
+        return 1
+    elif nonlinearity == "tanh":
+        return 5.0 / 3
+    elif nonlinearity == "relu":
+        return math.sqrt(2.0)
+    elif nonlinearity == "leaky_relu":
+        if param is None:
+            negative_slope = 0.01
+        elif (
+            not isinstance(param, bool)
+            and isinstance(param, int)
+            or isinstance(param, float)
+        ):
+            # True/False are instances of int, hence check above
+            negative_slope = param
+        else:
+            raise ValueError(f"negative_slope {param} not a valid number")
+        return math.sqrt(2.0 / (1 + negative_slope**2))
+    elif nonlinearity == "selu":
+        return (
+            3.0 / 4
+        )  # Value found empirically (https://github.com/pytorch/pytorch/pull/50664)
+    else:
+        raise ValueError(f"Unsupported nonlinearity {nonlinearity}")
+
+def kaimings(weight, a=math.sqrt(5), fan=4096):
+    nonlinearity = "leaky_relu"
+    generator = None
+    gain = calculate_gain(nonlinearity, a)
+    std = gain / math.sqrt(fan)
+    bound = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+    with torch.no_grad():
+        return weight.uniform_(-bound, bound, generator=generator)
+
+@torch.no_grad()
+def reinit_lora_modules(name, module, init_config, **kwargs):
+    r"""
+    Reinitialize the lora model with the given configuration.
+    """
+    lora_r1 = kwargs["lora_r1"]
+    lora_r2 = kwargs["lora_r2"]
+    lora_r = kwargs["lora_r"]
+    # lora_r1 = min(module.lora_A.default.weight.shape)
+    # lora_r2 = min(module.lora_B.default.weight.shape)
+    a_dim = max(module.lora_A.default.weight.shape)
+    b_dim = max(module.lora_B.default.weight.shape)
+    if init_config.mode == "simple":
+        match init_config.lora_A:
+            case "gaussian":
+                torch.nn.init.normal_(
+                    module.lora_A.default.weight, mean=0.0, std=init_config.lora_A_std
+                )
+            case "kaiming":
+                # https://github.com/microsoft/LoRA/blob/a0a92e0f26c067cf94747bdbf1ce73793fa44d19/loralib/layers.py#L124
+                torch.nn.init.kaiming_uniform_(module.lora_A.default.weight, a=math.sqrt(5))
+            case "kaimings":
+                kaimings(module.lora_A.default.weight, a=math.sqrt(5), fan=module.weight.size(1))
+            case "fan_out_kaiming":
+                torch.nn.init.kaiming_normal_(
+                    module.lora_A.default.weight, mode="fan_out"
+                )
+            case "xavier":
+                torch.nn.init.xavier_normal_(module.lora_A.default.weight)
+            case "zeros":
+                torch.nn.init.zeros_(module.lora_A.default.weight)
+            case "unit":
+                torch.nn.init.normal_(
+                    module.lora_A.default.weight, mean=0.0, std=1.0 / (a_dim**0.5)
+                )
+            case "orthogonal":
+                torch.nn.init.orthogonal_(module.lora_A.default.weight)
+            case _:
+                raise ValueError(f"Unknown lora_A initialization: {init_config.lora_A}")
+        match init_config.lora_B:
+            case "gaussian":
+                torch.nn.init.normal_(
+                    module.lora_B.default.weight, mean=0.0, std=init_config.lora_B_std
+                )
+            case "kaiming":
+                torch.nn.init.kaiming_normal_(module.lora_B.default.weight.T, a=math.sqrt(5))
+            case "fan_out_kaiming":
+                torch.nn.init.kaiming_normal_(
+                    module.lora_B.default.weight, mode="fan_out"
+                )
+            case "xavier":
+                torch.nn.init.xavier_normal_(module.lora_B.default.weight)
+            case "zeros":
+                torch.nn.init.zeros_(module.lora_B.default.weight)
+            case "unit":
+                torch.nn.init.normal_(
+                    module.lora_B.default.weight, mean=0.0, std=1.0 / (b_dim**0.5)
+                )
+            case "orthogonal":
+                torch.nn.init.orthogonal_(module.lora_B.default.weight)
+            case _:
+                raise ValueError(f"Unknown lora_B initialization: {init_config.lora_B}")
+        if init_config.get("scale", "") == "stable":
+            gamma = init_config.stable_gamma
+            #module.lora_B.default.weight.data *= (m**0.25) / gamma**0.5
+            #module.lora_A.default.weight.data *= (n**0.25) / gamma**0.5
+            #module.lora_B.default.weight.data *= (m**0.25)
+            #module.lora_A.default.weight.data *= (n**0.25)
+            module.lora_B.default.weight.data *= 1
+            module.lora_A.default.weight.data *= 1
+
+
+    elif init_config.mode == "svd":
+        U, S, V = torch.svd_lowrank(module.weight.float(), q=4 * lora_r, niter=4)
+        V = V.T
+        m, n = module.weight.shape
+        if init_config.scale == "default":
+            S = S / module.scaling["default"]
+            module.lora_B.default.weight = torch.nn.Parameter(
+                (U[:, :lora_r] * torch.sqrt(S[:lora_r])).contiguous()
+            )
+            module.lora_A.default.weight = torch.nn.Parameter(
+                (V[:lora_r, :].T * torch.sqrt(S[:lora_r])).T.contiguous()
+            )
+        elif init_config.scale == "stable":
+            gamma = init_config.stable_gamma
+            module.lora_B.default.weight = torch.nn.Parameter(
+                (U[:, :lora_r] * (m**0.25) / gamma**0.5).contiguous()
+            )
+            module.lora_A.default.weight = torch.nn.Parameter(
+                (V[:lora_r, :] * (n**0.25) / gamma**0.5).contiguous()
+            )
+        elif init_config.scale == "unit":
+            module.lora_B.default.weight = torch.nn.Parameter(
+                (U[:, :lora_r]).contiguous()
+            )
+            module.lora_A.default.weight = torch.nn.Parameter(
+                (V[:lora_r, :]).contiguous()
+            )
+        elif init_config.scale == "normalized":
+            S_sum = S[:lora_r].sum()
+            module.lora_B.default.weight = torch.nn.Parameter(
+                (U[:, :lora_r] * torch.sqrt(S[:lora_r])/torch.sqrt(S_sum)*lora_r**0.5).contiguous()
+            )
+            module.lora_A.default.weight = torch.nn.Parameter(
+                (V[:lora_r, :].T * torch.sqrt(S[:lora_r])/torch.sqrt(S_sum)*lora_r**0.5).T.contiguous()
+            )
+
+    elif init_config.mode == "qr":
+        W = module.weight.float()
+        k,d = W.shape
+        Q, R = torch.linalg.qr(W, mode="reduced")
+        diag = torch.sign(torch.diag(R))
+        diag[diag == 0] = 1.0
+
+        D = torch.diag(diag)
+
+        Q = Q @ D
+        R = D @ R
+        print(torch.min(torch.diag(R)))
+        lambda_vals = torch.abs(torch.diag(R))
+        perm = torch.argsort(lambda_vals, descending=True)
+
+        I1 = perm[:lora_r2] 
+        I2 = perm[lora_r2:lora_r1+lora_r2]  
+        Q1 = Q[:, I1]          # (m, r_high)
+        R1 = R[I1]
+        Q2 = Q[:, I2]
+        R2 = R[I2]
+        B = Q1[:k // lora_r1] @ R1[:, :lora_r2]      
+        A = (Q2[:d // lora_r2] @ R2[:, :lora_r1]).T
+        module.lora_B.default.weight = torch.nn.Parameter(B.contiguous().to(module.lora_B.default.weight.device))
+        module.lora_A.default.weight = torch.nn.Parameter(A.contiguous().to(module.lora_A.default.weight.device))
+
+    elif init_config.mode == "gradient":
+        named_grad = kwargs["named_grads"]
+        grad_name = ".".join(name.split(".")[2:]) + ".weight"
+        grads = named_grad[grad_name]
+        # print(grads.shape)
+        if lora_r1 == 1 and lora_r2 == 1:
+            U, S, V = torch.svd_lowrank(-grads.cuda().float(), q=512, niter=16)
+        else:
+            U, S, V = torch.svd_lowrank(rebuild(-grads.float(),lora_r1, lora_r2), q=4*lora_r, niter=16)
+        V = V.T
+        # set direction
+        if init_config.direction == "ArBr":
+            if lora_r1 == 1 and lora_r2 == 1:
+                B = U[:, :lora_r] @ torch.diag(torch.sqrt(S[:lora_r])) / torch.sqrt(S[0]) / 128.0 **0.5
+                A = torch.diag(torch.sqrt(S[:lora_r])) @ V[:lora_r, :] / torch.sqrt(S[0]) / 128.0 **0.5
+                module.lora_B.default.weight = torch.nn.Parameter(B.contiguous().to(module.lora_B.default.weight.device))
+                module.lora_A.default.weight = torch.nn.Parameter(A.contiguous().to(module.lora_A.default.weight.device))
+            else:
+                for i in range(lora_r):
+                    B = (S[i] / S[0] / 1024)**0.5 * V[i, :].reshape([lora_r2, grads.shape[0]//lora_r1]).T
+                    A = (S[i] / S[0] / 1024)**0.5 * U[:, i].reshape([grads.shape[1]//lora_r2,lora_r1]).T
+                    module.lora_A.default.weight[i::lora_r] = torch.nn.Parameter(A.contiguous().to(module.lora_A.default.weight.device))
+                    module.lora_B.default.weight[:,i::lora_r] = torch.nn.Parameter(B.contiguous().to(module.lora_B.default.weight.device))
+        elif init_config.direction == "A2rBr":
+            B = U[:, :lora_r]
+            A = V[lora_r : 2 * lora_r, :]
+        elif init_config.direction == "ArB2r":
+            B = U[:, lora_r : 2 * lora_r]
+            A = V[:lora_r, :]
+        scaling_factor = module.scaling["default"]
+        if init_config.scale == "gd":
+            A = A / scaling_factor
+            B = B / scaling_factor
+        elif init_config.scale == "unit":
+            # Because A,B is orthogonal, do not need to scale
+            pass
+        elif init_config.scale == "stable":
+            m, n = grads.shape # m: feature_out, n: feature_in
+            # the scale of output is only related to the feature_out
+            gamma = init_config.stable_gamma
+
+
+        elif init_config.scale == "weightS":
+            _, S, _ = torch.svd_lowrank(module.weight.float(), q=4 * lora_r, niter=4)
+            S = S / module.scaling["default"]
+            avg_s = torch.sqrt(S[:lora_r]).mean().to(A.device)
+            B = B * avg_s
+            A = A * avg_s
+        # module.lora_B.default.weight = torch.nn.Parameter(B.contiguous().to(module.lora_B.default.weight.device))
+        # module.lora_A.default.weight = torch.nn.Parameter(A.contiguous().to(module.lora_A.default.weight.device))
+
+    with torch.no_grad():
+        # consider dtype not in init_config
+        if "dtype" not in init_config:
+            pass
+        elif init_config.dtype == "bf16":
+            module.lora_A.default.weight.data = module.lora_A.default.weight.data.to(
+                torch.bfloat16
+            )
+            module.lora_B.default.weight.data = module.lora_B.default.weight.data.to(
+                torch.bfloat16
+            )
+        elif init_config.dtype == "fp32":
+            module.lora_A.default.weight.data = module.lora_A.default.weight.data.to(
+                torch.float32
+            )
+            module.lora_B.default.weight.data = module.lora_B.default.weight.data.to(
+                torch.float32
+            )
+        # If lora_A@lora_B is not zero, then we need to subtract lora_A@lora_B from the original weight matrix
+        if init_config.mode == "qr":
+            offset = (kron(module.lora_B.default.weight.contiguous(),module.lora_A.default.weight.contiguous())).to(
+            module.weight.data.device
+        )
+        else:
+            offset = 0
+        # offset = (module.lora_B.default.weight @ module.lora_A.default.weight).to(
+        #     module.weight.data.device
+        # )
+
+        scaling_factor = module.scaling["default"]
+        offset *= scaling_factor
+        if "norm_clip" in init_config and init_config.norm_clip:
+            # for numerical stability, offset's largest value must be less then weight's largest value
+            ratio = torch.max(torch.abs(module.weight.data)) / torch.max(
+                torch.abs(offset)
+            )
+            if ratio < 1:
+                offset *= ratio
+                module.lora_A.default.weight.data *= ratio**0.5
+                module.lora_B.default.weight.data *= ratio**0.5
+                log.warning(f"Clipping offset by {ratio}")
+        try:
+            module.weight.data -= offset
+        except:
+            breakpoint()
+
+
+def reinit_lora(model, init_config, **kwargs):
+    r"""
+    Reinitialize the lora model with the given configuration.
+    """
+    for name, module in tqdm(
+        model.named_modules(),
+        desc="Reinitializing Lora",
+        total=len(list(model.named_modules())),
+    ):
+        if isinstance(module, LoraLinear):
+            reinit_lora_modules(name, module, init_config, **kwargs)
+
+    return model
+
+
+def get_record_gradient_hook(model, record_dict):
+    def record_gradient_hook(grad):
+        for n, p in model.named_parameters():
+            if p.requires_grad and p.grad is not None:
+                if n not in record_dict:
+                    record_dict[n] = p.grad.cpu()
+                else:
+                    record_dict[n] += p.grad.cpu()
+                p.grad = None
+        return grad
+
+    return record_gradient_hook
+
+
+def estimate_gradient(
+    model, dataset, batch_size: int = 4
+) -> Dict[str, List[torch.Tensor]]:
+    r"""
+    Estimate the gradient of the model on the given dataset
+    """
+    log.info("Estimating gradient")
+    model.train()
+    named_grads = {}
+    hooks = []
+    for name, param in model.named_parameters():
+        hook = param.register_hook(get_record_gradient_hook(model, named_grads))
+        hooks.append(hook)
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
+    num = 0
+    for batch in tqdm(dataloader, desc="Estimating gradient"):
+        num += 1
+        batch = {k: v.to(model.device) for k, v in batch.items()}
+        outputs = model(**batch)
+        outputs.loss.backward()
+        get_record_gradient_hook(model, named_grads)(None)  # get gradient of last layer
+        # make sure the gradient is cleared
+        for n, p in model.named_parameters():
+            if p.grad is not None:
+                p.grad = None
+    for n, g in named_grads.items():
+        named_grads[n] /= num
+    for hook in hooks:
+        hook.remove()
+    torch.cuda.empty_cache()
+    return named_grads
+
+
+
+
+
+
+def extract_num(text):
+    # Regex pattern to find the number following '####'
+    pattern = r'####\s*(\d+)'
+    # Using re.search to find the first match
+    match = re.search(pattern, text)
+    if match:
+        result = match.group(1)
+        print(text)
+    else:
+        print(text)
+        result = ""
+    try:
+        return int(result.replace(",", ""))
+    except:
+        print(f"'{result}' can't be converted")
+        return 0
+
+
+def eval_gsm8k(model,tokenizer,model_type, test_set):
+    all = 0
+    correct = 0
+    t = tqdm(test_set)
+    for example in t:
+        # print(example['x'])
+        pred_text = model_inference(model, tokenizer, example['x'], model_type, max_target_length=512)
+        gt = extract_num(example["y"])
+        pred = extract_num(pred_text)
+        correct += int(gt == pred)
+        all += 1
+        t.set_description(f"Accuracy: {correct / all * 100:02f}%")
+
+    print("Acc:", correct / all)
+    # append to gsm8k_results.txt (create if not exists)
+    if not os.path.exists("gsm8k_results.txt"):
+        with open("gsm8k_results.txt", "w") as f:
+            f.write("Model Acc\n")
+    with open("gsm8k_results.txt", "a") as f:
+        f.write(f"{model_name} {correct / all}\n")
+
+@hydra.main(version_base="1.2", config_path="conf", config_name="config")
+def run_exp(cfg: DictConfig):
+    log.info(OmegaConf.to_yaml(cfg))
+    seed_everything(cfg.seed)
+    model_name = cfg.model.name
+    model_type = cfg.model.type
+    dataset_name = cfg.dataset_name
+    dataset_func = DATASET_MAP[dataset_name]
+    use_peft = cfg.peft.use_peft
+    if_use_rslora = cfg.peft.use_rslora
+    lora_r = cfg.peft.lora_r
+    lora_r1 = cfg.peft.lora_r1
+    lora_r2 = cfg.peft.lora_r2
+    lora_relative_r = cfg.peft.lora_relative_r
+    lora_target_modules = cfg.peft.lora_target_modules
+    train_embeddings = cfg.peft.train_embeddings
+    if cfg.dry_run:
+        return
+    if use_peft:
+        lora_r = cfg.peft.lora_r
+        lora_r1 = cfg.peft.lora_r1
+        lora_r2 = cfg.peft.lora_r2
+        lora_alpha = cfg.peft.lora_alpha
+        lora_relative_r = None
+        init = cfg.init.mode
+    else:
+        lora_r = None
+        lora_target_modules = None
+        lora_relative_r = None
+        train_embeddings = True
+    config = {
+        "model_name": model_name,
+        "dataset_name": dataset_name,
+        "use_peft": use_peft,
+        "lora_r1": lora_r1,
+        "lora_r2": lora_r2,
+        "lora_r": lora_r,
+        "lora_alpha": lora_alpha,
+        "init": init,
+        "lora_target_modules": str(lora_target_modules),
+        "lora_relative_r": lora_relative_r,
+        "train_embeddings": train_embeddings,
+    }
+    if cfg.wandb.name:
+        name = cfg.wandb.name
+    else:
+        name = "_".join([f"{k}={v}" for k, v in config.items()])
+    cfg.wandb.project += "_" + cfg.dataset_name
+    wandb.init(
+        project=cfg.wandb.project,
+        name=name,
+        config=config,
+    )
+    train_set, val_set, eval_set = dataset_func()
+    model, tokenizer = initialize_text_to_text_model(
+        model_name, model_type, cfg.model.bf16, cfg.peft.use_peft, flash_attention=True
+    )
+    additional_kwargs = {}
+    if use_peft and cfg.init.mode == "gradient":
+        if isinstance(train_set, list):
+            temp_set = train_set[: cfg.init.bsz * cfg.init.iters]
+        else:
+            temp_set = train_set.select(range(cfg.init.bsz * cfg.init.iters))
+        transform_dataset(
+            model_type=model_type,
+            dataset=temp_set,
+            tokenizer=tokenizer,
+            max_length=cfg.init.max_length,
+        )
+        # named_grads = estimate_layer_inputs(model, temp_set, cfg.init.bsz)
+        named_grads = estimate_gradient(model, temp_set, cfg.init.bsz)
+        additional_kwargs["named_grads"] = named_grads
+        
+    additional_kwargs["lora_r1"] = lora_r1
+    additional_kwargs["lora_r"] = lora_r
+    additional_kwargs["lora_r2"] = lora_r2
+
+    if lora_target_modules == "all":
+        lora_target_modules = find_all_linear_modules(model)
+    else:
+        lora_target_modules = list(lora_target_modules) if lora_target_modules else []
+    if lora_relative_r is not None:
+        hidden_size = find_hidden_state_size(model)
+        lora_r = int(hidden_size * lora_relative_r)
+        log.info(f"lora_r is set to {hidden_size} * {lora_relative_r} = {lora_r}")
+    if use_peft and cfg.peft.get("dora", False):
+        log.info("Using Dora")
+        peft_config = LoraConfig(
+            r1=lora_r1,
+            r2=lora_r2,
+            lora_alpha=cfg.peft.lora_alpha,
+            target_modules=lora_target_modules,
+            use_rslora=if_use_rslora,
+            use_dora=True,
+        )
+        orig_model_params = sum(p.numel() for p in model.parameters())
+        model = get_peft_model(model, peft_config)
+        trainable_params, all_param = model.get_nb_trainable_parameters()
+        rate = {
+            "trainable_params": trainable_params,
+            "orig_params": orig_model_params,
+            "all_params": all_param,
+            "trainable_ratio": trainable_params / all_param,
+            "param_ratio": trainable_params / orig_model_params,
+        }
+    elif use_peft and cfg.peft.get("adalora", False):
+        log.info("Using AdaLora")
+        peft_config = AdaLoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            target_r=lora_r,
+            lora_alpha=cfg.peft.lora_alpha,
+            target_modules=lora_target_modules,
+            total_step=int(len(train_set)/cfg.model.real_batch_size)*cfg.model.epochs,
+        )
+        orig_model_params = sum(p.numel() for p in model.parameters())
+        model = get_peft_model(model, peft_config)
+        trainable_params, all_param = model.get_nb_trainable_parameters()
+        rate = {
+            "trainable_params": trainable_params,
+            "orig_params": orig_model_params,
+            "all_params": all_param,
+            "trainable_ratio": trainable_params / all_param,
+            "param_ratio": trainable_params / orig_model_params,
+        }
+    elif use_peft:
+        peft_config = LoraConfig(
+            r1=lora_r1,
+            r2=lora_r2,
+            r= lora_r,
+            lora_alpha=cfg.peft.lora_alpha,
+            target_modules=lora_target_modules,
+            use_rslora=if_use_rslora,
+        )
+        orig_model_params = sum(p.numel() for p in model.parameters())
+        model = get_peft_model(model, peft_config)
+        reinit_lora(model, cfg.init, **additional_kwargs)
+        if train_embeddings:
+            model.lm_head.weight.requires_grad = True
+        trainable_params, all_param = model.get_nb_trainable_parameters()
+        rate = {
+            "trainable_params": trainable_params,
+            "orig_params": orig_model_params,
+            "all_params": all_param,
+            "trainable_ratio": trainable_params / all_param,
+            "param_ratio": trainable_params / orig_model_params,
+        }
+        save_dir = os.path.join(
+            "results", f"{cfg.wandb.project}/{name}/{cfg.seed}", "orig_checkpoint"
+        )
+        model.save_pretrained(save_dir)
+        adapter_config = json.load(open(os.path.join(save_dir, "adapter_config.json")))
+        adapter_config["lora_alpha"] = -adapter_config["lora_alpha"]
+        json.dump(
+            adapter_config, open(os.path.join(save_dir, "adapter_config.json"), "w")
+        )
+    else:
+        # full finetune
+        all_param = sum(p.numel() for p in model.parameters())
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        rate = {
+            "trainable_params": trainable_params,
+            "orig_params": all_param,
+            "all_params": all_param,
+            "trainable_ratio": trainable_params / all_param,
+            "param_ratio": 1,
+        }
+    log.info(rate)
+    wandb.summary.update(rate)
+    training_loop = train_text_to_text_model
+    global s
+    print(s)
+    
+    model = training_loop(
+        f"{cfg.wandb.project}/{name}",
+        train_set,
+        val_set,
+        model,
+        tokenizer,
+        model_type,
+        num_train_epochs=cfg.model.epochs,
+        per_device_batch_size=cfg.model.per_device_batch_size,
+        real_batch_size=cfg.model.real_batch_size,
+        bf16=cfg.model.bf16,
+        eval_epochs=cfg.model.eval_epochs,
+        early_stopping_patience=cfg.model.early_stopping_patience,
+        max_length=cfg.model.max_length,
+        logging_steps=cfg.model.logging_steps,
+        use_loraplus=cfg.peft.use_loraplus,
+        loraplus_lr_ratio=cfg.peft.loraplus_lr_ratio,
+        learning_rate=cfg.model.learning_rate,
+        # deepspeed=(
+        #     "z3_offload_all_bf16.json" if cfg.peft == False else None
+        # ),
+        gradient_checkpointing=cfg.get("gradient_checkpointing", False),
+        seed=cfg.seed,
+    )
+
+
+
+    save_dir = os.path.join(
+        "results", f"{cfg.wandb.project}/{name}/{cfg.seed}"
+    )
+    if not use_peft:
+        model.save_pretrained(save_dir)
+        tokenizer.save_pretrained(save_dir)
+    else:
+        # merge_llama(os.path.join("results", f"{cfg.wandb.project}/{name}/{cfg.seed}"))
+        pass
+    log.info(f"Saving model to {save_dir}")
+    if dataset_name == 'meta_math':
+        train_set, val_set, eval_set = load_gsm8k()
+        model.generation_config.pad_token_id = tokenizer.pad_token_id
+        eval_gsm8k(model,tokenizer,model_type,eval_set)
+    if dataset_name == 'codefeedback':
+        model.generation_config.pad_token_id = tokenizer.pad_token_id
+        humaneval(model,tokenizer,save_dir, model_type)
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    run_exp()

split.py

@@ -0,0 +1,136 @@
+import torch
+from typing import List, Tuple
+
+
+# def vec(K):
+#     return K.T.flatten().reshape(-1, 1)
+
+# def rebuild(K, r1, r2):
+#     """
+#     Implements the R(K) operation from the image.
+#     K: input matrix (k x d)
+#     r1: block height
+#     r2: number of block columns
+#     """
+#     k, d = K.shape
+#     num_block_rows = k // r1
+#     num_block_cols = r2
+#     bw = d // r2 # block width
+    
+#     blocks = []
+#     # R(K) stacks vec(Ki,j) as columns. 
+#     # The image shows column-major order through the blocks.
+#     for j in range(num_block_cols):
+#         for i in range(num_block_rows):
+#             # Extract block Ki,j
+#             Ki_j = K[i*r1:(i+1)*r1, j*bw:(j+1)*bw]
+#             # Vectorize (column-major) and add to list
+#             blocks.append(vec(Ki_j))
+    
+#     return torch.hstack(blocks)
+
+
+def rebuild(K, r1, r2):
+    k, d = K.shape
+    Br = k // r1      # number of block rows
+    bw = d // r2      # block width
+
+    # Step 1: reshape to (Br, r1, r2, bw)
+    K_view = K.view(Br, r1, r2, bw)
+
+    # Step 2: we want to vectorize each (r1, bw) block in COLUMN-MAJOR order.
+    # That is equivalent to transposing the block and flattening in row-major.
+    # So we permute to (Br, r2, bw, r1) and then flatten last two dims.
+    
+    # But better: move r1 and bw to end, then transpose those two
+    # Actually: to get column-major flatten of (r1, bw), we can do:
+    #   block.transpose(-2, -1).contiguous().view(-1)
+    # So let's transpose the last two dims of the block
+
+    # Current: (Br, r1, r2, bw) → we want to treat (r1, bw) as block → transpose to (bw, r1)
+    # So permute to (Br, r2, bw, r1)
+    K_transposed_blocks = K_view.permute(0, 2, 3, 1)  # (Br, r2, bw, r1)
+
+    # Now flatten the last two dims (bw, r1) → (bw * r1,) → this is column-major of original block
+    vecs = K_transposed_blocks.reshape(Br, r2, bw * r1)  # (Br, r2, vec_len)
+
+    # Now, we have vecs[i, j] = vectorized block (i,j)
+    # But we want to output columns in order: j=0: i=0,1,...,Br-1; j=1: i=0,...
+    # So we need to **transpose the first two dimensions** and then **flatten in row-major**
+    
+    # Transpose to (r2, Br, vec_len)
+    vecs = vecs.permute(1, 0, 2)  # (r2, Br, vec_len)
+
+    # Now flatten first two dims: (r2*Br, vec_len), then transpose to (vec_len, r2*Br)
+    result = vecs.reshape(r2 * Br, -1).t()
+
+    return result
+
+# def rebuild(grad, block_size: [int, int]):
+#     new_matrix_rows = []
+#     if grad.dim() == 2:  # 只处理二维梯度（矩阵）
+#         # 获取梯度矩阵的尺寸
+#         rows, cols = grad.size()
+#         # 遍历分块
+#         for j in range(0, cols, block_size[1]):
+#             for i in range(0, rows, block_size[0]):
+#                 # 获取当前块
+#                 block = grad[i:i + block_size[0], j:j + block_size[1]]
+#                 # 如果块的大小不足，填充零
+#                 if block.size(0) < block_size[0] or block.size(1) < block_size[1]:
+#                     padding = (
+#                         0, block_size[1] - block.size(1),  # 列填充
+#                         0, block_size[0] - block.size(0)  # 行填充
+#                     )
+#                     block = torch.nn.functional.pad(block, padding, "constant", 0)
+#                 # 向量化并添加到新矩阵的行中
+#                 new_matrix_rows.append(block.T.flatten())
+
+#         # 将所有行堆叠成一个新矩阵
+#     if new_matrix_rows:  # 如果有数据
+#         new_gad = torch.stack(new_matrix_rows)
+#     else:
+#         new_gad = torch.empty(0)  # 如果没有梯度数据，返回空矩阵
+
+#     return new_gad
+
+# if __name__ == "__main__":
+#     # 定义一个简单的模型
+#     class SimpleModel(torch.nn.Module):
+#         def __init__(self):
+#             super(SimpleModel, self).__init__()
+#             self.fc1 = torch.nn.Linear(10, 20)
+#             self.fc2 = torch.nn.Linear(20, 10)
+
+#         def forward(self, x):
+#             x = self.fc1(x)
+#             x = self.fc2(x)
+#             return x
+
+#     # 初始化模型和损失函数
+#     model = SimpleModel()
+#     criterion = torch.nn.MSELoss()
+#     optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+#     # 模拟输入数据
+#     inputs = torch.randn(5, 10)  # batch_size=5, input_size=10
+#     targets = torch.randn(5, 10)  # batch_size=5, output_size=10
+
+#     # 前向传播
+#     outputs = model(inputs)
+#     loss = criterion(outputs, targets)
+
+#     # 反向传播计算梯度
+#     loss.backward()
+
+#     # 调用函数将梯度分块并构造新矩阵
+#     for param in model.parameters():
+#         if param.grad is not None:  # 检查是否有梯度
+#             grad = param.grad  # 获取梯度
+#             print("旧矩阵的内容:\n", grad)
+#             print("新矩阵的形状:", grad.shape)
+#             block_size = (2, 2)  # 分块大小为 2x2
+#             new_matrix = rebuild(grad, block_size)
+#             print("新矩阵��形状:", new_matrix.shape)
+#             print("新矩阵的内容:\n", new_matrix)
+

utils.py

@@ -0,0 +1,427 @@
+import torch
+import os
+import typing as tp
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer,
+    EarlyStoppingCallback,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+)
+from transformers.trainer_utils import PredictionOutput
+from datasets import Dataset, load_dataset
+from torch.utils.data import DataLoader
+from transformers import AdamW, get_linear_schedule_with_warmup
+from lora_plus import LoraPlusTrainingArguments, LoraPlusTrainer
+from logTrainer import LogTrainer
+import logging
+import wandb
+from peft import PeftModel
+from data import load_alpaca
+
+log = logging.getLogger(__name__)
+
+
+def causalLMEncode(example, tokenizer, max_length=-1, ignore_masked_token=True):
+    is_list_input = isinstance(example["x"], list)
+    # Combine text and add EOS token
+    combined_text = (
+        [
+            x + " " + y + tokenizer.eos_token
+            for (x, y) in zip(example["x"], example["y"])
+        ]
+        if is_list_input
+        else example["x"] + " " + example["y"] + tokenizer.eos_token
+    )
+    # Tokenize combined text
+    encodings = tokenizer(
+        combined_text,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_length if max_length != -1 else None,
+    )
+    # Calculate input text length in tokens
+    input_text_length = (
+        [
+            len(tokenizer(example["x"][i], return_tensors="pt")["input_ids"][0])
+            for i in range(len(example["x"]))
+        ]
+        if is_list_input
+        else len(tokenizer(example["x"], return_tensors="pt")["input_ids"][0])
+    )
+    if input_text_length[0] >= max_length:
+        log.warning(
+            f"Input text length >= max_length: {input_text_length} >= {max_length}. "
+            "Consider increasing max_length to avoid truncation."
+        )
+    # Create labels
+    labels = encodings["input_ids"].clone()
+    if is_list_input:
+        for i, l in enumerate(input_text_length):
+            labels[i, :l] = -100
+    else:
+        labels[0, :input_text_length] = -100
+    if ignore_masked_token:
+        labels[encodings["attention_mask"] == 0] = -100
+    # Update example dictionary
+    results = {
+        "input_ids": encodings["input_ids"],
+        "attention_mask": encodings["attention_mask"],
+        "labels": labels,
+        # "input_text_length": input_text_length,
+    }
+
+    return results
+
+
+def SeqToSeqEncode(example, tokenizer, max_length=None, ignore_masked_token=False):
+    inputs = tokenizer(
+        example["x"],
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+    )
+    outputs = tokenizer(
+        example["y"],
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+    )
+
+    results = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+        "labels": outputs["input_ids"],
+        "decoder_attention_mask": outputs["attention_mask"],
+    }
+
+    if ignore_masked_token:
+        results["labels"][outputs["attention_mask"] == 0] = -100
+
+    return results
+
+
+def preprocess_dataset(
+    dataset: tp.Union[Dataset, tp.List[tp.Tuple[str, str]], tp.List[tp.Dict[str, str]]]
+) -> Dataset:
+    if isinstance(dataset, list) and isinstance(dataset[0], tuple):
+        dataset = Dataset.from_pandas(pd.DataFrame(dataset, columns=["x", "y"]))
+    elif isinstance(dataset, list) and isinstance(dataset[0], dict):
+        dataset = Dataset.from_dict(
+            {k: [dic[k] for dic in dataset] for k in dataset[0]}
+        )
+    elif isinstance(dataset, dict):
+        dataset = Dataset.from_dict(dataset)
+    elif isinstance(dataset, Dataset):
+        pass
+    else:
+        raise ValueError("Wrong format")
+    return dataset
+
+
+def initialize_text_to_text_model(
+    model_name: str,
+    model_type: str,
+    bf16: bool,
+    use_peft: bool = True,
+    tokenizer: str = None,
+    flash_attention: bool = False,
+):
+    if model_type == "CausalLM":
+        if flash_attention:
+            log.info("Using flash attention 2")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16 if bf16 else torch.float32,
+                device_map="auto" if use_peft else None,
+                attn_implementation="flash_attention_2",
+            )
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16 if bf16 else torch.float32,
+                device_map="auto" if use_peft else None,
+            )
+    elif model_type == "ConditionalGeneration":
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16 if bf16 else torch.float32,
+            device_map="auto" if use_peft else None,
+        )
+    if tokenizer:
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if tokenizer.eos_token is None:
+        tokenizer.add_special_tokens({"eos_token": "<|endoftext|>"})
+        model.resize_token_embeddings(len(tokenizer))
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    return model, tokenizer
+
+
+def compute_metrics(p: PredictionOutput):
+    predictions = p.predictions
+    label_ids = p.label_ids # shape (batch_size, seq_len)
+    if False:
+        # Hard metric: the model must output exactly the same as the target
+        # This should be the default evaluation metric for most tasks
+        pred = np.argmax(predictions[0], axis=-1)
+        num_correct = sum([np.array_equal(pred[i], label_ids[i]) for i in range(len(pred))])
+        accuracy = num_correct / len(pred)
+    else:
+        # Soft metric: we limit the output space to the target space
+        # i.e. the model classify the one with higher prob in positive and negative
+        # **Use it in cola and mrpc, because it's too hard for vanilla lora**
+        # Only suit for the binary classification with each label of 1 token
+        label_ids = label_ids[:, 0] # remove the eos token
+        unique_labels = np.unique(label_ids)
+        flipped_labels = np.ones_like(label_ids) * unique_labels.sum() - label_ids
+        predictions = predictions[0][:, 0, :] # remove the eos token # seq_len, tokens
+        label_prob = predictions[np.arange(len(predictions)), label_ids]
+        flipped_label_prob = predictions[np.arange(len(predictions)), flipped_labels]
+        num_correct = sum(label_prob > flipped_label_prob)
+        accuracy = num_correct / len(label_prob)
+
+    return {"accuracy": accuracy}
+
+
+def transform_dataset(model_type, tokenizer, dataset, max_length):
+    if model_type == "CausalLM":
+        dataset.set_transform(lambda x: causalLMEncode(x, tokenizer, max_length))
+    elif model_type == "ConditionalGeneration":
+        dataset.set_transform(lambda x: SeqToSeqEncode(x, tokenizer, max_length))
+    else:
+        raise ValueError("Wrong model type")
+    return dataset
+
+
+def train_text_to_text_model(
+    run_name: str,
+    train_dataset: Dataset,
+    valid_dataset: Dataset,
+    model: torch.nn.Module,
+    tokenizer: AutoTokenizer,
+    model_type: str,
+    per_device_batch_size: int = 1,
+    real_batch_size: int = 32,
+    max_length: int = None,
+    **kwargs,
+) -> torch.nn.Module:
+    # Preprocess the dataset
+    train_dataset = preprocess_dataset(train_dataset)
+    valid_dataset = preprocess_dataset(valid_dataset)
+
+    assert (
+        real_batch_size % per_device_batch_size == 0
+    ), "real_batch_size must be divisible by per_device_batch_size"
+    accu_step = real_batch_size // per_device_batch_size
+
+    train_dataset, valid_dataset = transform_dataset(
+        model_type, tokenizer, train_dataset, max_length
+    ), transform_dataset(model_type, tokenizer, valid_dataset, max_length)
+
+    eval_steps = (
+        int(len(train_dataset) * kwargs.get("eval_epochs", 1)) // real_batch_size
+    )
+    # Special for lorqplus
+    use_loraplus = kwargs.get("use_loraplus", False)
+    TrainingArgumentsClass = (
+        LoraPlusTrainingArguments if use_loraplus else Seq2SeqTrainingArguments
+    )
+    TrainerClass = LoraPlusTrainer if use_loraplus else LogTrainer
+    if use_loraplus:
+        additional_kwargs = {
+            "loraplus_lr_ratio": kwargs.get("loraplus_lr_ratio", 1.0),
+        }
+        log.info(
+            f"Begin training using LoraPlusTrainer with additional kwargs: {additional_kwargs}"
+        )
+    else:
+        additional_kwargs = {}
+        log.info("Begin training using Seq2SeqTrainer")
+
+    # Training arguments
+    output_dir = f"./results/{run_name}/{kwargs.get('seed')}"
+    training_args = TrainingArgumentsClass(
+        output_dir=output_dir,  # output directory
+        num_train_epochs=kwargs.get(
+            "num_train_epochs", 3
+        ),  # total number of training epochs
+        per_device_train_batch_size=per_device_batch_size,
+        per_device_eval_batch_size=per_device_batch_size,
+        gradient_accumulation_steps=accu_step,
+        logging_dir="./logs",  # directory for storing logs
+        logging_steps=kwargs.get("logging_steps", 10),  # when to print log
+        bf16=kwargs.get("bf16", False),
+        gradient_checkpointing=kwargs.get("gradient_checkpointing", False),
+        optim=kwargs.get("optim", "adamw_torch"),
+        evaluation_strategy="no",
+        eval_steps=eval_steps,
+        save_steps=eval_steps,
+        save_strategy="steps",
+        save_total_limit=1,  # No need for saving
+        load_best_model_at_end=False,
+        metric_for_best_model=kwargs.get("metric_for_best_model", "eval_loss"),
+        greater_is_better=kwargs.get("greater_is_better", False),
+        do_eval=False,
+        learning_rate=kwargs.get("learning_rate", 5e-5),
+        remove_unused_columns=False,  # We tokenize the dataset on the fly
+        eval_accumulation_steps=kwargs.get("eval_accumulation_steps", real_batch_size),
+        label_names=[
+            "labels"
+        ],  # Peft are not compatible with HF's default label names yet
+        # Ref: https://discuss.huggingface.co/t/eval-with-trainer-not-running-with-peft-lora-model/53286
+        # weight_decay = 0, # No weight decay
+        weight_decay = 5e-4,
+        warmup_ratio = 0.03,
+        lr_scheduler_type = "cosine",
+        seed = kwargs.get("seed", 42),
+        **additional_kwargs,
+    )
+
+    trainer = TrainerClass(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        compute_metrics=compute_metrics if "llama" not in run_name else None,
+        # callbacks=[
+        #     EarlyStoppingCallback(
+        #         early_stopping_patience=kwargs.get("early_stopping_patience", 1)
+        #     ),
+        # ],
+    )
+
+    trainer.train()
+    # eval_results = trainer.evaluate()
+    # eval_accuracy = eval_results.get("eval_accuracy", 0)  
+    # print(f"FINAL_EVAL_ACCURACY: {eval_accuracy:.4f}")
+    return model
+
+
+def model_inference(
+    model: torch.nn.Module,
+    tokenizer: AutoTokenizer,
+    input_text: str,
+    model_type: str,
+    max_source_length: str = 768,
+    max_target_length: str = 256,
+):
+    if model_type == "CausalLM":
+        inputs = tokenizer(
+            input_text + " ",
+            return_tensors="pt",
+            max_length=max_source_length,
+            truncation=True,
+            return_token_type_ids=False,
+        )
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                return_dict_in_generate=True,
+                output_scores=False,
+                max_new_tokens=max_target_length,
+                eos_token_id=tokenizer.eos_token_id,
+                top_p=0.95,
+                temperature=0.8,
+            )
+        pred_text = tokenizer.decode(
+            outputs.sequences[0][len(inputs["input_ids"][0]) :],
+            skip_special_tokens=True,
+        )
+    elif model_type == "ConditionalGeneration":
+        inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=max_target_length)
+        pred_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return pred_text
+
+
+def load_peft_model(model, peft_path: str):
+    peft_paths = [f"{peft_path}/{i}" for i in os.listdir(peft_path) if "merge" not in i]
+    for peft_path in peft_paths:
+        print(f"loading and merging from {peft_path}")
+        model: PeftModel = PeftModel.from_pretrained(model, peft_path)
+        model = model.merge_and_unload()
+    return model
+
+
+def test_train():
+    # Example usage using emo dataset
+    dataset = load_dataset("emo")
+    label_map = {0: "others", 1: "happy", 2: "sad", 3: "angry"}
+    dataset = dataset.map(lambda e: {"x": e["text"], "y": label_map[e["label"]]})
+    train_set = dataset["train"]
+    test_set = dataset["test"]
+
+    model_name = "t5-small"
+    model_type = "ConditionalGeneration"
+    model, tokenizer = initialize_text_to_text_model(model_name, model_type)
+
+    model = train_text_to_text_model(
+        train_set,
+        test_set,
+        model,
+        tokenizer,
+        model_type,
+        num_train_epochs=1,
+        per_device_batch_size=64,
+        real_batch_size=64,
+    )
+    # Use the model for inference in the testset, print the first 10 examples
+    for i in range(10):
+        print("Input:", test_set[i]["x"])
+        print("Target:", test_set[i]["y"])
+        print(
+            "Prediction:",
+            model_inference(model, tokenizer, test_set[i]["x"], model_type),
+        )
+        print()
+
+
+def test_llama_alpaca():
+    model_name = "meta-llama/Llama-2-7b-hf"
+    model_type = "CausalLM"
+    peft_path = "results/llama-alpaca_alpaca/gradient-ArB2r-adam/0"
+    model, tokenizer = initialize_text_to_text_model(model_name, model_type, True)
+    model = load_peft_model(model, peft_path)
+    _, _, test_set = load_alpaca()
+    for i in range(10):
+        print("Input:", test_set[i]["x"])
+        # print("Target:", test_set[i]["y"])
+        print(
+            "Prediction:",
+            model_inference(model, tokenizer, test_set[i]["x"], model_type),
+        )
+        print()
+
+
+def merge_llama(peft_path):
+    model_name = "meta-llama/Llama-2-7b-hf"
+    model_type = "CausalLM"
+    model, tokenizer = initialize_text_to_text_model(model_name, model_type, True)
+    model = load_peft_model(model, peft_path)
+    print("Save model to ", os.path.join(peft_path, "merged_checkpoint"))
+    model.save_pretrained(os.path.join(peft_path, "merged_checkpoint"))
+    tokenizer.save_pretrained(os.path.join(peft_path, "merged_checkpoint"))
+    del model, tokenizer
+
+
+if __name__ == "__main__":
+    merge_llama("results/llama-alpaca_alpaca/default/0")
+    # merge_llama("results/llama-alpaca_alpaca/gradient-ArB2r-adam/0")