Ajustar el Falcon 7B y otros LLMs en Amazon SageMaker con el decorador @remote

Ajustar Falcon 7B y otros LLMs en Amazon SageMaker con @remote

Hoy en día, los modelos de IA generativa abarcan una variedad de tareas, desde resúmenes de texto, preguntas y respuestas, hasta generación de imágenes y videos. Para mejorar la calidad de la salida, se utilizan enfoques como el aprendizaje n-shorts, la ingeniería de indicaciones, la Generación Mejorada por Búsqueda (RAG) y el ajuste fino. El ajuste fino te permite ajustar estos modelos de IA generativa para lograr un mejor rendimiento en tus tareas específicas del dominio.

Con Amazon SageMaker, ahora puedes ejecutar un trabajo de entrenamiento de SageMaker simplemente anotando tu código Python con el decorador @remote. El SDK de Python de SageMaker traduce automáticamente tu entorno de trabajo existente y cualquier código de procesamiento de datos y conjuntos de datos asociados en un trabajo de entrenamiento de SageMaker que se ejecuta en la plataforma de entrenamiento. Esto tiene la ventaja de escribir el código de una forma más natural y orientada a objetos, y aún así utiliza las capacidades de SageMaker para ejecutar trabajos de entrenamiento en un clúster remoto con cambios mínimos.

En esta publicación, mostramos cómo ajustar finamente los modelos de Falcon-7B Foundation Models (FM) usando el decorador @remote del SDK de Python de SageMaker. También utiliza la biblioteca de ajuste fino de Hugging Face y técnicas de cuantización a través de bitsandbytes para admitir el ajuste fino. El código presentado en este blog también se puede utilizar para ajustar finamente otros FMs, como Llama-2 13b.

Las representaciones de precisión completa de este modelo pueden tener desafíos para caber en la memoria de una o incluso varias Unidades de Procesamiento Gráfico (GPUs) —o incluso pueden necesitar una instancia más grande. Por lo tanto, para ajustar finamente este modelo sin aumentar el costo, utilizamos la técnica conocida como Adaptadores de Baja Rango (QLoRA) de LLMs Cuantizados. QLoRA es un enfoque eficiente de ajuste fino que reduce el uso de memoria de LLMs manteniendo un rendimiento muy bueno.

Ventajas de usar el decorador @remote

Antes de continuar, comprendamos cómo el decorador remoto mejora la productividad del desarrollador al trabajar con SageMaker:

• Crear un pipeline de clasificación con la clasificación personaliza...
• Estados Unidos desafía la dominancia de búsqueda de Google
• El Capital Privado recluta talento en Ciencia de Datos mientras la ...

• El decorador @remote desencadena directamente un trabajo de entrenamiento utilizando código Python nativo, sin la invocación explícita de Estimadores de SageMaker y canales de entrada de SageMaker
• Barrera baja para que los desarrolladores entrenen modelos en SageMaker.
• No es necesario cambiar de entornos de desarrollo integrados (IDEs). Continúa escribiendo código en tu IDE preferido e invoca trabajos de entrenamiento de SageMaker.
• No es necesario aprender sobre contenedores. Continúa proporcionando dependencias en un requirements.txt y suminístralo al decorador remoto.

Prerrequisitos

Se necesita una cuenta de AWS con un rol de AWS Identity and Access Management (IAM) que tenga permisos para administrar los recursos creados como parte de la solución. Para más detalles, consulta la creación de una cuenta de AWS.

En esta publicación, utilizamos Amazon SageMaker Studio con la imagen Data Science 3.0 y una instancia de lanzamiento rápido ml.t3.medium. Sin embargo, puedes utilizar cualquier entorno de desarrollo integrado (IDE) de tu elección. Solo necesitas configurar correctamente las credenciales de la AWS Command Line Interface (AWS CLI). Para más información, consulta la configuración de la AWS CLI.

Para el ajuste fino, se utiliza una instancia ml.g5.12xlarge en esta publicación. Asegúrate de tener suficiente capacidad para esta instancia en tu cuenta de AWS.

Necesitarás clonar este repositorio de Github para replicar la solución demostrada en esta publicación.

Descripción general de la solución

1. Instalar los prerequisitos para ajustar finamente el modelo Falcon-7B
2. Configurar las configuraciones del decorador remoto
3. Procesar el conjunto de datos que contiene las preguntas frecuentes de los servicios de AWS
4. Ajustar finamente el Falcon-7B en las preguntas frecuentes de los servicios de AWS
5. Probar los modelos ajustados finamente en preguntas de muestra relacionadas con los servicios de AWS

1. Instalar los prerequisitos para ajustar finamente el modelo Falcon-7B

Ejecuta el notebook falcon-7b-qlora-remote-decorator_qa.ipynb en SageMaker Studio seleccionando la imagen Data Science y el Kernel como Python 3. Instala todas las bibliotecas requeridas mencionadas en el requirements.txt. Algunas de las bibliotecas deben instalarse en la propia instancia del notebook. Realiza otras operaciones necesarias para el procesamiento del conjunto de datos y para desencadenar un trabajo de entrenamiento de SageMaker.

%pip install -r requirements.txt

%pip install -q -U transformers==4.31.0
%pip install -q -U datasets==2.13.1
%pip install -q -U peft==0.4.0
%pip install -q -U accelerate==0.21.0
%pip install -q -U bitsandbytes==0.40.2
%pip install -q -U boto3
%pip install -q -U sagemaker==2.154.0
%pip install -q -U scikit-learn

2. Configuración de los decoradores remotos

Cree un archivo de configuración donde se especifican todas las configuraciones relacionadas con el trabajo de entrenamiento de Amazon SageMaker. Este archivo es leído por el decorador @remote mientras se ejecuta el trabajo de entrenamiento. Este archivo contiene configuraciones como dependencias, imagen de entrenamiento, instancia y rol de ejecución a utilizar para el trabajo de entrenamiento. Para obtener una referencia detallada de todas las configuraciones admitidas por el archivo de configuración, consulte la guía Configuración y uso de valores predeterminados con el SDK de Python de SageMaker.

SchemaVersion: '1.0'
SageMaker:
  PythonSDK:
    Modules:
      RemoteFunction:
        Dependencies: ./requirements.txt
        ImageUri: '{aws_account_id}.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04'
        InstanceType: ml.g5.12xlarge
        RoleArn: arn:aws:iam::111122223333:role/ExampleSageMakerRole

No es obligatorio utilizar el archivo config.yaml para trabajar con el decorador @remote. Esta es solo una forma más limpia de proporcionar todas las configuraciones al decorador @remote. Esto mantiene los parámetros relacionados con SageMaker y AWS fuera del código con un esfuerzo único para configurar el archivo de configuración utilizado por los miembros del equipo. Todas las configuraciones también se pueden proporcionar directamente en los argumentos del decorador, pero esto reduce la legibilidad y la capacidad de mantener los cambios a largo plazo. Además, el archivo de configuración puede ser creado por un administrador y compartido con todos los usuarios en un entorno.

Preprocesar el conjunto de datos que contiene preguntas frecuentes sobre los servicios de AWS

El siguiente paso es cargar y preprocesar el conjunto de datos para prepararlo para el trabajo de entrenamiento. Primero, echemos un vistazo al conjunto de datos:

Muestra las preguntas frecuentes sobre uno de los servicios de AWS. Además de QLoRA, se utiliza bitsanbytes para convertir a precisión de 4 bits y cuantizar LLM congelado a 4 bits y adjuntar adaptadores LoRA en él.

Cree una plantilla de instrucciones para convertir cada muestra de pregunta frecuente a un formato de instrucción:

from random import randint

# Inicio personalizado de la plantilla de instrucciones
prompt_template = f"{{question}}\n---\nRespuesta:\n{{answer}}{{eos_token}}"

# Conjunto de datos de plantilla para agregar instrucción a cada muestra
def template_dataset(sample):
    sample["text"] = prompt_template.format(question=sample["question"],
                                            answer=sample["answers"],
                                            eos_token=tokenizer.eos_token)
    return sample

El siguiente paso es convertir las entradas (texto) a identificadores de tokens. Esto se hace mediante un tokenizador de Hugging Face Transformers.

from transformers import AutoTokenizer

model_id = "tiiuae/falcon-7b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
# Establecer el tokenizador Falcon
tokenizer.pad_token = tokenizer.eos_token

Ahora simplemente use la función prompt_template para convertir todas las preguntas frecuentes al formato de instrucción y configure los conjuntos de datos de entrenamiento y prueba.

4. Ajuste fino de Falcon-7B en las preguntas frecuentes de los servicios de AWS

Ahora puede preparar el script de entrenamiento y definir la función de entrenamiento train_fn y colocar el decorador @remote en la función.

La función de entrenamiento realiza lo siguiente:

• tokeniza y divide el conjunto de datos en fragmentos
• configura BitsAndBytesConfig, que especifica que el modelo debe cargarse en 4 bits, pero durante el cálculo debe convertirse a bfloat16.
• Carga el modelo
• Encuentra los módulos objetivo y actualiza las matrices necesarias utilizando el método de utilidad find_all_linear_names
• Crea configuraciones de LoRA que especifican la clasificación de las matrices de actualización (s), el factor de escala (lora_alpha), los módulos para aplicar las matrices de actualización de LoRA (target_modules), la probabilidad de abandono para las capas de Lora (lora_dropout), task_type, etc.
• Inicia el entrenamiento y la evaluación

import bitsandbytes as bnb

def encontrar_todos_los_nombres_lineales(hf_model):
    nombres_modulo_lora = set()
    for nombre, modulo in hf_model.named_modules():
        if isinstance(modulo, bnb.nn.Linear4bit):
            nombres = nombre.split(".")
            nombres_modulo_lora.add(nombres[0] if len(nombres) == 1 else nombres[-1])

    if "lm_head" in nombres_modulo_lora:
        nombres_modulo_lora.remove("lm_head")
    return list(nombres_modulo_lora)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from sagemaker.remote_function import remote
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import transformers

# Iniciar entrenamiento
@remote(volume_size=50)
def train_fn(
        model_name,
        train_ds,
        test_ds,
        lora_r=8,
        lora_alpha=32,
        lora_dropout=0.05,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        learning_rate=2e-4,
        num_train_epochs=1
):
    # tokenizar y dividir el conjunto de datos
    lm_train_dataset = train_ds.map(
        lambda muestra: tokenizer(muestra["text"]), batched=True, batch_size=24, remove_columns=list(train_dataset.features)
    )


    lm_test_dataset = test_ds.map(
        lambda muestra: tokenizer(muestra["text"]), batched=True, remove_columns=list(test_dataset.features)
    )

    # Imprimir el número total de muestras
    print(f"Número total de muestras de entrenamiento: {len(lm_train_dataset)}")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    # Falcon requiere permitir la ejecución remota de código. Esto se debe a que el modelo utiliza una nueva arquitectura que aún no forma parte de transformers.
    # El código es proporcionado por los autores del modelo en el repositorio.
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        quantization_config=bnb_config,
        device_map="auto")

    model.gradient_checkpointing_enable()
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)

    # obtener módulos objetivo de lora
    modulos = encontrar_todos_los_nombres_lineales(model)
    print(f"Encontrados {len(modulos)} módulos para cuantizar: {modulos}")

    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=modulos,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM"
    )

    model = get_peft_model(model, config)
    print_trainable_parameters(model)

    trainer = transformers.Trainer(
        model=model,
        train_dataset=lm_train_dataset,
        eval_dataset=lm_test_dataset,
        args=transformers.TrainingArguments(
            per_device_train_batch_size=per_device_train_batch_size,
            per_device_eval_batch_size=per_device_eval_batch_size,
            logging_steps=2,
            num_train_epochs=num_train_epochs,
            learning_rate=learning_rate,
            bf16=True,
            save_strategy="no",
            output_dir="outputs"
        ),
        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )
    model.config.use_cache = False

    trainer.train()
    trainer.evaluate()

    model.save_pretrained("/opt/ml/model")

E invoca la función train_fn()

train_fn(model_id, train_dataset, test_dataset)

El trabajo de ajuste se ejecutará en el clúster de entrenamiento de Amazon SageMaker. Espera a que el trabajo de ajuste termine.

5. Prueba los modelos de ajuste fino en preguntas de ejemplo relacionadas con los servicios de AWS

Ahora, es hora de ejecutar algunas pruebas en el modelo. Primero, carguemos el modelo:

from peft import PeftModel, PeftConfig
import torch
from transformers import AutoModelForCausalLM

device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'

config = PeftConfig.from_pretrained("./model")
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)
model = PeftModel.from_pretrained(model, "./model")
model.to(device)

Ahora carga una pregunta de ejemplo del conjunto de datos de entrenamiento para ver la respuesta original y luego haz la misma pregunta al modelo ajustado para ver la respuesta en comparación.