Dostrojenie modelu Gemma za pomocą funkcji Hugging Face Transformers i QloRA

Wyświetl na ai.google.dev Uruchom w Google Colab Uruchom w Kaggle Otwórz w Vertex AI Wyświetl źródło na GitHubie

Ten przewodnik zawiera instrukcje dotyczące dostrajania modelu Gemma na niestandardowym zbiorze danych tekst-SQL przy użyciu bibliotek Hugging Face Transformers i TRL. Dowiesz się:

  • Co to jest kwantyzowana adaptacja o niskim rzędzie (QLoRA)
  • Konfigurowanie środowiska programistycznego
  • Tworzenie i przygotowywanie zbioru danych do dostrajania
  • Dostrajanie modelu Gemma za pomocą TRL i SFTTrainer
  • Testowanie wnioskowania modelu i generowanie zapytań SQL

Co to jest kwantyzowana adaptacja o niskim rzędzie (QLoRA)

W tym przewodniku pokazujemy, jak używać skwantyzowanej adaptacji o niskim rzędzie (QLoRA), która stała się popularną metodą wydajnego dostrajania dużych modeli językowych, ponieważ zmniejsza wymagania dotyczące zasobów obliczeniowych przy zachowaniu wysokiej wydajności. W przypadku QLoRA wstępnie wytrenowany model jest kwantyzowany do 4 bitów, a wagi są zamrażane. Następnie dołączane są warstwy adaptera z możliwością trenowania (LoRA) i trenowane są tylko te warstwy. Następnie wagi adaptera można połączyć z modelem podstawowym lub zachować jako osobny adapter.

Konfigurowanie środowiska programistycznego

Pierwszym krokiem jest zainstalowanie bibliotek Hugging Face, w tym TRL i zbiorów danych, aby dostroić otwarty model, w tym różne techniki RLHF i dopasowywania.

# Install Pytorch & other libraries
%pip install torch tensorboard

# Install Transformers
%pip install transformers

# Install Hugging Face libraries
%pip install datasets accelerate evaluate bitsandbytes trl peft protobuf sentencepiece

# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100
#%pip install flash-attn

Uwaga: jeśli używasz procesora GPU z architekturą Ampere (np. NVIDIA L4) lub nowszą, możesz korzystać z szybkiego mechanizmu uwagi. Flash Attention to metoda, która znacznie przyspiesza obliczenia i zmniejsza wykorzystanie pamięci z kwadratowego do liniowego w przypadku długości sekwencji, co przyspiesza trenowanie nawet 3-krotnie. Więcej informacji znajdziesz w artykule FlashAttention.

Aby opublikować model, musisz mieć ważny token Hugging Face. Jeśli korzystasz z Google Colab, możesz bezpiecznie używać tokena Hugging Face za pomocą sekretów Colab. W przeciwnym razie możesz ustawić token bezpośrednio w metodzie login. Upewnij się, że Twój token ma też uprawnienia do zapisu, ponieważ podczas trenowania modelu będziesz przesyłać go do Hub.

# Login into Hugging Face Hub
from huggingface_hub import login
login()

Tworzenie i przygotowywanie zbioru danych do dostrajania

Podczas dostrajania dużych modeli językowych ważne jest, aby znać przypadek użycia i zadanie, które chcesz rozwiązać. Pomoże Ci to utworzyć zbiór danych do dostrajania modelu. Jeśli nie masz jeszcze zdefiniowanego przypadku użycia, warto wrócić do etapu projektowania.

W tym przewodniku skupimy się na tym przykładzie:

  • Dostrój model konwersji z języka naturalnego na SQL, aby zapewnić płynną integrację z narzędziem do analizy danych. Celem jest znaczne skrócenie czasu i zmniejszenie poziomu wiedzy wymaganych do generowania zapytań SQL, co umożliwi nawet użytkownikom bez wiedzy technicznej wyodrębnianie z danych przydatnych informacji.

Przekształcanie tekstu w SQL może być dobrym przykładem zastosowania dostrajania LLM, ponieważ jest to złożone zadanie, które wymaga dużej ilości (wewnętrznej) wiedzy o danych i języku SQL.

Gdy stwierdzisz, że dostrajanie jest odpowiednim rozwiązaniem, potrzebujesz zbioru danych do dostrajania. Zestaw danych powinien zawierać różnorodne przykłady zadań, które chcesz rozwiązać. Istnieje kilka sposobów utworzenia takiego zbioru danych, m.in.:

  • Korzystanie z istniejących zbiorów danych typu open source, takich jak Spider.
  • Korzystanie z syntetycznych zbiorów danych utworzonych przez duże modele językowe, np. Alpaca.
  • Korzystanie ze zbiorów danych utworzonych przez ludzi, takich jak Dolly.
  • Używanie kombinacji metod, np. Orca

Każda z tych metod ma swoje zalety i wady, a jej wybór zależy od budżetu, czasu i wymagań dotyczących jakości. Na przykład użycie istniejącego zbioru danych jest najprostsze, ale może nie być dostosowane do konkretnego przypadku użycia, a skorzystanie z pomocy ekspertów w danej dziedzinie może być najbardziej precyzyjne, ale czasochłonne i kosztowne. Można też łączyć kilka metod, aby utworzyć zbiór danych instrukcji, jak pokazano w artykule Orca: Progressive Learning from Complex Explanation Traces of GPT-4.

W tym przewodniku używamy istniejącego już zbioru danych (philschmid/gretel-synthetic-text-to-sql), czyli wysokiej jakości syntetycznego zbioru danych Text-to-SQL, który zawiera instrukcje w języku naturalnym, definicje schematów, rozumowanie i odpowiednie zapytania SQL.

Hugging Face TRL obsługuje automatyczne tworzenie szablonów formatów zbiorów danych do rozmów. Oznacza to, że musisz tylko przekonwertować zbiór danych na odpowiednie obiekty JSON, a trl zajmie się tworzeniem szablonów i przekształcaniem ich do właściwego formatu.

{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

Zbiór philschmid/gretel-synthetic-text-to-sql zawiera ponad 100 tys. próbek. Aby zmniejszyć rozmiar przewodnika, zmniejszyliśmy liczbę próbek do 10 000.

Możesz teraz użyć biblioteki Hugging Face Datasets, aby wczytać zbiór danych i utworzyć szablon prompta, który połączy instrukcję w języku naturalnym, definicję schematu i doda wiadomość systemową dla asystenta.

from datasets import load_dataset

# System message for the assistant
system_message = """You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."""

# User prompt that combines the user query and the schema
user_prompt = """Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.

<SCHEMA>
{context}
</SCHEMA>

<USER_QUERY>
{question}
</USER_QUERY>
"""
def create_conversation(sample):
  return {
    "messages": [
      {"role": "system", "content": system_message},
      {"role": "user", "content": user_prompt.format(question=sample["sql_prompt"], context=sample["sql_context"])},
      {"role": "assistant", "content": sample["sql"]}
    ]
  }

# Load dataset from the hub
dataset = load_dataset("philschmid/gretel-synthetic-text-to-sql", split="train")
dataset = dataset.shuffle().select(range(12500))

# Convert dataset to OAI messages
dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
# split dataset into 80% training samples and 20% test samples
dataset = dataset.train_test_split(test_size=0.2)

# Print formatted user prompt
for item in dataset["train"][0]["messages"]:
  print(item)

README.md:   0%|          | 0.00/737 [00:00<?, ?B/s]
synthetic_text_to_sql_train.snappy.parqu(…):   0%|          | 0.00/32.4M [00:00<?, ?B/s]
synthetic_text_to_sql_test.snappy.parque(…):   0%|          | 0.00/1.90M [00:00<?, ?B/s]
Generating train split:   0%|          | 0/100000 [00:00<?, ? examples/s]
Generating test split:   0%|          | 0/5851 [00:00<?, ? examples/s]
Map:   0%|          | 0/12500 [00:00<?, ? examples/s]
{'content': 'You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.', 'role': 'system'}
{'content': "Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.\n\n<SCHEMA>\nCREATE TABLE Menu (id INT PRIMARY KEY, name VARCHAR(255), category VARCHAR(255), price DECIMAL(5,2));\n</SCHEMA>\n\n<USER_QUERY>\nCalculate the average price of all menu items in the Vegan category\n</USER_QUERY>\n", 'role': 'user'}
{'content': "SELECT AVG(price) FROM Menu WHERE category = 'Vegan';", 'role': 'assistant'}

Dostrajanie modelu Gemma za pomocą TRL i SFTTrainer

Możesz teraz dostroić model. Biblioteka Hugging Face TRL SFTTrainer ułatwia nadzorowane dostrajanie otwartych modeli LLM. SFTTrainer jest podklasą Trainer z biblioteki transformers i obsługuje wszystkie te same funkcje, w tym rejestrowanie, ocenianie i tworzenie punktów kontrolnych, ale dodaje dodatkowe funkcje ułatwiające pracę, takie jak:

  • Formatowanie zbioru danych, w tym formaty konwersacyjne i instruktażowe
  • Trenowanie tylko na podstawie uzupełnień, z pominięciem promptów
  • Pakowanie zbiorów danych w celu wydajniejszego trenowania
  • Obsługa dostrajania konkretnych parametrów (PEFT), w tym QLoRA
  • Przygotowanie modelu i tokenizera do dostrajania konwersacyjnego (np. dodanie tokenów specjalnych)

Poniższy kod wczytuje model Gemma i tokenizator z Hugging Face oraz inicjuje konfigurację kwantyzacji.

import torch
from transformers import AutoTokenizer, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-4-E2B" # @param ["google/gemma-4-E2B","google/gemma-4-E4B"] {"allow-input":true}

# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] >= 8:
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float16

# Define model init arguments
model_kwargs = dict(
    dtype=torch_dtype,
    device_map="auto", # Let torch decide how to load the model
)

# BitsAndBytesConfig: Enables 4-bit quantization to reduce model size/memory usage
model_kwargs["quantization_config"] = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=model_kwargs['dtype'],
    bnb_4bit_quant_storage=model_kwargs['dtype'],
)

# Load model and tokenizer
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E2B-it") # Load the Instruction Tokenizer to use the official Gemma template

config.json: 0.00B [00:00, ?B/s]
model.safetensors:   0%|          | 0.00/10.2G [00:00<?, ?B/s]
Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
generation_config.json:   0%|          | 0.00/181 [00:00<?, ?B/s]
config.json: 0.00B [00:00, ?B/s]
tokenizer_config.json: 0.00B [00:00, ?B/s]
tokenizer.json:   0%|          | 0.00/32.2M [00:00<?, ?B/s]
chat_template.jinja: 0.00B [00:00, ?B/s]

SFTTrainer obsługuje wbudowaną integrację z peft, co ułatwia efektywne dostrajanie LLM za pomocą QLoRA. Wystarczy, że utworzysz LoraConfig i przekażesz go trenerowi.

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=["lm_head", "embed_tokens"], # make sure to save the lm_head and embed_tokens as you train the special tokens
    ensure_weight_tying=True,
)

Zanim rozpoczniesz trenowanie, musisz zdefiniować hiperparametr, którego chcesz użyć w instancji SFTConfig.

import torch
from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-text-to-sql",         # directory to save and repository id
    max_length=512,                         # max length for model and packing of the dataset
    num_train_epochs=3,                     # number of training epochs
    per_device_train_batch_size=1,          # batch size per device during training
    optim="adamw_torch_fused",              # use fused adamw optimizer
    logging_steps=10,                       # log every 10 steps
    save_strategy="epoch",                  # save checkpoint every epoch
    eval_strategy="epoch",                  # evaluate checkpoint every epoch
    learning_rate=5e-5,                     # learning rate
    fp16=True if model.dtype == torch.float16 else False,   # use float16 precision
    bf16=True if model.dtype == torch.bfloat16 else False,   # use bfloat16 precision
    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper
    lr_scheduler_type="constant",           # use constant learning rate scheduler
    push_to_hub=True,                           # push model to hub
    report_to="tensorboard",                # report metrics to tensorboard
    dataset_kwargs={
        "add_special_tokens": False, # Template with special tokens
        "append_concat_token": True, # Add EOS token as separator token between examples
    }
)

Masz już wszystkie elementy składowe potrzebne do utworzenia SFTTrainer, aby rozpocząć trenowanie modelu.

from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    peft_config=peft_config,
    processing_class=tokenizer,
)

Tokenizing train dataset:   0%|          | 0/10000 [00:00<?, ? examples/s]
Tokenizing eval dataset:   0%|          | 0/2500 [00:00<?, ? examples/s]

Rozpocznij trenowanie, wywołując metodę train().

# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()

# Save the final model again to the Hugging Face Hub
trainer.save_model()

The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 1, 'bos_token_id': 2, 'pad_token_id': 0}.

Zanim przetestujesz model, zwolnij pamięć.

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

W przypadku QLoRA trenujesz tylko adaptery, a nie cały model. Oznacza to, że podczas zapisywania modelu w trakcie trenowania zapisujesz tylko wagi adaptera, a nie cały model. Jeśli chcesz zapisać pełny model, co ułatwia korzystanie z platform do udostępniania, takich jak vLLM czy TGI, możesz scalić wagi adaptera z wagami modelu za pomocą metody merge_and_unload, a następnie zapisać model za pomocą metody save_pretrained. Spowoduje to zapisanie modelu domyślnego, który można wykorzystać do wnioskowania.

from peft import PeftModel

# Load Model base model
model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)

# Merge LoRA and base model and save
peft_model = PeftModel.from_pretrained(model, args.output_dir)
merged_model = peft_model.merge_and_unload()
merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")

processor = AutoTokenizer.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")

Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
Writing model shards:   0%|          | 0/5 [00:00<?, ?it/s]
('merged_model/tokenizer_config.json',
 'merged_model/chat_template.jinja',
 'merged_model/tokenizer.json')

Testowanie wnioskowania modelu i generowanie zapytań SQL

Po zakończeniu trenowania warto ocenić i przetestować model. Możesz wczytać różne próbki ze zbioru danych testowych i ocenić na nich model.

import torch
from transformers import pipeline

model_id = "merged_model"

# Load Model with PEFT adapter
model = AutoModelForImageTextToText.from_pretrained(
  model_id,
  device_map="auto",
  dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

Loading weights:   0%|          | 0/2012 [00:00<?, ?it/s]
The tied weights mapping and config for this model specifies to tie model.language_model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints with different values, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning.

Wczytajmy losową próbkę z testowego zbioru danych i wygenerujmy polecenie SQL.

from random import randint
import re
from transformers import pipeline, GenerationConfig

config = GenerationConfig.from_pretrained(model_id)
config.max_new_tokens = 256

# Load the model and tokenizer into the pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Load a random sample from the test dataset
rand_idx = randint(0, len(dataset["test"]))
test_sample = dataset["test"][rand_idx]

# Convert as test example into a prompt with the Gemma template
prompt = pipe.tokenizer.apply_chat_template(test_sample["messages"][:2], tokenize=False, add_generation_prompt=True)
print(prompt)

# Generate our SQL query.
outputs = pipe(prompt, generation_config=config)

# Extract the user query and original answer
print(f"Context:\n", re.search(r'<SCHEMA>\n(.*?)\n</SCHEMA>', test_sample['messages'][1]['content'], re.DOTALL).group(1).strip())
print(f"Query:\n", re.search(r'<USER_QUERY>\n(.*?)\n</USER_QUERY>', test_sample['messages'][1]['content'], re.DOTALL).group(1).strip())
print(f"Original Answer:\n{test_sample['messages'][2]['content']}")
print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")

<bos><|turn>system
You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.<turn|>
<|turn>user
Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.

<SCHEMA>
CREATE TABLE broadband_plans (plan_id INT, plan_name VARCHAR(255), download_speed INT, upload_speed INT, price DECIMAL(5,2));
</SCHEMA>

<USER_QUERY>
Delete a broadband plan from the 'broadband_plans' table
</USER_QUERY><turn|>
<|turn>model

Context:
 CREATE TABLE broadband_plans (plan_id INT, plan_name VARCHAR(255), download_speed INT, upload_speed INT, price DECIMAL(5,2));
Query:
 Delete a broadband plan from the 'broadband_plans' table
Original Answer:
DELETE FROM broadband_plans WHERE plan_id = 3001;
Generated Answer:
DELETE FROM broadband_plans
WHERE plan_name = 'Basic';

Podsumowanie i dalsze kroki

W tym samouczku pokazaliśmy, jak dostroić model Gemma za pomocą TRL i QLoRA. Zapoznaj się z tymi dokumentami: