정보엔지니어 (맥스웰의 도깨비, 양자컴퓨터)

Gemma 공개후 곧바로 파인튜닝 가능한 가이드가 공개되었다.

https://adithyask.medium.com/a-beginners-guide-to-fine-tuning-gemma-0444d46d821c

gemma-2b-it 모델의 quantized 버전 정도는 16GB GPU로도 파인튜닝 테스트를 해볼 수 있다.

상기 가이드대로 GPU환경에서 python을 설치한다.

$ conda create -n gemma python=3.11

$ conda activate gemma

$ pip install bitsandbytes==0.42.0
$ pip install peft==0.8.2
$ pip install trl==0.7.10
$ pip install accelerate==0.27.1
$ pip install datasets==2.17.0
$ pip install transformers==4.38.0

아래 소스를 통해 샘플 데이터 셋을 허깅페이스에서 다운받아 파인튜닝해서 결과를 얻어볼 수 있다.

이때 허깅 페이스에 미리 계정을 만든 후 WRITE 가능한 token으로 로그인을 미리 해두자. WRITE 권한으로 생성한 토큰으로는 모델을 huggingface에 업로드 하는 것도 가능하다(아래 소스의 맨 마지막 두줄)

$ huggingface-cli login
...
Token: 

$ huggingface-cli whoami
Justik (필자의 profile이다)

아래 코드를 실행한다. LORA를 통한 경량형 튜닝 정도는 할 수 있게 된다. gemma-7b-it의 경우는 16gb에서는 메모리가 부족하다. (본래 글에서는 A100을 추천하고 있다.) gemma-2b-it의 경우는 아래의 코드로 10~14gb정도의 메모리를 소모한다.

import json
import pandas as pd
import torch
from datasets import Dataset, load_dataset
from huggingface_hub import notebook_login
from peft import LoraConfig, PeftModel
from peft import LoraConfig, get_peft_model
lora_config = LoraConfig(
    r=64,
    lora_alpha=32,
    target_modules=['o_proj', 'q_proj', 'up_proj', 'v_proj', 'k_proj', 'down_proj', 'gate_proj'],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
) 
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from trl import SFTTrainer

notebook_login()

#model_id = "google/gemma-7b-it"
# model_id = "google/gemma-7b"
model_id = "google/gemma-2b-it"
#model_id = "google/gemma-2b"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

#model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=lora_config, device_map={"":0})
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0})
#model = AutoModelForCausalLM.from_pretrained(model_id, device_map={"":0})
tokenizer = AutoTokenizer.from_pretrained(model_id, add_eos_token=True)
{
    "instruction": "Create a function to calculate the sum of a sequence of integers.",
    "input":"[1, 2, 3, 4, 5]",
    "output": "# Python code def sum_sequence(sequence): sum = 0 for num in sequence: sum += num return sum"
}

dataset = load_dataset("TokenBender/code_instructions_122k_alpaca_style", split="train")

def generate_prompt(data_point):
    """Gen. input text based on a prompt, task instruction, (context info.), and answer

    :param data_point: dict: Data point
    :return: dict: tokenzed prompt
    """
    prefix_text = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n'
    # Samples with additional context into.
    if data_point['input']:
        text = f"""<start_of_turn>user {prefix_text} {data_point["instruction"]} here are the inputs {data_point["input"]} <end_of_turn>\\n<start_of_turn>model{data_point["output"]} <end_of_turn>"""
    # Without
    else:
        text = f"""<start_of_turn>user {prefix_text} {data_point["instruction"]} <end_of_turn>\\n<start_of_turn>model{data_point["output"]} <end_of_turn>"""
    return text

# add the "prompt" column in the dataset
text_column = [generate_prompt(data_point) for data_point in dataset]
dataset = dataset.add_column("prompt", text_column)
dataset = dataset.shuffle(seed=1234)  # Shuffle dataset here
dataset = dataset.map(lambda samples: tokenizer(samples["prompt"]), batched=True)

dataset = dataset.train_test_split(test_size=0.2)
train_data = dataset["train"]
test_data = dataset["test"]

model = get_peft_model(model, lora_config)
trainable, total = model.get_nb_trainable_parameters()
print(f"Trainable: {trainable} | total: {total} | Percentage: {trainable/total*100:.4f}%")

import transformers

from trl import SFTTrainer

tokenizer.pad_token = tokenizer.eos_token
torch.cuda.empty_cache()
trainer = SFTTrainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=test_data,
    dataset_text_field="prompt",
    peft_config=lora_config,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=0.03,
        max_steps=100,
        learning_rate=2e-4,
        logging_steps=1,
        output_dir="outputs",
        optim="paged_adamw_8bit",
        save_strategy="epoch",
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

# Start the training process
trainer.train()

new_model = "gemma-2b-it-finetune" #Name of the model you will be pushing to huggingface model hub
# Save the fine-tuned model
trainer.model.save_pretrained(new_model)

# Merge the model with LoRA weights
base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map={"": 0},
)
merged_model= PeftModel.from_pretrained(base_model, new_model)
merged_model= merged_model.merge_and_unload()

# Save the merged model
merged_model.save_pretrained("merged_model",safe_serialization=True)
tokenizer.save_pretrained("merged_model")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"


def get_completion(query: str, model, tokenizer) -> str:
  device = "cuda:0"
  prompt_template = """
  <start_of_turn>user
  Below is an instruction that describes a task. Write a response that appropriately completes the request.
  {query}
  <end_of_turn>\\n<start_of_turn>model
  
  """
  prompt = prompt_template.format(query=query)
  encodeds = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
  model_inputs = encodeds.to(device)
  generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True, pad_token_id=tokenizer.eos_token_id)
  # decoded = tokenizer.batch_decode(generated_ids)
  decoded = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
  return (decoded)

result = get_completion(query="code the fibonacci series in python using reccursion", model=merged_model, tokenizer=tokenizer)
print(result)


# Push the model and tokenizer to the Hugging Face Model Hub
# merged_model.push_to_hub(new_model, use_temp_dir=False)
# tokenizer.push_to_hub(new_model, use_temp_dir=False)

이렇게 hugging face hub에 올려진 모델은 아래와 같다.

https://huggingface.co/Justik/gemma-2b-it-finetune