Gemma mit Hugging Face Transformers und QloRA optimieren

Auf ai.google.dev ansehen In Google Colab ausführen In Kaggle ausführen In Vertex AI öffnen Quelle auf GitHub ansehen

In dieser Anleitung erfahren Sie, wie Sie Gemma mit Hugging Face Transformers und TRL auf einem benutzerdefinierten Text-zu-SQL-Dataset optimieren. Lerninhalte:

  • Was ist Quantized Low-Rank Adaptation (QLoRA)?
  • Entwicklungsumgebung einrichten
  • Dataset für die Feinabstimmung erstellen und vorbereiten
  • Gemma mit TRL und SFTTrainer optimieren
  • Modellinferenz testen und SQL-Abfragen generieren

Was ist Quantized Low-Rank Adaptation (QLoRA)?

In dieser Anleitung wird die Verwendung von Quantized Low-Rank Adaptation (QLoRA) veranschaulicht. Diese Methode hat sich als beliebte Methode zur effizienten Feinabstimmung von LLMs etabliert, da sie die Anforderungen an Rechenressourcen reduziert und gleichzeitig eine hohe Leistung beibehält. Bei QLoRA wird das vortrainierte Modell auf 4 Bit quantisiert und die Gewichte werden eingefroren. Anschließend werden trainierbare Adapterschichten (LoRA) angehängt und nur die Adapterschichten werden trainiert. Danach können die Adaptergewichte mit dem Basismodell zusammengeführt oder als separater Adapter beibehalten werden.

Entwicklungsumgebung einrichten

Im ersten Schritt müssen Sie Hugging Face-Bibliotheken, einschließlich TRL und Datasets, installieren, um das offene Modell zu optimieren. Dazu gehören verschiedene RLHF- und Alignment-Techniken.

# 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

Hinweis: Wenn Sie eine GPU mit Ampere-Architektur (z. B. NVIDIA L4) oder höher verwenden, können Sie Flash Attention verwenden. Flash Attention ist eine Methode, die Berechnungen erheblich beschleunigt und die Arbeitsspeichernutzung von quadratisch auf linear in der Sequenzlänge reduziert, was zu einer bis zu dreifachen Beschleunigung des Trainings führt. Weitere Informationen finden Sie unter FlashAttention.

Sie benötigen ein gültiges Hugging Face-Token, um Ihr Modell zu veröffentlichen. Wenn Sie in Google Colab arbeiten, können Sie Ihr Hugging Face-Token sicher mit Colab-Secrets verwenden. Andernfalls können Sie das Token direkt in der Methode login festlegen. Achten Sie darauf, dass Ihr Token auch Schreibzugriff hat, da Sie Ihr Modell während des Trainings in den Hub hochladen.

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

Dataset für die Feinabstimmung erstellen und vorbereiten

Bei der Feinabstimmung von LLMs ist es wichtig, den Anwendungsfall und die Aufgabe zu kennen, die Sie lösen möchten. So können Sie ein Dataset erstellen, um Ihr Modell zu optimieren. Wenn Sie Ihren Anwendungsfall noch nicht definiert haben, sollten Sie noch einmal von vorn beginnen.

In dieser Anleitung wird beispielsweise der folgende Anwendungsfall behandelt:

  • Ein Modell für die Umwandlung von natürlicher Sprache in SQL für die nahtlose Einbindung in ein Datenanalysetool optimieren. Ziel ist es, den Zeit- und Fachwissenaufwand für die Generierung von SQL-Abfragen erheblich zu reduzieren, sodass auch Nutzer ohne technisches Vorwissen aussagekräftige Informationen aus Daten gewinnen können.

Die Umwandlung von Text in SQL kann ein guter Anwendungsfall für die Feinabstimmung von LLMs sein, da es sich um eine komplexe Aufgabe handelt, die viel (internes) Wissen über die Daten und die SQL-Sprache erfordert.

Sobald Sie festgestellt haben, dass die Feinabstimmung die richtige Lösung ist, benötigen Sie ein Dataset für die Feinabstimmung. Das Dataset sollte eine vielfältige Reihe von Demonstrationen der Aufgaben enthalten, die Sie lösen möchten. Es gibt verschiedene Möglichkeiten, ein solches Dataset zu erstellen, darunter:

  • Vorhandene Open-Source-Datasets wie Spider verwenden
  • Synthetische Datasets verwenden, die von LLMs wie Alpaca erstellt wurden
  • Von Menschen erstellte Datasets wie Dolly verwenden
  • Eine Kombination der Methoden verwenden, z. B. Orca

Jede Methode hat ihre eigenen Vor- und Nachteile und hängt vom Budget, der Zeit und den Qualitätsanforderungen ab. Die Verwendung eines vorhandenen Datasets ist beispielsweise die einfachste Methode, ist aber möglicherweise nicht auf Ihren spezifischen Anwendungsfall zugeschnitten. Die Verwendung von Fachexperten ist möglicherweise die genaueste Methode, kann aber zeitaufwendig und teuer sein. Es ist auch möglich, mehrere Methoden zu kombinieren, um ein Dataset mit Anweisungen zu erstellen, wie in Orca: Progressive Learning from Complex Explanation Traces of GPT-4. gezeigt.

In dieser Anleitung wird ein bereits vorhandenes Dataset verwendet (philschmid/gretel-synthetic-text-to-sql). Es handelt sich um ein hochwertiges synthetisches Text-zu-SQL-Dataset mit Anweisungen in natürlicher Sprache, Schemadefinitionen, Begründungen und der entsprechenden SQL-Abfrage.

Hugging Face TRL unterstützt die automatische Vorlagenerstellung für Konversationsdataset-Formate. Das bedeutet, dass Sie Ihr Dataset nur in die richtigen JSON-Objekte konvertieren müssen. trl kümmert sich um die Vorlagenerstellung und die Formatierung.

{"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": "..."}]}

Das Dataset philschmid/gretel-synthetic-text-to-sql enthält über 100.000 Beispiele. Um die Anleitung kurz zu halten, wird es auf 10.000 Beispiele reduziert.

Sie können jetzt die Hugging Face Datasets-Bibliothek verwenden, um das Dataset zu laden und eine Prompt-Vorlage zu erstellen, um die Anweisung in natürlicher Sprache und die Schemadefinition zu kombinieren und eine Systemnachricht für Ihren Assistenten hinzuzufügen.

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'}

Gemma mit TRL und SFTTrainer optimieren

Jetzt können Sie Ihr Modell optimieren. Mit Hugging Face TRL SFTTrainer können Sie offene LLMs ganz einfach optimieren. Der SFTTrainer ist eine Unterklasse von Trainer aus der Bibliothek transformers und unterstützt alle Funktionen, einschließlich Logging, Bewertung und Checkpointing. Außerdem bietet er zusätzliche Funktionen, die die Nutzung erleichtern:

  • Dataset-Formatierung, einschließlich Konversations- und Anweisungsformaten
  • Training nur mit Vervollständigungen, Prompts werden ignoriert
  • Datasets für ein effizienteres Training packen
  • Unterstützung für parametereffiziente Feinabstimmung (PEFT), einschließlich QLoRA
  • Vorbereiten des Modells und des Tokenizers für die Konversationsabstimmung (z. B. Hinzufügen von Sondertokens)

Der folgende Code lädt das Gemma-Modell und den Tokenizer von Hugging Face und initialisiert die Quantisierungskonfiguration.

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]

Der SFTTrainer unterstützt eine integrierte Einbindung von peft, mit der Sie LLMs effizient mit QLoRA optimieren können. Sie müssen nur eine LoraConfig erstellen und sie dem Trainer zur Verfügung stellen.

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,
)

Bevor Sie mit dem Training beginnen können, müssen Sie die Hyperparameter definieren, die Sie in einer SFTConfig-Instanz verwenden möchten.

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
    }
)

Jetzt haben Sie alle Bausteine, die Sie zum Erstellen Ihres SFTTrainer benötigen, um mit dem Training Ihres Modells zu beginnen.

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]

Starten Sie das Training mit der Methode 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}.

Bevor Sie Ihr Modell testen können, müssen Sie den Arbeitsspeicher freigeben.

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

Bei der Verwendung von QLoRA werden nur Adapter und nicht das vollständige Modell trainiert. Wenn Sie das Modell während des Trainings speichern, werden daher nur die Adaptergewichte und nicht das vollständige Modell gespeichert. Wenn Sie das vollständige Modell speichern möchten, um es einfacher mit Bereitstellungsstacks wie vLLM oder TGI zu verwenden, können Sie die Adaptergewichte mit der Methode merge_and_unload in die Modellgewichte zusammenführen und das Modell dann mit der Methode save_pretrained speichern. Dadurch wird ein Standardmodell gespeichert, das für die Inferenz verwendet werden kann.

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')

Modellinferenz testen und SQL-Abfragen generieren

Nach Abschluss des Trainings sollten Sie Ihr Modell bewerten und testen. Sie können verschiedene Beispiele aus dem Test-Dataset laden und das Modell anhand dieser Beispiele bewerten.

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.

Laden wir ein zufälliges Beispiel aus dem Test-Dataset und generieren wir einen SQL-Befehl.

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';

Zusammenfassung und nächste Schritte

In dieser Anleitung wurde beschrieben, wie Sie ein Gemma-Modell mit TRL und QLoRA optimieren. Weitere Informationen finden Sie in den folgenden Dokumenten: