if is_qlora and not is_sharded_qlora:

if is_qlora and not is_sharded_qlora and not isinstance(model, PeftModel):

import torch
from datasets import load_dataset
from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

repo_id = 'microsoft/Phi-3-mini-4k-instruct'

nf4_config = BitsAndBytesConfig(
   load_in_4bit=True,
)

model_q4 = AutoModelForCausalLM.from_pretrained(repo_id,
                                                device_map='cuda:0',
                                                quantization_config=nf4_config)

model_q4 = prepare_model_for_kbit_training(model_q4,
                                           use_gradient_checkpointing=False,
                                           gradient_checkpointing_kwargs={})

config = LoraConfig(
    r=16,
    target_modules=['o_proj', 'qkv_proj', 'gate_up_proj', 'down_proj'],
    lora_alpha=32,
    task_type="CAUSAL_LM",
)
peft_model = get_peft_model(model_q4, config)

tokenizer = AutoTokenizer.from_pretrained(repo_id)
tokenizer.pad_token = tokenizer.unk_token
tokenizer.pad_token_id = tokenizer.unk_token_id

dataset = load_dataset("dvgodoy/yoda_sentences", split="train")
dataset = dataset.rename_column("sentence", "prompt")
dataset = dataset.rename_column("translation_extra", "completion")
dataset = dataset.remove_columns(["translation"])

def format_dataset(examples):
    if isinstance(examples["prompt"], list):
        output_texts = []
        for i in range(len(examples["prompt"])):
            converted_sample = [
                {"role": "user", "content": examples["prompt"][i]},
                {"role": "assistant", "content": examples["completion"][i]},
            ]
            output_texts.append(converted_sample)
        return {'messages': output_texts}
    else:
        converted_sample = [
            {"role": "user", "content": examples["prompt"]},
            {"role": "assistant", "content": examples["completion"]},
        ]
        return {'messages': converted_sample}

dataset = dataset.map(format_dataset).remove_columns(['prompt', 'completion'])

# Using own PEFT model, LoRA already applied
trainer1 = SFTTrainer(
    model=peft_model,
    processing_class=tokenizer,
    train_dataset=dataset,
    data_collator=None,
    args=SFTConfig(
        output_dir="./future_name_on_the_hub",
        packing=False,
        report_to='none',
        bf16=False,
        max_length=64,
    )
)

# Using base model and LoRA config
base_model = AutoModelForCausalLM.from_pretrained(repo_id,
                                                device_map='cuda:0',
                                                quantization_config=nf4_config)

trainer2 = SFTTrainer(
    model=base_model,
    peft_config=config,
    processing_class=tokenizer,
    train_dataset=dataset,
    data_collator=None,
    args=SFTConfig(
        output_dir="./future_name_on_the_hub",
        packing=False,
        report_to='none',
        bf16=False,
        max_length=64,
    )
)

trainer1.model.print_trainable_parameters()
trainer2.model.print_trainable_parameters()