Primeros pasos con Gemma en Ray en Vertex AI

JUN 24, 2024
Ju-yeong Ji Gemma DevRel
Ivan Nardini Customer Engineer Google Cloud

Gemma es una familia de modelos abiertos desarrollados a partir de la misma investigación y tecnología utilizadas para crear los modelos Gemini. Los modelos Gemma son capaces de realizar una amplia gama de tareas, incluida la generación de texto, la finalización y generación de código, el ajuste para tareas específicas y la ejecución en varios dispositivos.

Ray es un framework de código abierto para escalamiento de aplicaciones de IA y Python que ofrece la infraestructura para realizar computación distribuida y procesamiento paralelo para tu flujo de trabajo de aprendizaje automático (AA).

Al final de este tutorial, tendrás un conocimiento sólido de cómo usar el ajuste supervisado de Gemma en Ray en Vertex AI para entrenar y entregar modelos de aprendizaje automático de manera eficiente y eficaz.

Puedes explorar el bloc de notas del tutorial “Comenzar con Gemma en Ray en Vertex AI” de GitHub para obtener más información sobre Gemma en Ray. Todo el código que aparece a continuación está en ese bloc de notas para que tu recorrido sea más fácil.


Requisito previo

Son necesarios los siguientes pasos, independientemente de tu entorno.

  1. Selecciona o crea un proyecto de Google Cloud.

2. Asegúrate de que la facturación esté habilitada para tu proyecto.

3. Habilita las APIs.

Si ejecutas este tutorial de forma local, debes instalar el SDK de Cloud.


Costos

Este tutorial utiliza componentes facturables de Google Cloud:

  • Vertex AI

  • Cloud Build

  • Artifact Registry

  • Cloud Storage

Obtén más información sobre los precios y usa la Calculadora de precios para generar una estimación de costos basada en tu uso proyectado.


Qué necesitas

Conjunto de datos

Utilizaremos el conjunto de datos Extreme Summarization (XSum), que es un conjunto de datos sobre sistemas de resumen abstracto de un solo documento.


Bucket de Cloud Storage

Debes crear un bucket de almacenamiento para almacenar los artefactos intermedios, como los conjuntos de datos.

Create a storage bucket to store intermediate artifacts

O con Google Cloud CLI

gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}
 
# for example: gsutil mb -l asia-northeast1 -p test-bebechien gs://test-bebechien-ray-bucket

Repositorio de imágenes de Docker

Para almacenar la imagen personalizada del clúster, crea un repositorio de Docker en Artifact Registry.

Create a Docker repository in the Artifact Registry.

O con Google Cloud CLI

gcloud artifacts repositories create your-repo --repository-format=docker --location=your-region --description="Tutorial repository"

Instancia de TensorBoard en Vertex AI

Una instancia de TensorBoard sirve para rastrear y monitorizar tus trabajos de ajuste. Puedes crear una a partir de Experimentos.

Create a TensorBoard instance from Experiments

O con Google Cloud CLI

ggcloud ai tensorboards create --display-name your-tensorboard --project your-project --region your-region

Cómo configurar un clúster de Ray en Vertex AI

Compila la imagen personalizada del clúster

Para comenzar con Ray en Vertex AI, puedes optar por crear un Dockerfile para una imagen personalizada desde cero o utilizar una de las imágenes básicas de Ray ya compiladas. Una de esas imágenes básicas está disponible aquí.

En primer lugar, prepara el archivo de requisitos que incluye las dependencias que tu aplicación de Ray debe ejecutar.

Luego, crea el Dockerfile para la imagen personalizada utilizando una de las imágenes básicas de Ray en Vertex AI ya compiladas.

Por último, compila la imagen personalizada del clúster de Ray con Cloud Build.

gcloud builds submit --region=your-region
--tag=your-region-docker.pkg.dev/your-project/your-repo/train --machine-type=E2_HIGHCPU_32 ./dockerfile-path

Si todo va bien, verás que la imagen personalizada se ha enviado correctamente a tu repositorio de imágenes de Docker.

The custom image has been successfully pushed to your docker image repository.

También en tu Artifact Registry

The custom image has been successfully pushed to your artifact registry

Crea el clúster de Ray

Puedes crear el clúster de Ray desde Ray en Vertex AI

Ray on Vertex AI

o usar el SDK Python de Vertex AI para crear un clúster de Ray con una imagen personalizada y para personalizar la configuración del clúster. Para obtener más información sobre la configuración del clúster, consulta la documentación.

A continuación, un ejemplo de código Python para crear el clúster de Ray con la configuración personalizada predefinida.

NOTA: la creación de un clúster puede demorar varios minutos, de acuerdo con su configuración.

# Set up Ray on Vertex AI
import vertex_ray
from google.cloud import aiplatform as vertex_ai
from vertex_ray import NodeImages, Resources
 
# Retrieves an existing managed tensorboard given a tensorboard ID
tensorboard = vertex_ai.Tensorboard(your-tensorboard-id, project=your-project, location=your-region)
 
# Initialize the Vertex AI SDK for Python for your project
vertex_ai.init(project=your-project, location=your-region, staging_bucket=your-bucket-uri, experiment_tensorboard=tensorboard)
 
 
HEAD_NODE_TYPE = Resources(
    machine_type= "n1-standard-16",
    node_count=1,
)
 
WORKER_NODE_TYPES = [
    Resources(
        machine_type="n1-standard-16",
        node_count=1,
        accelerator_type="NVIDIA_TESLA_T4",
        accelerator_count=2,
    )
]
 
CUSTOM_IMAGES = NodeImages(
    head="your-region-docker.pkg.dev/your-project/your-repo/train",
    worker="your-region-docker.pkg.dev/your-project/your-repo/train",
)
 
ray_cluster_name = vertex_ray.create_ray_cluster(
    head_node_type=HEAD_NODE_TYPE,
    worker_node_types=WORKER_NODE_TYPES,
    custom_images=CUSTOM_IMAGES,
    cluster_name=your-cluster-name,
)

Ahora puedes obtener el clúster de Ray con get_ray_cluster(). Usa list_ray_clusters() si quieres ver todos los clústeres asociados con tu proyecto.

ray_clusters = vertex_ray.list_ray_clusters()
ray_cluster_resource_name = ray_clusters[-1].cluster_resource_name
ray_cluster = vertex_ray.get_ray_cluster(ray_cluster_resource_name)
 
print("Ray cluster on Vertex AI:", ray_cluster_resource_name)

Ajusta Gemma con Ray en Vertex AI

Para ajustar Gemma con Ray en Vertex AI, puedes usar Ray Train para distribuir transformadores HuggingFace con entrenamiento PyTorch, como puedes ver a continuación.

Con Ray Train, defines una función de entrenamiento que contiene tu código de HuggingFace Transformers para ajustar el tipo de Gemma que quieres distribuir. Luego, defines la configuración de escalamiento para especificar la cantidad deseada de trabajadores e indicar si el proceso de entrenamiento distribuido requiere GPU. Además, puedes definir una configuración de tiempo de ejecución para especificar comportamientos de control y sincronización. Por último, envías el ajuste iniciando un TorchTrainer y lo ejecutas utilizando su método de ajuste.

En este tutorial, ajustaremos Gemma 2B (gemma-2b-it) para resumir artículos de periódicos utilizando un transformador de HuggingFace en Ray en Vertex AI. Escribimos una secuencia de comandos trainer.py simple de Python y la enviaremos al clúster de Ray.


Prepara secuencias de comandos de Python

Preparemos la secuencia de comandos de entrenamiento. A continuación, un ejemplo de secuencia de comandos de Python para inicializar el ajuste de Gemma utilizando la biblioteca HuggingFace TRL.

Luego, prepara la secuencia de comandos de entrenamiento distribuido. A continuación, un ejemplo de secuencia de comandos de Python para ejecutar el trabajo de entrenamiento distribuido de Ray.

Ahora enviamos la secuencia de comandos al clúster de Ray utilizando la API de Ray Jobs a través de la dirección del panel de control de Ray. También puedes encontrar la dirección del panel de control en la página de detalles del clúster, como se muestra a continuación.

Cluster details page showing the dashboard address

En primer lugar, inicia el cliente para que envíe el trabajo.

import ray
from ray.job_submission import JobSubmissionClient
 
client = JobSubmissionClient(
    address="vertex_ray://{}".format(ray_cluster.dashboard_address)
)

Establezcamos una configuración de trabajo que incluya la ruta del modelo, la identificación del trabajo, el punto de entrada de predicción y más.

import random, string, datasets, transformers
from etils import epath
from huggingface_hub import login
 
 
# Initialize some libraries settings
login(token=your-hf-token)
datasets.disable_progress_bar()
transformers.set_seed(8)
 
train_experiment_name = your-experiment-name
train_submission_id = your-submission-id
train_entrypoint = f"python3 trainer.py --experiment-name={train_experiment_name} --logging-dir=”your-bucket-uri/logs” --num-workers=2 --use-gpu"
train_runtime_env = {
    "working_dir": "your-working-dir",
    "env_vars": {"HF_TOKEN": your-hf-token, "TORCH_NCCL_ASYNC_ERROR_HANDLING": "3"},
}

Envía el trabajo

train_job_id = client.submit_job(
    submission_id=train_submission_id,
    entrypoint=train_entrypoint,
    runtime_env=train_runtime_env,
)

Comprueba el estado del trabajo desde el panel de control de OSS.

OSS dashboard showing job status

Comprueba los artefactos de entrenamiento y supervisa el entrenamiento

El uso de Ray en Vertex AI para desarrollar aplicaciones de IA/AA ofrece varios beneficios. En este escenario, puedes usar el almacenamiento de Cloud para almacenar cómodamente puntos de control de modelos, métricas y más. Esto te permite consumir rápidamente el modelo para tareas downstream de IA/AA, incluida la monitorización del proceso de entrenamiento mediante Vertex AI TensorBoard o la generación de predicciones por lotes mediante Ray Data.

Durante el trabajo de entrenamiento de Ray y una vez completado este, verás los artefactos del modelo en la ubicación de Cloud Storage con Google Cloud CLI.

gsutil ls -l your-bucket-uri/your-experiments/your-experiment-name

Puedes usar Vertex AI TensorBoard para validar tu trabajo de entrenamiento registrando las métricas resultantes.

vertex_ai.upload_tb_log(
    tensorboard_id=tensorboard.name,
    tensorboard_experiment_name=train_experiment_name,
    logdir=./experiments,
)
Vertex AI TensorBoard

Valida el entrenamiento de Gemma en Vertex AI

Suponiendo que tu entrenamiento se ejecute con éxito, puedes generar predicciones de forma local para validar el modelo ajustado.

En primer lugar, descarga todos los puntos de control resultantes del trabajo de Ray con Google Cloud CLI.

# copy all artifacts
gsutil ls -l your-bucket-uri/your-experiments/your-experiment-name ./your-experiment-path

Usa el método ExperimentAnalysis para recuperar el mejor punto de control de acuerdo con las métricas y el modo relevantes.

import ray
from ray.tune import ExperimentAnalysis
 
experiment_analysis = ExperimentAnalysis(./your-experiment-path)
log_path = experiment_analysis.get_best_trial(metric="eval_rougeLsum", mode="max")
best_checkpoint = experiment_analysis.get_best_checkpoint(
    log_path, metric="eval_rougeLsum", mode="max"
)

Ahora ya sabes cuál es el mejor punto de control. A continuación, puedes ver un ejemplo de resultado.

example output - code snippet

Carga el modelo ajustado como se describe en la documentación de Hugging Face.

A continuación, te ofrecemos un ejemplo de código Python para cargar el modelo básico y combinar los adaptadores en él para que puedas usarlo como un modelo de transformadores normal. Puedes encontrar el modelo ajustado guardado en tuned_model_path. Por ejemplo, “tutorial/models/xsum-tuned-gemma-it”.

import torch
from etils import epath
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
 
base_model_path = "google/gemma-2b-it"
peft_model_path = epath.Path(best_checkpoint.path) / "checkpoint"
tuned_model_path = models_path / "xsum-tuned-gemma-it"
 
tokenizer = AutoTokenizer.from_pretrained(base_model_path)
tokenizer.padding_side = "right"
 
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path, device_map="auto", torch_dtype=torch.float16
)
peft_model = PeftModel.from_pretrained(
    base_model,
    peft_model_path,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    is_trainable=False,
)
 
tuned_model = peft_model.merge_and_unload()
tuned_model.save_pretrained(tuned_model_path)

Dato de interés: dado que ajustaste un modelo, también puedes publicarlo en Hugging Face Hub utilizando esta única línea de código.

tuned_model.push_to_hub("my-awesome-model")

Para generar resúmenes con el modelo ajustado, usemos el conjunto de validación del conjunto de datos del instructivo.

El siguiente ejemplo de código Python demuestra cómo muestrear un artículo de un conjunto de datos para resumir. Luego, genera el resumen asociado e imprime el resumen de referencia del conjunto de datos y el resumen generado, uno al lado del otro.

import random, datasets
from transformers import pipeline
 
dataset = datasets.load_dataset(
    "xsum", split="validation", cache_dir=./data, trust_remote_code=True
)
 
sample = dataset.select([random.randint(0, len(dataset) - 1)])
document = sample["document"][0]
reference_summary = sample["summary"][0]
 
messages = [
    {
        "role": "user",
        "content": f"Summarize the following ARTICLE in one sentence.\n###ARTICLE: {document}",
    },
]
 
prompt = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
 
tuned_gemma_pipeline = pipeline(
    "text-generation", model=tuned_model, tokenizer=tokenizer, max_new_tokens=50
)
 
generated_tuned_gemma_summary = tuned_gemma_pipeline(
    prompt, do_sample=True, temperature=0.1, add_special_tokens=True
)[0]["generated_text"][len(prompt) :]
 
print(f"Reference summary: {reference_summary}")
print("-" * 100)
print(f"Tuned generated summary: {generated_tuned_gemma_summary}")

A continuación, un ejemplo de resultado del modelo ajustado. Ten en cuenta que el resultado ajustado puede requerir una mayor definición. Para lograr una calidad óptima, es necesario iterar el proceso varias veces, ajustando factores como la tasa de aprendizaje y la cantidad de pasos de entrenamiento.

Example of output from the tuned model

Evalúa el modelo ajustado

Como paso adicional, puedes evaluar el modelo ajustado. Para hacerlo, se comparan los modelos tanto cualitativa como cuantitativamente.

En un caso, comparas las respuestas generadas por el modelo Gemma básico con las generadas por el modelo Gemma ajustado. En el otro caso, calculas las métricas de ROUGE y sus mejoras, lo que te da una idea de la medida en la que el modelo ajustado es capaz de reproducir los resúmenes de referencia correctamente con respecto al modelo básico.

A continuación, se muestra un código Python para evaluar modelos mediante la comparación de los resúmenes generados.

gemma_pipeline = pipeline(
    "text-generation", model=base_model, tokenizer=tokenizer, max_new_tokens=50
)
 
generated_gemma_summary = gemma_pipeline(
    prompt, do_sample=True, temperature=0.1, add_special_tokens=True
)[0]["generated_text"][len(prompt) :]
 
print(f"Reference summary: {reference_summary}")
print("-" * 100)
print(f"Base generated summary: {generated_gemma_summary}")
print("-" * 100)
print(f"Tuned generated summary: {generated_tuned_gemma_summary}")

A continuación, un ejemplo de resultado del modelo básico y del modelo ajustado.

Example output from the base model and tuned model.

Y también te ofrecemos un código para evaluar modelos mediante el cálculo de las métricas de ROUGE y sus mejoras.

import evaluate
 
rouge = evaluate.load("rouge")
 
gemma_results = rouge.compute(
    predictions=[generated_gemma_summary],
    references=[reference_summary],
    rouge_types=["rouge1", "rouge2", "rougeL", "rougeLsum"],
    use_aggregator=True,
    use_stemmer=True,
)
 
tuned_gemma_results = rouge.compute(
    predictions=[generated_tuned_gemma_summary],
    references=[reference_summary],
    rouge_types=["rouge1", "rouge2", "rougeL", "rougeLsum"],
    use_aggregator=True,
    use_stemmer=True,
)
 
improvements = {}
for rouge_metric, gemma_rouge in gemma_results.items():
    tuned_gemma_rouge = tuned_gemma_results[rouge_metric]
    if gemma_rouge != 0:
        improvement = ((tuned_gemma_rouge - gemma_rouge) / gemma_rouge) * 100
    else:
        improvement = None
    improvements[rouge_metric] = improvement
 
print("Base Gemma vs Tuned Gemma - ROUGE improvements")
for rouge_metric, improvement in improvements.items():
    print(f"{rouge_metric}: {improvement:.3f}%")

Y el resultado de ejemplo para la evaluación.

Base Gemma v Tuned Gemma - Rouge improvements

Entregar el modelo Gemma ajustado con Ray Data para predicciones sin conexión

Para generar predicciones sin conexión a escala con Gemma en Ray en Vertex AI, puedes usar Ray Data, una biblioteca de procesamiento de datos escalable para cargas de trabajo de AA.

Al usar Ray Data para generar predicciones sin conexión con Gemma, debes definir una clase de Python para cargar el modelo ajustado en Hugging Face Pipeline. Luego, en función de tu fuente de datos y su formato, usa Ray Data para realizar la lectura de datos distribuidos y un método de conjunto de datos de Ray para aplicar la clase de Python para realizar predicciones en paralelo en múltiples lotes de datos.


Predicción por lotes con Ray Data

Para generar predicciones por lotes con el modelo ajustado utilizando Ray Data en Vertex AI, necesitas un conjunto de datos para generar predicciones y el modelo ajustado almacenado en el bucket de Cloud.

Luego, puedes hacer uso de Ray Data, que proporciona una API fácil de usar para la inferencia por lotes sin conexión.

En primer lugar, carga el modelo ajustado en el almacenamiento de Cloud con Google Cloud CLI

gsutil -q cp -r ./models your-bucket-uri/models

Prepara el archivo de secuencia de comandos de entrenamiento de predicción por lotes para ejecutar el trabajo de predicción por lotes de Ray.

Una vez más, puedes iniciar al cliente para que envíe el trabajo como se indica a continuación con la API de Ray Jobs a través de la dirección del panel de control de Ray.

import ray
from ray.job_submission import JobSubmissionClient
 
client = JobSubmissionClient(
    address="vertex_ray://{}".format(ray_cluster.dashboard_address)
)

Establezcamos una configuración de trabajo que incluya la ruta del modelo, la ID del trabajo, el punto de entrada de predicción y más.

import random, string
 
batch_predict_submission_id = "your-batch-prediction-job"
tuned_model_uri_path = "/gcs/your-bucket-uri/models"
batch_predict_entrypoint = f"python3 batch_predictor.py --tuned_model_path={tuned_model_uri_path} --num_gpus=1 --output_dir=”your-bucket-uri/predictions”"
batch_predict_runtime_env = {
    "working_dir": "tutorial/src",
    "env_vars": {"HF_TOKEN": your-hf-token},
}

Puedes especificar la cantidad de GPU que se utilizarán con el argumento "--num_gpus". Este debe ser un valor que sea igual o menor que el número de GPU disponibles en tu clúster de Ray.

Ya puedes enviar el trabajo.

batch_predict_job_id = client.submit_job(
    submission_id=batch_predict_submission_id,
    entrypoint=batch_predict_entrypoint,
    runtime_env=batch_predict_runtime_env,
)

Echemos un vistazo rápido a los resúmenes generados utilizando un Pandas DataFrame.

import io
import pandas as pd
from google.cloud import storage
 
def read_json_files(bucket_name, prefix=None):
    """Reads JSON files from a cloud storage bucket and returns a Pandas DataFrame"""
 
    # Set up storage client
    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blobs = bucket.list_blobs(prefix=prefix)
 
    dfs = []
 
    for blob in blobs:
        if blob.name.endswith(".json"):
            file_bytes = blob.download_as_bytes()
            file_string = file_bytes.decode("utf-8")
            with io.StringIO(file_string) as json_file:
                df = pd.read_json(json_file, lines=True)
            dfs.append(df)
 
    return pd.concat(dfs, ignore_index=True)
 
predictions_df = read_json_files(prefix="predictions/", bucket_name=your-bucket-uri)
predictions_df = predictions_df[
    ["id", "document", "prompt", "summary", "generated_summary"]
]
predictions_df.head()

A continuación, puedes ver un resultado de ejemplo. La cantidad predeterminada de artículos para resumir es de 20. Puedes especificar la cantidad con el argumento "--sample_size".

Example output of summaries

Resumen

Ya has aprendido muchas cosas, entre ellas, las siguientes:

  • Cómo crear un clúster de Ray en Vertex AI

  • Cómo ajustar Gemma con Ray Train en Vertex AI

  • Cómo validar el entrenamiento de Gemma en Vertex AI

  • Cómo evaluar el modelo Gemma ajustado

  • Cómo entregar Gemma con Ray Data para predicciones sin conexión

Esperamos que este tutorial haya sido esclarecedor y que te haya brindado información valiosa.

Considera unirte al servidor de Discord de la Comunidad de desarrolladores de Google. Te ofrece la oportunidad de compartir tus proyectos, conectarte con otros desarrolladores y participar en debates colaborativos.

No olvides borrar todos los recursos de Google Cloud utilizados en este proyecto. Puedes simplemente eliminar el proyecto de Google Cloud que utilizaste para el tutorial o, de lo contrario, eliminar los recursos individuales que creaste.

# Delete tensorboard
tensorboard_list = vertex_ai.Tensorboard.list()
for tensorboard in tensorboard_list:
    tensorboard.delete()
 
# Delete experiments
experiment_list = vertex_ai.Experiment.list()
for experiment in experiment_list:
    experiment.delete()
 
# Delete ray on vertex cluster
ray_cluster_list = vertex_ray.list_ray_clusters()
for ray_cluster in ray_cluster_list:
    vertex_ray.delete_ray_cluster(ray_cluster.cluster_resource_name)
# Delete artifacts repo
gcloud artifacts repositories delete your-repo -q
 
# Delete Cloud Storage objects that were created
gsutil -q -m rm -r your-bucker-uri