Ajuste Falcon 7B y otros LLM en Amazon SageMaker con @remote decorador

Hoy en día, los modelos de IA generativa cubren una variedad de tareas, desde resúmenes de texto, preguntas y respuestas y generación de imágenes y videos. Para mejorar la calidad del resultado, se utilizan enfoques como el aprendizaje n-corto, la ingeniería rápida, Recuperación de Generación Aumentada (RAG) y se utilizan ajustes finos. El ajuste fino le permite ajustar estos modelos de IA generativa para lograr un mejor rendimiento en las tareas específicas de su dominio.

Con Amazon SageMakerahora puede ejecutar un trabajo de capacitación de SageMaker simplemente anotando su código Python con @decorador remoto. El SDK de Python de SageMaker traduce automáticamente su entorno de espacio de trabajo existente, y cualquier código y conjunto de datos de procesamiento de datos asociados, en un trabajo de capacitación de SageMaker que se ejecuta en la plataforma de capacitación. Esto tiene la ventaja de escribir el código de una manera más natural, orientada a objetos, y aún 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 un Halcón-7B Modelos de base (FM) utilizando el decorador @remote de SageMaker Python SDK. También utiliza Ajuste fino eficiente en parámetros de Hugging Face (PEFT) y técnicas de cuantificación a través de bits y bytes para apoyar el ajuste. El código presentado en este blog también se puede utilizar para ajustar otros FM, como Llama-2 13b.

Las representaciones de precisión total de este modelo pueden tener dificultades para caber en la memoria en uno o incluso varios. Unidades de procesamiento gráfico (GPU) – o incluso puede necesitar una instancia más grande. Por lo tanto, para perfeccionar este modelo sin aumentar el costo, utilizamos la técnica conocida como LLM cuantificados con adaptadores de bajo rango (QLoRA). QLoRA es un enfoque de ajuste eficiente que reduce el uso de memoria de los LLM y al mismo tiempo mantiene un rendimiento muy bueno.

Ventajas de usar @remote decorador

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

• @remote decorador activa un trabajo de entrenamiento directamente usando código Python nativo, sin la invocación explícita de los estimadores de SageMaker y los canales de entrada de SageMaker.
• Barrera de entrada baja para desarrolladores de modelos de formación en SageMaker.
• No es necesario cambiar Entornos de desarrollo integrados (IDE). Continúe escribiendo código en el IDE que elija e invoque los trabajos de capacitación de SageMaker.
• No es necesario aprender sobre contenedores. Continuar brindando dependencias en un requirements.txt y entregárselo al decorador remoto.

Requisitos previos

Se necesita una cuenta de AWS con una Gestión de acceso e identidad de AWS (AWS IAM) role que tiene permisos para administrar los recursos creados como parte de la solución. Para más detalles, consulte Crear una cuenta de AWS.

En esta publicación utilizamos Amazon SageMaker Estudio con el Data Science 3.0 imagen y un ml.t3.medium instancia de lanzamiento rápido. Sin embargo, puede utilizar cualquier entorno de desarrollo integrado (IDE) de su elección. Sólo necesitas configurar tu Interfaz de línea de comandos de AWS (AWS CLI) correctamente. Para obtener más información, consulte Configurar la CLI de AWS.

Para realizar ajustes, el Falcon-7B, un ml.g5.12xlarge La instancia se utiliza en esta publicación. Asegúrese de tener suficiente capacidad para esta instancia en la cuenta de AWS.

Necesitas clonar esto repositorio de github por replicar la solución demostrada en esta publicación.

Descripción general de la solución

1. Instale requisitos previos para ajustar el modelo Falcon-7B
2. Configurar configuraciones de decorador remoto
3. Preprocesar el conjunto de datos que contiene las preguntas frecuentes sobre los servicios de AWS
4. Preguntas frecuentes sobre cómo ajustar Falcon-7B en los servicios de AWS
5. Pruebe los modelos de ajuste fino en preguntas de muestra relacionadas con los servicios de AWS

1. Instale los requisitos previos para ajustar el modelo Falcon-7B

Inicie el cuaderno falcon-7b-qlora-decorador-remoto_qa.ipynb en SageMaker Studio seleccionando Imago como Data Science y Núcleo como Python 3. Instale todas las bibliotecas requeridas mencionadas en el requirements.txt. Pocas bibliotecas deben instalarse en la propia instancia del portátil. Realice otras operaciones necesarias para procesar conjuntos de datos y activar un trabajo de capacitación 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. Configurar configuraciones de decorador remoto

Cree un archivo de configuración donde se especifiquen todas las configuraciones relacionadas con el trabajo de capacitación de Amazon SageMaker. El decorador @remote lee este archivo mientras ejecuta el trabajo de entrenamiento. Este archivo contiene configuraciones como dependencias, imagen de entrenamiento, instancia y la función de ejecución que se utilizará para el trabajo de entrenamiento. Para obtener una referencia detallada de todas las configuraciones admitidas por el archivo de configuración, consulte Configuración y uso de valores predeterminados con SageMaker Python SDK.

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 config.yaml archivo 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 todos los miembros del equipo. Todas las configuraciones también podrían proporcionarse directamente en los argumentos del decorador, pero eso reduce la legibilidad y la mantenibilidad de los cambios a largo plazo. Además, un administrador puede crear el archivo de configuración y compartirlo con todos los usuarios de un entorno.

Preprocesar el conjunto de datos que contiene las 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 preguntas frecuentes sobre uno de los servicios de AWS. Además de QLoRA, bitsanbytes se utiliza para convertir a precisión de 4 bits para cuantificar el LLM congelado a 4 bits y adjuntar lora adaptadores en él.

Cree una plantilla de mensaje para convertir cada muestra de preguntas frecuentes a un formato de mensaje:

from random import randint

# custom instruct prompt start
prompt_template = f"{{question}}\n---\nAnswer:\n{{answer}}{{eos_token}}"

# template dataset to add prompt to each sample
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) en ID de token. Esto lo hace un Tokenizador de transformadores de cara abrazada.

from transformers import AutoTokenizer

model_id = "tiiuae/falcon-7b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
# Set the Falcon tokenizer
tokenizer.pad_token = tokenizer.eos_token

Ahora simplemente usa el prompt_template función para convertir todas las preguntas frecuentes al formato de solicitud y configurar conjuntos de datos de prueba y entrenamiento.

4. Preguntas frecuentes sobre cómo ajustar Falcon-7B en los servicios de AWS

Ahora puedes preparar el guión de entrenamiento y definir la función de entrenamiento. train_fn y ponga @remote decorador en la función.

La función de entrenamiento hace lo siguiente:

• tokeniza y fragmenta el conjunto de datos
• configuración BitsAndBytesConfigque especifica que el modelo debe cargarse en 4 bits pero mientras que el cálculo debe convertirse a bfloat16.
• Cargar el modelo
• Encuentre los módulos de destino y actualice las matrices necesarias utilizando el método de utilidad find_all_linear_names
• Crear lora configuraciones que especifican la clasificación de matrices de actualización (s), factor de escala (lora_alpha), los módulos para aplicar las matrices de actualización LoRA (target_modules), probabilidad de abandono para las capas de Lora (lora_dropout), task_typeetc.
• Iniciar la formación y evaluación.

import bitsandbytes as bnb

def find_all_linear_names(hf_model):
    lora_module_names = set()
    for name, module in hf_model.named_modules():
        if isinstance(module, bnb.nn.Linear4bit):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if "lm_head" in lora_module_names:
        lora_module_names.remove("lm_head")
    return list(lora_module_names)
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

# Start training
@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
):
    # tokenize and chunk dataset
    lm_train_dataset = train_ds.map(
        lambda sample: tokenizer(sample["text"]), batched=True, batch_size=24, remove_columns=list(train_dataset.features)
    )


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

    # Print total number of samples
    print(f"Total number of train samples: {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 requires you to allow remote code execution. This is because the model uses a new architecture that is not part of transformers yet.
    # The code is provided by the model authors in the repo.
    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)

    # get lora target modules
    modules = find_all_linear_names(model)
    print(f"Found {len(modules)} modules to quantize: {modules}")

    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=modules,
        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 invocar el train_fn()

train_fn(model_id, train_dataset, test_dataset)

El trabajo de ajuste se ejecutaría en el clúster de capacitación de Amazon SageMaker. Espere a que finalice el trabajo de ajuste.

5. Pruebe los modelos de ajuste fino en preguntas de muestra relacionadas con los servicios de AWS

Ahora es el momento de realizar 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 cargue una pregunta de muestra del conjunto de datos de entrenamiento para ver la respuesta original y luego haga la misma pregunta desde el modelo optimizado para ver la respuesta en comparación.

Aquí hay un ejemplo de una pregunta del conjunto de capacitación y la respuesta original:

Ahora, se hace la misma pregunta sobre el modelo Falcon-7B sintonizado:

Con esto concluye la implementación del ajuste fino de Falcon-7B en el conjunto de datos de preguntas frecuentes de los servicios de AWS utilizando el decorador @remote del SDK de Python de Amazon SageMaker.

Limpiar

Complete los siguientes pasos para limpiar sus recursos:

• Cierre las instancias de Amazon SageMaker Studio para evitar incurrir en costos adicionales.
• Limpie su sistema de archivos elástico de Amazon (AmazonEFS) borrando el directorio de caché de Hugging Face:

  rm -R ~/.cache/huggingface/hub

Conclusión

En esta publicación, le mostramos cómo usar efectivamente las capacidades del decorador @remote para ajustar el modelo Falcon-7B usando QLoRA, Hugging Face PEFT con bitsandbtyes sin aplicar cambios significativos en el cuaderno de capacitación y utilizó las capacidades de Amazon SageMaker para ejecutar trabajos de capacitación en un clúster remoto.

Todo el código que se muestra como parte de esta publicación para ajustar Falcon-7B está disponible en el repositorio de GitHub. El repositorio también contiene un cuaderno que muestra cómo ajustar Llama-13B.

Como siguiente paso, le recomendamos que consulte el @funcionalidad decoradora remota y API del SDK de Python y utilícelo en el entorno e IDE que elija. Hay ejemplos adicionales disponibles en el ejemplos-de-amazon-sagemaker repositorio para comenzar rápidamente. También puedes consultar las siguientes publicaciones:

Sobre los autores

Bruno Pistone es un arquitecto de soluciones especializado en IA/ML para AWS con sede en Milán. Trabaja con grandes clientes ayudándoles a comprender en profundidad sus necesidades técnicas y a diseñar soluciones de inteligencia artificial y aprendizaje automático que aprovechen al máximo la nube de AWS y la pila de aprendizaje automático de Amazon. Su experiencia incluye: aprendizaje automático de extremo a extremo, industrialización del aprendizaje automático e inteligencia artificial generativa. Le gusta pasar tiempo con sus amigos y explorar nuevos lugares, además de viajar a nuevos destinos.

vikesh pandey es un arquitecto de soluciones especializado en aprendizaje automático en AWS y ayuda a clientes de industrias financieras a diseñar y crear soluciones en IA generativa y aprendizaje automático. Fuera del trabajo, a Vikesh le gusta probar diferentes cocinas y practicar deportes al aire libre.