Vertex AI: Usa el empaquetado automático para ajustar BERT con Hugging Face en Vertex AI Training

1. Descripción general

En este lab, aprenderás a ejecutar un trabajo de entrenamiento personalizado en Vertex AI Training con la función de empaquetado automático. Los trabajos de entrenamiento personalizados en Vertex AI usan contenedores. Si no quieres compilar tu propia imagen, puedes usar auotpackaging, el cual compilará una imagen personalizada de Docker basada en tu código, la enviará a Container Registry y, luego, iniciará un CustomJob basado en la imagen.

Qué aprenderá

Aprenderás a hacer lo siguiente:

• Usa el modo local para probar tu código.
• Configura y, luego, inicia un trabajo de entrenamiento personalizado con el empaquetado automático.

El costo total de la ejecución de este lab en Google Cloud es de aproximadamente $2.

2. Descripción general del caso de uso

Con las bibliotecas de Hugging Face, ajustarás un modelo de Bert en el conjunto de datos de IMDB. El modelo predecirá si la opinión sobre una película es positiva o negativa. El conjunto de datos se descargará de la biblioteca de conjuntos de datos de Hugging Face y el modelo de Bert de la biblioteca de transformadores de Hugging Face.

3. Introducción a Vertex AI

En este lab, se utiliza la oferta de productos de IA más reciente de Google Cloud. Vertex AI integra las ofertas de AA de Google Cloud en una experiencia de desarrollo fluida. Anteriormente, se podía acceder a los modelos personalizados y a los entrenados con AutoML mediante servicios independientes. La nueva oferta combina ambos en una sola API, junto con otros productos nuevos. También puedes migrar proyectos existentes a Vertex AI. Para enviarnos comentarios, visita la página de asistencia.

Vertex AI incluye muchos productos distintos para respaldar flujos de trabajo de AA de extremo a extremo. Este lab se enfocará en Training y Workbench.

4. Configura el entorno

Para ejecutar este codelab, necesitarás un proyecto de Google Cloud Platform que tenga habilitada la facturación. Para crear un proyecto, sigue estas instrucciones.

Paso 1: Habilita la API de Compute Engine

Ve a Compute Engine y selecciona Habilitar (si aún no está habilitada).

Paso 2: Habilita la API de Vertex AI

Navegue hasta la sección de Vertex AI en la consola de Cloud y haga clic en Habilitar API de Vertex AI.

Paso 3: Habilita la API de Container Registry

Navega a Container Registry y selecciona Habilitar si aún no lo has hecho. La usarás para crear un contenedor para tu trabajo de entrenamiento personalizado.

Paso 4: Crea una instancia de Vertex AI Workbench

En la sección Vertex AI de Cloud Console, haz clic en Workbench:

Allí, haz clic en NOTEBOOKS ADMINISTRADOS (MANAGED NOTEBOOKS):

Luego, selecciona NUEVO NOTEBOOK (NEW NOTEBOOK).

Asígnale un nombre al notebook y, luego, haz clic en Configuración avanzada (Advanced settings).

En Configuración avanzada (Advanced settings), activa la opción Habilitar el cierre inactivo (Enable Idle Shutdown) y establece la cantidad de minutos en 60. Esto provocará que el notebook se cierre automáticamente cuando no esté en uso para que no se generen costos innecesarios.

Puedes dejar el resto de la configuración avanzada tal como está.

Luego, haz clic en Crear.

Una vez que se crea la instancia, selecciona Abrir JupyterLab (Open JupyterLab).

La primera vez que uses una instancia nueva, se te solicitará que te autentiques.

5. Escribe el código de entrenamiento

Para comenzar, en el menú Launcher, abre una ventana de terminal en tu instancia de notebook:

Crea un directorio nuevo llamado autopkg-codelab y ábrelo con el comando cd.

mkdir autopkg-codelab
cd autopkg-codelab

En la terminal, ejecuta el siguiente comando para crear un directorio para el código de entrenamiento y un archivo de Python en el que agregarás el código.

mkdir trainer
touch trainer/task.py

Ahora, deberías tener lo siguiente en el directorio autopkg-codelab/:

+ trainer/
    + task.py

A continuación, abre el archivo task.py que acabas de crear y copia el código que aparece más abajo.

import argparse

import tensorflow as tf
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import TFAutoModelForSequenceClassification

CHECKPOINT = "bert-base-cased"

def get_args():
  '''Parses args.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--epochs',
      required=False,
      default=3,
      type=int,
      help='number of epochs')
  parser.add_argument(
      '--job_dir',
      required=True,
      type=str,
      help='bucket to store saved model, include gs://')
  args = parser.parse_args()
  return args


def create_datasets():
    '''Creates a tf.data.Dataset for train and evaluation.'''

    raw_datasets = load_dataset('imdb')
    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
    tokenized_datasets = raw_datasets.map((lambda examples: tokenize_function(examples, tokenizer)), batched=True)

    # To speed up training, we use only a portion of the data.
    # Use full_train_dataset and full_eval_dataset if you want to train on all the data.
    small_train_dataset = tokenized_datasets['train'].shuffle(seed=42).select(range(1000))
    small_eval_dataset = tokenized_datasets['test'].shuffle(seed=42).select(range(1000))
    full_train_dataset = tokenized_datasets['train']
    full_eval_dataset = tokenized_datasets['test']

    tf_train_dataset = small_train_dataset.remove_columns(['text']).with_format("tensorflow")
    tf_eval_dataset = small_eval_dataset.remove_columns(['text']).with_format("tensorflow")

    train_features = {x: tf_train_dataset[x] for x in tokenizer.model_input_names}
    train_tf_dataset = tf.data.Dataset.from_tensor_slices((train_features, tf_train_dataset["label"]))
    train_tf_dataset = train_tf_dataset.shuffle(len(tf_train_dataset)).batch(8)

    eval_features = {x: tf_eval_dataset[x] for x in tokenizer.model_input_names}
    eval_tf_dataset = tf.data.Dataset.from_tensor_slices((eval_features, tf_eval_dataset["label"]))
    eval_tf_dataset = eval_tf_dataset.batch(8)

    return train_tf_dataset, eval_tf_dataset


def tokenize_function(examples, tokenizer):
    '''Tokenizes text examples.'''

    return tokenizer(examples['text'], padding='max_length', truncation=True)


def main():
    args = get_args()
    train_tf_dataset, eval_tf_dataset = create_datasets()
    model = TFAutoModelForSequenceClassification.from_pretrained(CHECKPOINT, num_labels=2)

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=tf.metrics.SparseCategoricalAccuracy(),
    )

    model.fit(train_tf_dataset, validation_data=eval_tf_dataset, epochs=args.epochs)
    model.save(f'{args.job_dir}/model_output')


if __name__ == "__main__":
    main()

Ten en cuenta lo siguiente sobre el código:

• CHECKPOINT es el modelo que queremos ajustar. En este caso, usamos Bert.
• El método TFAutoModelForSequenceClassification cargará la arquitectura del modelo de lenguaje especificado y los pesos en TensorFlow, y agregará un encabezado de clasificación en la parte superior con pesos inicializados de forma aleatoria. En este caso, tenemos un problema de clasificación binaria (positivo o negativo), por lo que especificamos num_labels=2 para este clasificador.

6. Crear contenedores y ejecutar código de entrenamiento de manera local

Puedes usar el comando gcloud ai custom-jobs local-run para compilar una imagen de contenedor de Docker basada en tu código de entrenamiento y ejecutarla como un contenedor en tu máquina local. La ejecución local de un contenedor ejecuta tu código de entrenamiento de manera similar a la que se ejecuta en el entrenamiento de Vertex AI y puede ayudarte a depurar problemas antes de realizar un entrenamiento personalizado en Vertex AI.

En nuestro trabajo de entrenamiento, exportaremos el modelo entrenado a un bucket de Cloud Storage. Desde tu terminal, ejecuta lo siguiente para definir una variable de entorno para tu proyecto y asegúrate de reemplazar your-cloud-project con el ID de tu proyecto:

Para obtener el ID del proyecto, ejecuta gcloud config list --format 'value(core.project)' en tu terminal.

PROJECT_ID='your-cloud-project'

Luego, crea un bucket. Si ya tienes un bucket, puedes usarlo en su lugar.

BUCKET_NAME="gs://${PROJECT_ID}-bucket"
gsutil mb -l us-central1 $BUCKET_NAME

Cuando ejecutemos el trabajo de entrenamiento personalizado en el entrenamiento de Vertex AI, usaremos una GPU. Sin embargo, como no especificamos nuestra instancia de Workbench con GPUs, usaremos una imagen basada en CPU para las pruebas locales. En este ejemplo, usamos un contenedor creado previamente de Vertex AI Training.

Ejecuta el siguiente comando para establecer el URI de una imagen de Docker que se usará como base del contenedor.

BASE_CPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-7:latest

Luego, establece un nombre para la imagen resultante de Docker que compiló el comando de ejecución local.

OUTPUT_IMAGE=$PROJECT_ID-local-package-cpu:latest

Nuestro código de entrenamiento usa los conjuntos de datos y las bibliotecas de transformadores de Hugging Face. Estas bibliotecas no se incluyen en la imagen que seleccionamos como imagen base, por lo que deberemos proporcionarlas como requisitos. Para ello, crearemos un archivo requirements.txt en nuestro directorio autopkg-codelab.

Asegúrate de estar en el directorio autopkg-codelab y escribe lo siguiente en la terminal.

touch requirements.txt

Ahora, deberías tener lo siguiente en el directorio autopkg-codelab:

+ requirements.txt
+ trainer/
    + task.py

Abre el archivo de requisitos y pega lo siguiente:

datasets==1.18.2
transformers==4.16.2

Por último, ejecuta el comando gcloud ai custom-jobs local-run para iniciar el entrenamiento en nuestra instancia administrada de Workbench.

gcloud ai custom-jobs local-run \
--executor-image-uri=$BASE_CPU_IMAGE \
--python-module=trainer.task \
--output-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME

Deberías ver que se compila la imagen de Docker. Las dependencias que agregamos al archivo requirements.txt se instalarán con pip. La primera vez que ejecutes este comando, es posible que tarde unos minutos en completarse. Una vez que se compile la imagen, comenzará a ejecutarse el archivo task.py y verás el entrenamiento del modelo. Debería ver algo como esto:

Debido a que no estamos usando una GPU de forma local, el entrenamiento del modelo llevará mucho tiempo. Puedes presionar Ctrl + C y cancelar el entrenamiento local en lugar de esperar a que se complete el trabajo.

Ten en cuenta que, si quieres realizar más pruebas, también puedes ejecutar directamente la imagen compilada anteriormente, sin volver a empaquetarla.

gcloud beta ai custom-jobs local-run \
--executor-image-uri=$OUTPUT_IMAGE \
-- \
--job_dir=$BUCKET_NAME \
--epochs=1

7. Crea un trabajo personalizado

Ahora que probamos el modo local, usaremos la función de empaquetado automático para iniciar nuestro trabajo de entrenamiento personalizado en el entrenamiento de Vertex AI. Con un solo comando, esta función hará lo siguiente:

• Compilar una imagen de Docker personalizada basada en tu código
• Envía la imagen a Container Registry.
• Iniciar una CustomJob basada en la imagen

Regresa a la terminal y usa el comando cd para subir un nivel por encima del directorio autopkg-codelab.

+ autopkg-codelab
  + requirements.txt
  + trainer/
      + task.py

Especifica la imagen de GPU de TensorFlow previamente compilada de Vertex AI Training como la imagen base del trabajo de entrenamiento personalizado.

BASE_GPU_IMAGE=us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-7:latest

A continuación, ejecuta el comando gcloud ai custom-jobs create. Primero, este comando compilará una imagen de Docker personalizada basada en el código de entrenamiento. La imagen base es el contenedor precompilado de Vertex AI Training que establecimos como BASE_GPU_IMAGE. Luego, la función de empaquetado automático instalará los conjuntos de datos y las bibliotecas de transformadores con pip, como se especifica en nuestro archivo requirements.txt.

gcloud ai custom-jobs create \
--region=us-central1 \
--display-name=fine_tune_bert \
--args=--job_dir=$BUCKET_NAME \
--worker-pool-spec=machine-type=n1-standard-4,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,executor-image-uri=$BASE_GPU_IMAGE,local-package-path=autopkg-codelab,python-module=trainer.task

Veamos el argumento worker-pool-spec. Esto define la configuración del grupo de trabajadores que usa el trabajo personalizado. Puedes especificar varias especificaciones de grupo de trabajadores a fin de crear un trabajo personalizado con varios grupos de trabajadores para el entrenamiento distribuido. En este ejemplo, solo especificamos un grupo de trabajadores, ya que nuestro código de entrenamiento no está configurado para el entrenamiento distribuido.

Estos son algunos de los campos clave de esta especificación:

• machine-type (obligatorio): Es el tipo de máquina. Haz clic aquí para conocer los tipos admitidos.
• replica-count: Es la cantidad de réplicas de trabajadores que se deben usar para este grupo de trabajadores. De forma predeterminada, el valor es 1.
• accelerator-type: Es el tipo de GPU. Haz clic aquí para conocer los tipos admitidos. En este ejemplo, especificamos una GPU NVIDIA Tesla V100.
• accelerator-count: Es la cantidad de GPUs que debe usar cada VM del grupo de trabajo. De forma predeterminada, el valor es 1.
• executor-image-uri: Es el URI de una imagen de contenedor que ejecutará el paquete proporcionado. Esto se establece en nuestra imagen base.
• local-package-path: Es la ruta de acceso local de una carpeta que contiene el código de entrenamiento.
• python-module: Es el nombre del módulo de Python que se ejecutará dentro del paquete proporcionado.

Al igual que cuando ejecutaste el comando local, verás que se compila la imagen de Docker y, luego, se inicia el trabajo de entrenamiento. Sin embargo, en lugar de ver el resultado del trabajo de entrenamiento, verás el siguiente mensaje que confirma que se inició. Ten en cuenta que la primera vez que ejecutes el comando custom-jobs create, la compilación y el envío de la imagen pueden tardar unos minutos.

Regresa a la sección Entrenamiento de Vertex AI de la consola de Cloud y, en TRABAJOS PERSONALIZADOS, deberías ver que se está ejecutando tu trabajo.

El trabajo tardará unos 20 minutos en completarse.

Una vez que se complete, deberías ver los siguientes artefactos del modelo guardados en el directorio model_output de tu bucket.

🎉 ¡Felicitaciones! 🎉

Aprendiste a usar Vertex AI para hacer lo siguiente:

• Crear contenedores y ejecutar código de entrenamiento de manera local
• Envía trabajos de entrenamiento a Vertex AI Training con el empaquetado automático

Si quieres obtener más información sobre los diferentes componentes de Vertex AI, consulta la documentación.

8. Realiza una limpieza

Debido a que configuramos el notebook para que se agote el tiempo de espera después de 60 minutos de inactividad, no tenemos que preocuparnos por cerrar la instancia. Si quieres cerrar la instancia de forma manual, haz clic en el botón Detener (Stop) en la sección Vertex AI Workbench de la consola. Si quieres borrar el notebook por completo, haz clic en el botón Borrar (Delete).

Para borrar el bucket de almacenamiento, en el menú de navegación de Cloud Console, navega a Almacenamiento, selecciona tu bucket y haz clic en Borrar (Delete):