CausalLM Parte 2: Ajuste fino de un modelo |  de Theo Lebryk |  marzo de 2024

Tres formas de ajustar un modelo CausalLM en datos de chat

En este tutorial, ajustaremos un modelo CausalLM para realizar una traducción simple. Foto por Rob Wilson en desempaquetar

En el ultima publicación, hablamos sobre qué es CausalLM y cómo Hugging Face espera que se formatee los datos. En esta publicación, analizaremos un cuaderno resumido con tres formas de formatear los datos para ajustar un modelo. El primero es un enfoque sencillo que se basa en la intuición de la publicación anterior, simplemente copiando input_ids en etiquetas. El segundo enfoque utiliza enmascaramiento para aprender partes seleccionadas del texto. El tercer enfoque utiliza una biblioteca separada, TRLpara que no tengamos que enmascarar manualmente los datos.

Dejaré de lado algunas definiciones de funciones para que sean legibles, por lo que es mejor hacer referencia el cuaderno completok para obtener todo el código.

Ajuste fino con etiquetas copiadas de identificadores de entrada

vamos a estar usando Floración-560mun modelo multilingüe lo suficientemente pequeño como para que podamos ajustarlo en un portátil estándar.

model_name = "bigscience/bloom-560m"
tokenizer = AutoTokenizer.from_pretrained(
model_name, trust_remote_code=True, padding_side="right"
) # padding side should be right for CausalLM models
# overfit to 5 made up examples
str1 = '\n\n### Human: How do you say "dog" in Spanish?\n\n### Assistant: perro'
str2 = '\n\n### Human: How do you say "water" in Spanish?\n\n### Assistant: agua'
str3 = '\n\n### Human: How do you say "hello" in Spanish?\n\n### Assistant: hola'
str4 = '\n\n### Human: How do you say "tree" in Spanish?\n\n### Assistant: árbol'
str5 = '\n\n### Human: How do you say "mother" in Spanish?\n\n### Assistant: madre'
train_data = {
"text": [str1, str2, str3, str4, str5],
}
dataset_text = Dataset.from_dict(train_data)

# to test if we learn how to generate an unknown word.
holdout_str = (
'\n\n### Human: How do you say "day" in Spanish?\n\n### Assistant:<s>' # día
)
device = "cuda" if torch.cuda.is_available() else "cpu"
holdout_input = tokenizer(holdout_str, return_tensors="pt").to(device)

Comencemos haciendo un preprocesamiento. Vamos a agregar algunos tokens especiales, a saber, “fin de secuencia” (eos) y “comienzo de secuencia” (bos). Estos tokens especiales pueden ser útiles para que el modelo sepa cuándo debe comenzar y dejar de generar texto.

INSTRUCTION_TEMPLATE_BASE = "\n\n### Human:"
RESPONSE_TEMPLATE_BASE = "\n\n### Assistant:"
def add_special_tokens(
example: Dict,
tokenizer: PreTrainedTokenizerBase,
) -> Dict:
# add eos_token before human text and bos_token before assistant text
example["text"] = (
example["text"]
.replace(
INSTRUCTION_TEMPLATE_BASE, tokenizer.eos_token + INSTRUCTION_TEMPLATE_BASE
)
.replace(RESPONSE_TEMPLATE_BASE, RESPONSE_TEMPLATE_BASE + tokenizer.bos_token)
)
if not example["text"].endswith(tokenizer.eos_token):
example["text"] += tokenizer.eos_token
# Remove leading EOS tokens
while example["text"].startswith(tokenizer.eos_token):
example["text"] = example["text"][len(tokenizer.eos_token) :]
return example

dataset_text = dataset_text.map(lambda x: add_special_tokens(x, tokenizer))
print(f"{dataset_text=}")
print(f"{dataset_text[0]=}")
>>> dataset_text=Dataset({
features: ['text'],
num_rows: 5
})
>>> dataset_text[0]={'text': '\n\n### Human: How do you say "dog" in Spanish?\n\n### Assistant:<s> perro</s>'}

Ahora vamos a hacer lo que aprendimos en la última sesión: crear una entrada con una clave de etiquetas copiada de input_ids.

# tokenize the text
dataset = dataset_text.map(
lambda example: tokenizer(example["text"]), batched=True, remove_columns=["text"]
)
# copy the input_ids to labels
dataset = dataset.map(lambda x: {"labels": x["input_ids"]}, batched=True)
print(f"{dataset=}")
print(f"{dataset[0]['input_ids']=}")
print(f"{dataset[0]['labels']=}")
>>> dataset=Dataset({
features: ['input_ids', 'attention_mask', 'labels'],
num_rows: 5
})
>>> dataset[0]['input_ids']=[603, 105311, 22256, 29, 7535, 727, 1152, 5894, 20587, 744, 5, 361, 49063, 7076, 105311, 143005, 29, 1, 82208, 2]
>>> dataset[0]['labels']=[603, 105311, 22256, 29, 7535, 727, 1152, 5894, 20587, 744, 5, 361, 49063, 7076, 105311, 143005, 29, 1, 82208, 2]

Para empezar, las etiquetas y input_ids son idénticos. Veamos qué sucede cuando entrenamos un modelo como ese.

# training code inspired by
#https://mlabonne.github.io/blog/posts/Fine_Tune_Your_Own_Llama_2_Model_in_a_Colab_Notebook.html
model = load_model(model_name)
output_dir = "./results"
# How many times to iterate over the entire dataset
num_train_epochs = 15
# We're not aligning the sequence length (ie padding or truncating)
# so batch training won't work for our toy example.
per_device_train_batch_size = 1

training_arguments = TrainingArguments(
output_dir=output_dir,
num_train_epochs=num_train_epochs,
per_device_train_batch_size=per_device_train_batch_size,
seed=1,
)
trainer = Trainer(
model=model,
train_dataset=dataset,
args=training_arguments,
)
training1 = trainer.train()

# Sample generate prediction on holdout set
“\n\n### Human: How do you say "good" in Spanish?\n\n### Assistant:”
# the correct output is “bueno</s>”

sample_generate(model, tokenizer, holdout_inputs, max_new_tokens=5)
>>> ‘</s>’

Después de 15 épocas, todavía estamos un poco confundidos. Generamos ‘‘, que está cerca, pero realmente queremos generar “perro“. Aprendamos otras 15 épocas.

trainer.train()
sample_generate(model, tokenizer, holdout_input, max_new_tokens=5)
>>> bueno </s>

¡Después de 30 épocas aprendimos lo que se suponía que debíamos aprender!

Simulemos lo que sucede en el entrenamiento prediciendo de forma iterativa el mensaje, un token a la vez, en función de los tokens anteriores.

print_iterative_generate(model, tokenizer, inputs)
>>>
#
: How do you say "how morning in Spanish?

### Assistant: gu buenopu

Eso se acerca bastante al mensaje real, como esperábamos. Pero la tarea es la traducción, por lo que realmente no nos importa poder predecir el mensaje del usuario. ¿Hay alguna manera de aprender solo la parte de respuesta?

Enfoque enmascarado

Hugging Face te permite aprender a predecir ciertos tokens solo “enmascarando” los tokens que no te interesan en “etiquetas”. Esto es diferente de la máscara de atención, que oculta anterior tokens que utilizamos para generar un nuevo token. Enmascarar las etiquetas oculta el token que se supone que debes generar en un índice determinado de la función de pérdida. Tenga en cuenta la redacción: Hugging Face lo ha implementado de manera que durante el entrenamiento, todavía generamos predicciones para ese token enmascarado. Sin embargo, debido a que ocultamos la etiqueta verdadera para comparar las predicciones, no aprendemos directamente cómo mejorar esa predicción.

Creamos la “máscara” volteando esos tokens a -100 en la clave de etiquetas.

def create_special_mask(example: Dict) -> Dict:
"""Mask human text and keep assistant text as it is.

Args:
example (Dict): Result of tokenizing some text

Returns:
Dict: The dict with the label masked
"""
# setting a token to -100 is how we "mask" a token
# and tell the model to ignore it when calculating the loss
mask_token_id = -100
# assume we always start with a human text
human_text = True
for idx, tok_id in enumerate(example["labels"]):
if human_text:
# mask all human text up until and including the bos token
example["labels"][idx] = mask_token_id
if tok_id == tokenizer.bos_token_id:
human_text = False
elif not human_text and tok_id == tokenizer.eos_token_id:
# don’t mask the eos token, but the next token will be human text to mask
human_text = True
elif not human_text:
# leave example['labels'] text as it is when assistant text
continue
return example

dataset_masked = dataset.map(create_special_mask)
# convert dataset from lists to torch tensors
dataset_masked.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
print(f"{dataset_masked[0]["labels"]=}")

>>> dataset[0]["labels"]=tensor([ -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 82208, 2])

model = load_model(model_name)
trainer = Trainer(
model=model,
train_dataset=dataset_masked,
args=training_arguments,
)

training2 = trainer.train()

print(f"{training2.metrics['train_runtime']=}")
print(f"{training1.metrics['train_runtime'] =}")
print(
f"{100*round((training1.metrics['train_runtime'] - training2.metrics['train_runtime']) / training1.metrics['train_runtime'] , 2)}%"
)

>>> training2.metrics['train_runtime']=61.7164
>>> training1.metrics['train_runtime'] =70.8013
>>> 13.0%

En primer lugar, esta vez fuimos más rápidos en más de un 10%. Presumiblemente, el hecho de que tengamos menos cálculos de pérdidas hace que las cosas sean un poco más rápidas.

No confiaría en que la aceleración sea tan grande: nuestro ejemplo está bastante desequilibrado con mucho más texto humano que texto generado. Pero cuando los tiempos de entrenamiento son horas, cada pequeño porcentaje es útil.

La gran pregunta: ¿aprendimos la tarea?

sample_generate(model, tokenizer, holdout_input, max_new_tokens=5)
>>> bueno </s>

Esta vez solo necesitamos 15 épocas para aprender la tarea. Volvamos a cómo son las cosas bajo el capó durante el entrenamiento.

print_iterative_generate(model, tokenizer, inputs)
>>>#include
code
to I get "we" in English?
A: Spanish: How bueno

La predicción iterativa del mensaje conduce a una tontería en comparación con nuestro primer enfoque de entrenamiento. Esto es cierto: enmascaramos el mensaje durante el entrenamiento y, por lo tanto, no aprendemos a predecir nada hasta nuestro objetivo real: la respuesta del asistente.

Usando el entrenador de ajuste fino supervisado de TRL

Hugging Face lanzó recientemente una biblioteca TRL (aprendizaje por refuerzo de transformadores) para agregar soporte de extremo a extremo para el proceso de capacitación LLM. Una característica es el ajuste fino supervisado. Usando las clases DataCollatorForCompletionOnlyLM y SFTTrainer, podemos crear las etiquetas como lo hicimos con crear_mascarilla_especial con solo unas pocas configuraciones.

model = load_model(model_name)

# a hugging face function to do the copying of labels for you.
# using the instruction and response templates will mask everything between the instruction template and the start of the response_template
collator = DataCollatorForCompletionOnlyLM(
instruction_template=tokenizer.eos_token,
response_template=tokenizer.bos_token,
tokenizer=tokenizer,
)

trainersft = SFTTrainer(
model,
train_dataset=dataset_text,
dataset_text_field="text",
data_collator=collator,
args=training_arguments,
tokenizer=tokenizer,
)
sftrain = trainersft.train()

sample_generate(model, tokenizer, holdout_input, max_new_tokens=5)
>>> ' perro</s>'

¡Éxito! Si profundizas más, el entrenamiento en realidad llevó más tiempo con SFT. Esto podría atribuirse al hecho de que tenemos que tokenizar en el momento del entrenamiento en lugar de como un paso de preprocesamiento en el enfoque enmascarado. Sin embargo, este enfoque nos brinda procesamiento por lotes gratuito (tendría que modificar el proceso de tokenización para usar el enfoque enmascarado para procesar por lotes correctamente), lo que debería acelerar las cosas a largo plazo.

El cuaderno completo explora algunas otras cosas, como entrenar en chats de varios turnos y usar tokens especiales para indicar texto humano versus texto de chat.

Obviamente, este ejemplo es un poco básico. Sin embargo, es de esperar que pueda comenzar a ver el poder de usar CausalLM: puede imaginar tomar interacciones de un modelo grande y confiable y usar las técnicas anteriores para ajustar un modelo más pequeño en los resultados del modelo grande. A esto se le llama destilación del conocimiento.

Si hemos aprendido algo en los últimos dos años de LLM, es que podemos hacer algunas cosas sorprendentemente inteligentes simplemente entrenándonos en la predicción del próximo token. Los modelos de lenguaje causal están diseñados para hacer precisamente eso. Incluso si la clase Hugging Face es un poco confusa al principio, una vez que te acostumbras, tienes una interfaz muy poderosa para entrenar tus propios modelos generativos.