0dstqkcyx56nl4qyu.png

Una técnica de ajuste unificada más barata y rápida

Imagen generada con DALL-E 3 por autor

ORPO es un nueva y emocionante técnica de ajuste que combina las tradicionales etapas supervisadas de ajuste fino y alineación de preferencias en un solo proceso. Esto reduce los recursos computacionales y el tiempo requerido para la capacitación. Además, los resultados empíricos demuestran que ORPO supera a otros métodos de alineación en varios tamaños de modelos y puntos de referencia.

En este artículo, ajustaremos el nuevo modelo Llama 3 8B usando ORPO con la biblioteca TRL. El código está disponible en colaboración de google y en el Curso de Maestría en Derecho en GitHub.

El ajuste de instrucciones y la alineación de preferencias son técnicas esenciales para adaptar los modelos de lenguaje grande (LLM) a tareas específicas. Tradicionalmente, esto implica un proceso de varias etapas: 1/ Ajuste supervisado (SFT) sobre instrucciones para adaptar el modelo al dominio objetivo, seguido de 2/ métodos de alineación de preferencias como el aprendizaje por refuerzo con retroalimentación humana (RLHF) o la optimización de preferencias directas (DPO) para aumentar la probabilidad de generar respuestas preferidas sobre las rechazadas.

Imagen del autor

Sin embargo, los investigadores han identificado una limitación en este enfoque. Si bien SFT adapta efectivamente el modelo al dominio deseado, sin darse cuenta aumenta la probabilidad de generar respuestas indeseables junto con los preferidos. Esta es la razón por la que la etapa de alineación de preferencias es necesaria para ampliar la brecha entre las probabilidades de resultados preferidos y rechazados.

Observe cómo la probabilidad de respuestas rechazadas aumenta durante el ajuste supervisado (imagen del artículo de ORPO).

Presentado por Hong y Lee (2024), ORPO ofrece una solución elegante a este problema al combinar el ajuste de instrucciones y la alineación de preferencias en un proceso de capacitación único y monolítico. ORPO modifica el objetivo de modelado del lenguaje estándar, combinando la pérdida de probabilidad logarítmica negativa con un término de odds ratio (OR). Esta pérdida de OR penaliza débilmente las respuestas rechazadas mientras recompensa fuertemente las preferidas, lo que permite que el modelo aprenda simultáneamente la tarea objetivo y se alinee con las preferencias humanas.

ORPO se ha implementado en las principales bibliotecas de ajuste, como TRL, Ajolotey LLaMA-Fábrica. En la siguiente sección, veremos cómo utilizar TRL.

Llama 3 es la última familia de LLM desarrollada por Meta. Los modelos fueron entrenados en un extenso conjunto de datos de 15 billones de fichas (en comparación con las fichas 2T de Llama 2). Se han lanzado dos tamaños de modelo: un modelo de 70 mil millones de parámetros y un modelo más pequeño de 8 mil millones de parámetros. El modelo 70B ya ha demostrado un rendimiento impresionante, con una puntuación de 82 en el punto de referencia MMLU y 81,7 en el punto de referencia HumanEval.

Los modelos Llama 3 también aumentaron la longitud del contexto hasta 8192 tokens (4096 tokens para Llama 2) y potencialmente escalar hasta 32k con RoPE. Además, los modelos utilizan un nuevo tokenizador con un vocabulario de 128.000 tokens, lo que reduce la cantidad de tokens necesarios para codificar texto en un 15 %. Este vocabulario también explica el salto de los parámetros 7B a 8B.

Muestras de ORPO-DPO-mix-40k (imagen del autor).

ORPO requiere un conjunto de datos de preferencias, que incluye una indicación, una respuesta elegida y una respuesta rechazada. En este ejemplo, usaremos mlabonne/orpo-dpo-mix-40kuna combinación de los siguientes conjuntos de datos de DPO de alta calidad:

Gracias a argilla, desalineación, M4-aiy jondurbin para proporcionar los conjuntos de datos de origen.

Como de costumbre, comencemos instalando las bibliotecas necesarias:

pip install -U transformers datasets accelerate peft trl bitsandbytes wandb

Una vez instalado, podemos importar las bibliotecas necesarias e iniciar sesión en W&B (opcional):

import gc
import os

import torch
import wandb
from datasets import load_dataset
from google.colab import userdata
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
TrainingArguments,
pipeline,
)
from trl import ORPOConfig, ORPOTrainer, setup_chat_format
wb_token = userdata.get('wandb')
wandb.login(key=wb_token)

Si tiene una GPU reciente, también debería poder usar la Biblioteca de atención flash para reemplazar la implementación predeterminada de atención ansiosa por una más eficiente.

if torch.cuda.get_device_capability()[0] >= 8:
!pip install -qqq flash-attn
attn_implementation = "flash_attention_2"
torch_dtype = torch.bfloat16
else:
attn_implementation = "eager"
torch_dtype = torch.float16

A continuación, cargaremos el modelo Llama 3 8B con precisión de 4 bits gracias a bits y bytes. Luego configuramos la configuración de LoRA usando PEFT para QLoRA. También estoy usando el conveniente setup_chat_format() función para modificar el modelo y el tokenizador para ChatML apoyo. Aplica automáticamente esta plantilla de chat, agrega tokens especiales y cambia el tamaño de la capa de incrustación del modelo para que coincida con el nuevo tamaño del vocabulario.

# Model
base_model = "meta-llama/Meta-Llama-3-8B"
new_model = "OrpoLlama-3-8B"

# QLoRA config
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch_dtype,
bnb_4bit_use_double_quant=True,
)

# LoRA config
peft_config = LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
target_modules=['up_proj', 'down_proj', 'gate_proj', 'k_proj', 'q_proj', 'v_proj', 'o_proj']
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Load model
model = AutoModelForCausalLM.from_pretrained(
base_model,
quantization_config=bnb_config,
device_map="auto",
attn_implementation=attn_implementation
)
model, tokenizer = setup_chat_format(model, tokenizer)
model = prepare_model_for_kbit_training(model)

Ahora que el modelo está listo para entrenar, podemos encargarnos del conjunto de datos. cargamos mlabonne/orpo-dpo-mix-40k y usar el apply_chat_template() función para convertir las columnas «elegidas» y «rechazadas» al formato ChatML. Tenga en cuenta que solo estoy usando 1000 muestras y no todo el conjunto de datos, ya que llevaría demasiado tiempo ejecutarlo.

dataset_name = "mlabonne/orpo-dpo-mix-40k"
dataset = load_dataset(dataset_name, split="all")
dataset = dataset.shuffle(seed=42).select(range(10))

def format_chat_template(row):
row["chosen"] = tokenizer.apply_chat_template(row["chosen"], tokenize=False)
row["rejected"] = tokenizer.apply_chat_template(row["rejected"], tokenize=False)
return row

dataset = dataset.map(
format_chat_template,
num_proc= os.cpu_count(),
)
dataset = dataset.train_test_split(test_size=0.01)

Primero, necesitamos establecer algunos hiperparámetros:

  • learning_rate: ORPO utiliza tasas de aprendizaje muy bajas en comparación con las SFT tradicionales o incluso con las DPO. Este valor de 8e-6 proviene del artículo original y corresponde aproximadamente a una tasa de aprendizaje SFT de 1e-5 y una tasa de aprendizaje DPO de 5e-6. Recomendaría aumentarlo alrededor de 1e-6 para un ajuste realmente fino.
  • beta: Es el parámetro $\lambda$ del artículo, con un valor predeterminado de 0,1. Un apéndice del artículo original muestra cómo se seleccionó con un estudio de ablación.
  • Otros parámetros, como max_length y el tamaño del lote están configurados para utilizar tanta VRAM como esté disponible (~20 GB en esta configuración). Lo ideal sería entrenar el modelo durante 3 a 5 épocas, pero aquí nos limitaremos a 1.

Finalmente, podemos entrenar el modelo usando ORPOTrainer, que actúa como contenedor.

orpo_args = ORPOConfig(
learning_rate=8e-6,
beta=0.1,
lr_scheduler_type="linear",
max_length=1024,
max_prompt_length=512,
per_device_train_batch_size=2,
per_device_eval_batch_size=2,
gradient_accumulation_steps=4,
optim="paged_adamw_8bit",
num_train_epochs=1,
evaluation_strategy="steps",
eval_steps=0.2,
logging_steps=1,
warmup_steps=10,
report_to="wandb",
output_dir="./results/",
)

trainer = ORPOTrainer(
model=model,
args=orpo_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
peft_config=peft_config,
tokenizer=tokenizer,
)

trainer.train()
trainer.save_model(new_model)

Entrenar el modelo con estas 1000 muestras tomó aproximadamente 2 horas en una GPU L4. Revisemos los gráficos de W&B:

Si bien la pérdida disminuye, la diferencia entre las respuestas elegidas y rechazadas no es clara: el margen promedio y la precisión están sólo ligeramente por encima de cero y 0,5, respectivamente.

En el artículo original, los autores entrenaron modelos en el Anthropic/hh-rlhf conjunto de datos (161k muestras) durante 10 épocas, que es mucho más largo que nuestra ejecución rápida. También experimentaron con Llama 3 y amablemente compartieron sus registros conmigo (gracias Ji Woo Hong).

Para finalizar este tutorial, fusionemos el adaptador QLoRA con el modelo base y lo empujemos al Hugging Face Hub.

# Flush memory
del trainer, model
gc.collect()
torch.cuda.empty_cache()

# Reload tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
base_model,
low_cpu_mem_usage=True,
return_dict=True,
torch_dtype=torch.float16,
device_map="auto",
)
model, tokenizer = setup_chat_format(model, tokenizer)

# Merge adapter with base model
model = PeftModel.from_pretrained(model, new_model)
model = model.merge_and_unload()
model.push_to_hub(new_model, use_temp_dir=False)
tokenizer.push_to_hub(new_model, use_temp_dir=False)

Felicitaciones, terminamos este rápido ajuste de Llama 3: mlabonne/OrpoLlama-3–8B. Puedes jugar con él usando esto. Abrazando el espacio de la cara (Aquí está un computadora portátil para hacer el tuyo propio). Aunque el modelo no está suficientemente entrenado, como lo destacan las curvas W&B, realicé algunas evaluaciones en el conjunto de pruebas comparativas de Nous utilizando LLM AutoEval.

Nuestro ajuste ORPO es bastante decente y mejora el rendimiento del modelo base en cada punto de referencia. Esto es alentador y probablemente signifique que un ajuste fino en las 40.000 muestras completas produciría excelentes resultados.

Este es un momento emocionante para la comunidad de código abierto, con el lanzamiento de cada vez más modelos de peso abierto de alta calidad. La brecha entre los modelos de código cerrado y de peso abierto se está cerrando lentamente, y el ajuste es una herramienta esencial para obtener el mejor rendimiento para sus casos de uso.

Imagen del autor

En este artículo, presentamos el algoritmo ORPO y explicamos cómo unifica las etapas SFT y de alineación de preferencias en un solo proceso. Luego, utilizamos TRL para ajustar un modelo Llama 3 8B en un conjunto de datos de preferencias personalizado. El modelo final muestra resultados alentadores y destaca el potencial de ORPO como un nuevo paradigma de ajuste.

Espero que haya sido útil y recomiendo ejecutar el Cuaderno de colaboración para afinar tus propios modelos Llama 3. En artículos futuros veremos cómo crear conjuntos de datos de alta calidad, un punto que a menudo se pasa por alto. Si te gustó este artículo, sígueme en abrazando la cara y Twitter @maximelabonne.