<h2 align="center">
    <p>Instruction for ROLoRA</p>
</h2>

ROLoRA is edited on top of PEFT (HuggingFace). ROLoRA currently supports AdaLoRA and SoRA.

<h3> Fine-tune with ROLoRA via AdaLoRA </h3>

Import ROLoRA

```python
from peft import get_peft_model, RoLoraConfig, RoLoraModel, PeftModel, TaskType
```

First, write Config for ROLoRA, for example:

```python
peft_config = RoLoraConfig(
    task_type=TaskType.SEQ_CLS, inference_mode=False, lora_alpha=8, lora_dropout=0.1, 
    target_r=8,
    tinit=0, tfinal=1000, 
    init_r=8, total_step=num_training_steps,
    orth_reg_weight=0.1,
    deltaT=1,
    target_modules=target_modules
)
```

ROLoRA requires save and load models via accelerate package. First, write save and load hooks for accelerator.

```python
def save_model_hook(models, weights, output_dir):
    for model in models:
        model.save_pretrained(output_dir)
        weights.pop()
def load_model_hook(models, input_dir):
    while len(models) > 0:
        model = models.pop()
        adapter_names = list(model.peft_config.keys())
        for key in adapter_names:
            PeftModel.from_pretrained(model.base_model.model,output_dir, key, adapter_name=key)
```

In each iteration, during training, ROLoRA does two steps:

1. Sparsification, by calling

```python
model.update_and_allocate(iteration_step)
```

2. Extending adapters, by calling

```python
RoLoraModel.extend_modules(model, adapter_name, it, num_repeat)
```

Once extended, use accelerator to save and load new states; reinitialize optimizer.

<h3> Fine-tune with ROLoRA via SoRA </h3>

Config

```python
peft_config = SoraConfig(
    task_type=TaskType.SEQ_CLS, inference_mode=False, lora_alpha=8, lora_dropout=0.1, 
    r=8,
    target_modules=target_modules,
)
```

SoRA use sparse optimizers

```python
from peft.optimizers.sparse_optimizer_multiply_lr import SparseAdamW

adamW_group = [p for n, p in model.named_parameters() if "lora_E" not in n and p.requires_grad]
sparse_adamW_group = [p for n, p in model.named_parameters() if "lora_E" in n and p.requires_grad]

sparse_optimizer = SparseAdamW(params=sparse_adamW_group, sparse_lambda=args.threshold, lr=learning_rate, weight_decay=weight_decay)
optimizer = torch.optim.AdamW(adamW_group, lr=learning_rate, weight_decay=weight_decay)
```

In each iteration, during training, ROLoRA does two steps:

1. Sparsification, by calling the sparse optimzer

```python
optimizer.step()
sparse_optimizer.step()
```

2. Extending adapters, by calling

```python
SoraModel.extend_modules(model, adapter_name, it)
```

Once extended, use accelerator to save and load new states; reinitialize optimizer.

<h3> Inference with fine-tuned models with ROLoRA </h3>

First load config from saved model's adapter config. Example

```python
with open(f"{best_checkpoint}/default_0/adapter_config.json", "r") as f:
            config = json.load(f)
peft_config = RoLoraConfig(
    task_type=TaskType.SEQ_CLS, inference_mode=False, lora_alpha=lora_alpha, lora_dropout=lora_dropout, 
    target_r=config["target_r"],
    tinit=config["tinit"], tfinal=config["tfinal"], 
    init_r=config["init_r"], total_step=10,
    orth_reg_weight=config["orth_reg_weight"],
    deltaT=config["deltaT"],
    target_modules=config["target_modules"]
)
```
Then load model by accelerator (load model hook required). Example:

```python
model = AutoModelForSequenceClassification.from_pretrained("microsoft/deberta-v3-base",  num_labels=num_labels, return_dict=True)
model = get_peft_model(model, peft_config, adapter_name="default_0")
model = accelerator.prepare(model)
accelerator.load_state(f"{best_checkpoint}")
```
This loaded model can be used for inference as usual.

<h2> Changes from PEFT </h2>

1. ROLoRA (via AdaLoRA) is added to the list of tuners. 
2. Model via SoRA (https://github.com/TsinghuaC3I/SoRA) is added to list of tuners, with support for adapter rank extension.
3. Sparse optimizer (from https://github.com/TsinghuaC3I/SoRA) is added.