Affiner Gemma à l'aide de Hugging Face Transformers et de QloRA

 Afficher sur ai.google.dev  Exécuter dans Google Colab Exécuter dans Kaggle Ouvrir dans Vertex AI  Afficher la source sur GitHub

Ce guide vous explique comment affiner Gemma sur un ensemble de données texte-vers-SQL personnalisé à l'aide de Transformers et TRL de Hugging Face. Vous pourrez découvrir :

  • Qu'est-ce que l'adaptation à faible rang quantifiée (QLoRA) ?
  • Configurer l'environnement de développement
  • Créer et préparer l'ensemble de données de réglage fin
  • Affiner Gemma à l'aide de TRL et de SFTTrainer
  • Tester l'inférence du modèle et générer des requêtes SQL

Qu'est-ce que l'adaptation à faible rang quantifiée (QLoRA) ?

Ce guide explique comment utiliser Quantized Low-Rank Adaptation (QLoRA), une méthode populaire pour affiner efficacement les LLM, car elle réduit les besoins en ressources de calcul tout en maintenant des performances élevées. Dans QloRA, le modèle pré-entraîné est quantifié à 4 bits et les pondérations sont figées. Des calques d'adaptateur entraînables (LoRA) sont ensuite associés, et seuls les calques d'adaptateur sont entraînés. Les pondérations de l'adaptateur peuvent ensuite être fusionnées avec le modèle de base ou conservées en tant qu'adaptateur distinct.

Configurer l'environnement de développement

La première étape consiste à installer les bibliothèques Hugging Face, y compris TRL et les ensembles de données, pour affiner le modèle ouvert, y compris différentes techniques de RLHF et d'alignement.

# 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

Remarque : Si vous utilisez un GPU avec une architecture Ampere (comme NVIDIA L4) ou plus récente, vous pouvez utiliser Flash Attention. Flash Attention est une méthode qui accélère considérablement les calculs et réduit l'utilisation de la mémoire de quadratique à linéaire en termes de longueur de séquence, ce qui permet d'accélérer l'entraînement jusqu'à trois fois. Pour en savoir plus, consultez FlashAttention.

Pour publier votre modèle, vous avez besoin d'un jeton Hugging Face valide. Si vous exécutez le code dans Google Colab, vous pouvez utiliser votre jeton Hugging Face de manière sécurisée à l'aide des secrets Colab. Sinon, vous pouvez définir le jeton directement dans la méthode login. Assurez-vous que votre jeton dispose également d'un accès en écriture, car vous allez transférer votre modèle vers le Hub pendant l'entraînement.

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

Créer et préparer l'ensemble de données de réglage fin

Lorsque vous ajustez des LLM, il est important de connaître votre cas d'utilisation et la tâche que vous souhaitez résoudre. Cela vous aide à créer un ensemble de données pour affiner votre modèle. Si vous n'avez pas encore défini votre cas d'utilisation, vous pouvez revenir en arrière.

Par exemple, ce guide se concentre sur le cas d'utilisation suivant :

  • Ajustez un modèle de langage naturel vers SQL pour l'intégrer de manière fluide dans un outil d'analyse de données. L'objectif est de réduire considérablement le temps et l'expertise nécessaires à la génération de requêtes SQL, afin de permettre même aux utilisateurs non techniques d'extraire des insights pertinents à partir des données.

La conversion de texte en SQL peut être un bon cas d'utilisation pour l'affinage des LLM, car il s'agit d'une tâche complexe qui nécessite de nombreuses connaissances (internes) sur les données et le langage SQL.

Une fois que vous avez déterminé que le réglage fin est la solution appropriée, vous avez besoin d'un ensemble de données pour le réglage fin. L'ensemble de données doit être un ensemble varié de démonstrations des tâches que vous souhaitez résoudre. Il existe plusieurs façons de créer un tel ensemble de données, y compris :

  • Utiliser des ensembles de données Open Source existants, tels que Spider
  • Utilisation d'ensembles de données synthétiques créés par des LLM, tels qu'Alpaca
  • Utiliser des ensembles de données créés par des humains, comme Dolly.
  • Utiliser une combinaison de méthodes, comme Orca

Chacune de ces méthodes présente ses propres avantages et inconvénients, et dépend des exigences en termes de budget, de temps et de qualité. Par exemple, l'utilisation d'un ensemble de données existant est la plus simple, mais peut ne pas être adaptée à votre cas d'utilisation spécifique. L'utilisation d'experts du domaine peut être la plus précise, mais peut prendre du temps et être coûteuse. Il est également possible de combiner plusieurs méthodes pour créer un ensemble de données d'instructions, comme indiqué dans Orca: Progressive Learning from Complex Explanation Traces of GPT-4.

Ce guide utilise un ensemble de données existant (philschmid/gretel-synthetic-text-to-sql), un ensemble de données Text-to-SQL synthétiques de haute qualité incluant des instructions en langage naturel, des définitions de schéma, un raisonnement et la requête SQL correspondante.

Hugging Face TRL est compatible avec la création automatique de modèles pour les formats d'ensembles de données de conversation. Cela signifie que vous n'avez qu'à convertir votre ensemble de données en objets JSON appropriés. trl s'occupe de la création de modèles et de la mise en forme.

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

Le jeu de données philschmid/gretel-synthetic-text-to-sql contient plus de 100 000 échantillons. Pour que le guide reste petit, il est sous-échantillonné pour n'utiliser que 10 000 échantillons.

Vous pouvez désormais utiliser la bibliothèque Hugging Face Datasets pour charger l'ensemble de données et créer un modèle d'invite afin de combiner l'instruction en langage naturel et la définition du schéma, et d'ajouter un message système pour votre assistant.

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

Affiner Gemma à l'aide de TRL et de SFTTrainer

Vous êtes maintenant prêt à affiner votre modèle. Hugging Face TRL SFTTrainer facilite le réglage fin supervisé des LLM Open Source. SFTTrainer est une sous-classe de Trainer de la bibliothèque transformers et prend en charge les mêmes fonctionnalités, y compris la journalisation, l'évaluation et la création de points de contrôle, mais ajoute des fonctionnalités supplémentaires de qualité de vie, y compris :

  • Mise en forme des ensembles de données, y compris les formats conversationnels et d'instructions
  • Entraînement sur les complétions uniquement, en ignorant les requêtes
  • Regrouper les ensembles de données pour un entraînement plus efficace
  • Prise en charge de l'affinage des paramètres avec optimisation (PEFT), y compris QLoRA
  • Préparer le modèle et le tokenizer pour l'affinage conversationnel (par exemple, en ajoutant des jetons spéciaux)

Le code suivant charge le modèle et le tokenizer Gemma depuis Hugging Face et initialise la configuration de quantification.

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 est compatible avec une intégration intégrée à peft, ce qui permet d'ajuster efficacement les LLM à l'aide de QLoRA. Il vous suffit de créer un LoraConfig et de le fournir au formateur.

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

Avant de pouvoir commencer votre entraînement, vous devez définir l'hyperparamètre que vous souhaitez utiliser dans une instance 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
    }
)

Vous disposez désormais de tous les éléments nécessaires pour créer votre SFTTrainer et commencer l'entraînement de votre modèle.

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]

Commencez l'entraînement en appelant la méthode 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}.

Avant de pouvoir tester votre modèle, assurez-vous de libérer de la mémoire.

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

Lorsque vous utilisez QLoRA, vous n'entraînez que les adaptateurs et non le modèle complet. Cela signifie que, lorsque vous enregistrez le modèle pendant l'entraînement, vous n'enregistrez que les pondérations de l'adaptateur et non le modèle complet. Si vous souhaitez enregistrer le modèle complet, ce qui facilite son utilisation avec des piles de diffusion telles que vLLM ou TGI, vous pouvez fusionner les pondérations de l'adaptateur dans les pondérations du modèle à l'aide de la méthode merge_and_unload, puis enregistrer le modèle avec la méthode save_pretrained. Cela permet d'enregistrer un modèle par défaut, qui peut être utilisé pour l'inférence.

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

Tester l'inférence du modèle et générer des requêtes SQL

Une fois l'entraînement terminé, vous devez évaluer et tester votre modèle. Vous pouvez charger différents échantillons à partir de l'ensemble de données de test et évaluer le modèle sur ces échantillons.

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.

Chargeons un échantillon aléatoire de l'ensemble de données de test et générons une commande 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';

Résumé et étapes suivantes

Ce tutoriel explique comment affiner un modèle Gemma à l'aide de TRL et QLoRA. Consultez ensuite les documents suivants :