Dostosowanie modelu Llama-2 7B Chat do generowania kodu w Pythonie: Wykorzystanie QLoRA, SFTTrainer i Gradient Checkpointing na zbiorze danych Alpaca-14k

Skuteczne dostosowanie modelu Llama-2 7B Chat do generowania kodu w Pythonie

W dzisiejszych czasach sztuczna inteligencja coraz częściej wykorzystywana jest do automatyzacji pisania kodu programistycznego. Jednym z najnowszych osiągnięć w tej dziedzinie jest model Llama-2 7B Chat, który można dostosować do generowania kodu w języku Python przy użyciu zaawansowanych technik, takich jak QLoRA, gradient checkpointing oraz nadzorowane dostrajanie (Supervised Fine-Tuning, SFT) z wykorzystaniem SFTTrainer. Dzięki temu model może produkować wysokiej jakości kod Python, minimalizując jednocześnie wymagania sprzętowe.

Instalacja niezbędnych bibliotek

Aby rozpocząć proces dostosowania modelu, należy zainstalować odpowiednie biblioteki z repozytorium Python Package Index (PyPI). Obejmuje to takie pakiety jak PLACEHOLDER62f79cf895f74ed2, PLACEHOLDER3fd1828663df3913, PLACEHOLDER9df7cb706aae07b5 oraz PLACEHOLDER134cb3ab151a2832. Instalację można przeprowadzić za pomocą poniższej komendy:

python
!pip install -q accelerate
!pip install -q peft
!pip install -q transformers
!pip install -q trl

Flaga -q (quiet mode) zapewnia minimalizację wyświetlania zbędnych komunikatów podczas instalacji.

Importowanie kluczowych modułów

Do przeprowadzenia treningu modelu niezbędne są różne moduły, które umożliwiają m.in. załadowanie zbioru danych, konfigurację modelu i jego dostosowanie. Oto kod importujący wymagane biblioteki:

python
import os
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

Wybór modelu bazowego i zbioru danych

W tym przypadku wykorzystujemy model bazowy PLACEHOLDER1c0fa7545bba5f2a dostępny w Hugging Face Hub oraz zbiór danych PLACEHOLDER530393a1f8068d13. Nazwa nowego dostrojonego modelu to /kaggle/working/llama-2-7b-codeAlpaca.

python
model_name = "NousResearch/llama-2-7b-chat-hf"
dataset_name = "user/minipython-Alpaca-14k"
new_model = "/kaggle/working/llama-2-7b-codeAlpaca"

Konfiguracja parametrów LoRA

Low-Rank Adaptation (LoRA) to technika, która pozwala na efektywne dostrajanie dużych modeli językowych poprzez modyfikację niewielkiej liczby parametrów. W tym celu definiujemy kluczowe ustawienia:

python
lora_r = 64  # Wymiar atencji LoRA
lora_alpha = 16  # Parametr skali dla LoRA
lora_dropout = 0.1  # Prawdopodobieństwo dropout dla warstw LoRA

Dzięki tym parametrom model może efektywnie dostosowywać się do nowych zadań, jednocześnie minimalizując zużycie pamięci.

Ustawienia procesu treningowego

Podczas treningu modelu istotne jest określenie szeregu parametrów, takich jak liczba epok, rozmiar partii danych oraz strategia optymalizacyjna:

python
output_dir = "/kaggle/working/llama-2-7b-codeAlpaca"
num_train_epochs = 1
fp16 = True  # Trening w trybie FP16
per_device_train_batch_size = 8
per_device_eval_batch_size = 8
gradient_accumulation_steps = 2
gradient_checkpointing = True
max_grad_norm = 0.3
learning_rate = 2e-4
weight_decay = 0.001
optim = "adamw_torch"
lr_scheduler_type = "constant"
group_by_length = True
warmup_ratio = 0.03
save_steps = 100
logging_steps = 10

Dzięki tym ustawieniom model może efektywnie uczyć się na nowych danych i generować lepsze wyniki.

Sprawdzenie konfiguracji sprzętowej

Przed uruchomieniem treningu warto zweryfikować dostępność niezbędnych zasobów sprzętowych:

python
import torch
print("PyTorch Version:", torch.__version__)
print("CUDA Version:", torch.version.cuda)

!nvidia-smi  # Wyświetlenie informacji o GPU

Załadowanie modelu i zbioru danych

Model oraz zbiór danych są ładowane do pamięci przed rozpoczęciem procesu dostrajania:

python
dataset = load_dataset(dataset_name, split="train")

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
)

model.gradient_checkpointing_enable()
model.enable_input_require_grads()

Implementacja LoRA

LoRA jest stosowana w modelu z pomocą funkcji get_peft_model:

python
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, peft_config)

Uruchomienie treningu

Po konfiguracji modelu możemy uruchomić proces treningowy:

python
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=None,
    tokenizer=tokenizer,
    args=TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_train_epochs,
        per_device_train_batch_size=per_device_train_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        optim=optim,
        save_steps=save_steps,
        logging_steps=logging_steps,
        learning_rate=learning_rate,
        weight_decay=weight_decay,
        fp16=fp16,
        max_grad_norm=max_grad_norm,
        warmup_ratio=warmup_ratio,
        group_by_length=True,
        lr_scheduler_type=lr_scheduler_type,
    ),
    packing=False,
)

trainer.train()
trainer.model.save_pretrained(new_model)

Generowanie kodu Python przy użyciu dostrojonego modelu

Po zakończeniu treningu można sprawdzić działanie dostrojonego modelu, generując kod Python dla konkretnego zapytania:

python
prompt = "Jak napisać program w Pythonie, który oblicza średnią, odchylenie standardowe i współczynnik zmienności z pliku CSV?"
pipe = pipeline(task="text-generation", model=trainer.model, tokenizer=tokenizer, max_length=400)
result = pipe(f"[INST] {prompt} [/INST]")
print(result[0]['generated_text'])

Dzięki temu model może generować precyzyjne i użyteczne fragmenty kodu w odpowiedzi na zadane pytania.

Zwolnienie pamięci GPU

Po zakończeniu pracy z modelem warto zwolnić zasoby GPU, aby uniknąć przeciążenia pamięci:

python
del tokenizer
import gc
gc.collect()
torch.cuda.empty_cache()

Podsumowanie

Dzięki temu przewodnikowi udało się skutecznie dostroić model Llama-2 7B Chat do generowania kodu w Pythonie. Wykorzystanie technik takich jak QLoRA, gradient checkpointing oraz SFT pozwala zoptymalizować model pod względem pamięci i wydajności, co jest kluczowe przy pracy z dużymi modelami językowymi.

Dalsze badania i eksperymenty mogą pozwolić na jeszcze lepsze dostosowanie modelu do specjalistycznych zadań programistycznych, co otwiera nowe możliwości w automatyzacji kodowania.