Entrena HuggingFace GPT2 en Cloud TPU (TF 2.x)


Si no conoces Cloud TPU, te recomendamos que consultes guía de inicio rápido para aprender a crear una VM de TPU.

En este instructivo, se muestra cómo entrenar el modelo GPT2 de HuggingFace en Cloud TPU.

Objetivos

  • Crear una Cloud TPU
  • Instala dependencias
  • Ejecuta el trabajo de entrenamiento

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

  • Compute Engine
  • Cloud TPU

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios. Es posible que los usuarios nuevos de Google Cloud califiquen para obtener una prueba gratuita.

Antes de comenzar

Antes de comenzar este instructivo, verifica que tu proyecto de Google Cloud esté configurado correctamente.

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

  6. En esta explicación, se usan componentes facturables de Google Cloud. Consulta la página de precios de Cloud TPU para calcular los costos. Asegúrate de limpiar los recursos que crees cuando hayas terminado de usarlos para evitar cargos innecesarios.

Entrena HuggingFace GPT2 con Cloud TPU

  1. Abre una ventana de Cloud Shell.

    Abra Cloud Shell

  2. Crea una variable de entorno para el ID de tu proyecto.

    export PROJECT_ID=your-project-id
  3. Configura Google Cloud CLI para usar tu proyecto de Google Cloud en la que quieres crear una Cloud TPU.

    gcloud config set project ${PROJECT_ID}

    La primera vez que ejecutes este comando en una VM de Cloud Shell nueva, se mostrará la página Authorize Cloud Shell. Haz clic en Authorize en la parte inferior de la página para permitir que gcloud realice llamadas a la API de Google Cloud con tus credenciales de Google Cloud.

  4. Crea una cuenta de servicio para el proyecto de Cloud TPU.

    Las cuentas de servicio permiten que el servicio de Cloud TPU acceda a otros Google Cloud de Google Cloud.

    $ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

    El comando muestra una cuenta de servicio de Cloud TPU con el siguiente formato:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

Crear una Cloud TPU

  1. Crea una VM de Cloud TPU con el comando gcloud. El siguiente comando crea una TPU v4-8. También puedes crear una TPU Podslice si configuras la La marca --accelerator-type para un tipo de porción de Pod, por ejemplo, v4-32.

    $ gcloud compute tpus tpu-vm create hf-gpt2 \
      --zone=us-central2-b \
      --accelerator-type=v4-8 \
      --version=tpu-vm-tf-2.17.0-pjrt

    Descripciones de las marcas de comandos

    zone
    Es la zona en la que deseas crear la Cloud TPU.
    accelerator-type
    El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
    version
    La versión de software de Cloud TPU.
  2. Conéctate a la VM de TPU con SSH. Cuando estés conectado a la VM, el símbolo del sistema de shell cambia de username@projectname a username@vm-name:

    gcloud compute tpus tpu-vm ssh hf-gpt2 --zone=us-central2-b

Instala dependencias

  1. Clona el repositorio de Transformers de HuggingFace:

    (vm)$ cd /tmp
    (vm)$ git clone https://github.com/huggingface/transformers.git
    (vm)$ cd transformers
  2. Instala las dependencias:

    (vm)$ pip install .
    (vm)$ pip install -r examples/tensorflow/_tests_requirements.txt
    (vm)$ cd /tmp/transformers/examples/tensorflow/language-modeling
    (vm)$ pip install -r requirements.txt
  3. Crea un directorio temporal:

    (vm)$ mkdir /tmp/gpt2-wikitext
  4. Cuando creas tu TPU, si estableces el parámetro --version en una versión que termina en -pjrt, establece las siguientes variables de entorno para habilitar el entorno de ejecución de PJRT:

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
      (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

Ejecuta la secuencia de comandos de entrenamiento

(vm)$ python3 run_clm.py \
  --model_name_or_path distilgpt2 \
  --max_train_samples 1000 \
  --max_eval_samples 100 \
  --num_train_epochs 1 \
  --output_dir /tmp/gpt2-wikitext \
  --dataset_name wikitext \
  --dataset_config_name wikitext-103-raw-v1

Descripciones de las marcas de comandos

model_name_or_path
El nombre del modelo que se entrenará.
max_train_samples
Es la cantidad máxima de muestras que se usarán para el entrenamiento.
max_eval_samples
Es la cantidad máxima de muestras que se usarán para la evaluación.
num_train_epochs
La cantidad de ciclos de entrenamiento para entrenar el modelo.
output_dir
El directorio de salida para la secuencia de comandos de entrenamiento.
dataset_name
Es el nombre del conjunto de datos que se usará.
dataset_config_name
El nombre de la configuración del conjunto de datos

Cuando se complete el entrenamiento, aparecerá un mensaje similar al siguiente:

  125/125 [============================] - ETA: 0s - loss: 3.61762023-07-07 21:38:17.902850: W tensorflow/core/framework/dataset.cc:956] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
  125/125 [============================] - 763s 6s/step - loss: 3.6176 - val_loss: 3.4233
  Configuration saved in /tmp/gpt2-wikitext/config.json
  Configuration saved in /tmp/gpt2-wikitext/generation_config.json
  Model weights saved in /tmp/gpt2-wikitext/tf_model.h5
  D0707 21:38:45.640973681   12027 init.cc:191]                          grpc_shutdown starts clean-up now
  

Limpia

  1. Desconéctate de la instancia de VM de TPU:

    (vm)$ exit

    El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

  2. Borra el recurso TPU.

    $ gcloud compute tpus tpu-vm delete hf-gpt2 \
    --zone=us-central2-b

¿Qué sigue?

Prueba uno de los otros modelos de referencia admitidos.