Entrena FairSeq Transformer en Cloud TPU con PyTorch

Este instructivo se enfoca en la versión FairSeq de Transformer y la tarea de traducción de WMT 18 de inglés a alemán.

Objetivos

  • Preparar el conjunto de datos
  • Ejecutar el trabajo de entrenamiento
  • Verificar los resultados de salida

Costos

En este instructivo, se usan componentes facturables de Google Cloud, que incluyen los siguientes:

  • Compute Engine
  • Cloud TPU

Usa la calculadora de precios para generar una estimación de los costos según el uso previsto. Los usuarios nuevos de Google Cloud pueden ser elegibles para obtener una prueba gratuita.

Antes de comenzar

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

  1. Accede a tu Cuenta de Google.

    Si todavía no tienes una cuenta, regístrate para obtener una nueva.

  2. En la página de selección de proyectos de Cloud Console, selecciona o crea un proyecto de Cloud.

    Ir a la página Selector de proyectos

  3. Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud. Obtén información sobre cómo confirmar que tienes habilitada la facturación para tu proyecto.

  4. 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.

Configura una instancia de Compute Engine

  1. Abre una ventana de Cloud Shell.

    Abre Cloud Shell

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

    export PROJECT_ID=project-id
    
  3. Configura la herramienta de línea de comandos gcloud para usar el proyecto en el que deseas crear Cloud TPU.

    gcloud config set project ${PROJECT_ID}
    
  4. Desde Cloud Shell, inicia el recurso de Compute Engine necesario para este instructivo.

    gcloud compute --project=${PROJECT_ID} instances create transformer-tutorial \
    --zone=us-central1-a  \
    --machine-type=n1-standard-16  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=200GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    
  5. Conéctate a la instancia nueva de Compute Engine.

    gcloud compute ssh transformer-tutorial --zone=us-central1-a
    

Inicia un recurso de Cloud TPU

  1. En la máquina virtual de Compute Engine, inicia un recurso de Cloud TPU con el siguiente comando:

    (vm) $ gcloud compute tpus create transformer-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.6 \
    --accelerator-type=v3-8
    
  2. Identifica la dirección IP para el recurso Cloud TPU.

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

    La dirección IP se encuentra en la columna NETWORK_ENDPOINTS. Necesitarás esta dirección IP cuando crees y configures el entorno PyTorch.

Descarga los datos

  1. Crea un directorio, pytorch-tutorial-data para almacenar los datos del modelo.

    (vm) $ mkdir $HOME/pytorch-tutorial-data
    
  2. Navega al directorio pytorch-tutorial-data.

    (vm) $ cd $HOME/pytorch-tutorial-data
    
  3. Descarga los datos del modelo.

    (vm) $ wget https://dl.fbaipublicfiles.com/fairseq/data/wmt18_en_de_bpej32k.zip
    
  4. Extrae los datos.

    (vm) $ sudo apt-get install unzip && \
    unzip wmt18_en_de_bpej32k.zip
    

Crea y configura el entorno PyTorch

  1. Inicia un entorno conda.

    (vm) $ conda activate torch-xla-1.6
    
  2. Configura variables de entorno para el recurso de Cloud TPU.

    (vm) $ export TPU_IP_ADDRESS=ip-address; \
    export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

Entrena el modelo

Para entrenar el modelo, ejecuta la siguiente secuencia de comandos:

(vm) $ python /usr/share/torch-xla-1.6/tpu-examples/deps/fairseq/train.py \
  $HOME/pytorch-tutorial-data/wmt18_en_de_bpej32k \
  --save-interval=1 \
  --arch=transformer_vaswani_wmt_en_de_big \
  --max-target-positions=64 \
  --attention-dropout=0.1 \
  --no-progress-bar \
  --criterion=label_smoothed_cross_entropy \
  --source-lang=en \
  --lr-scheduler=inverse_sqrt \
  --min-lr 1e-09 \
  --skip-invalid-size-inputs-valid-test \
  --target-lang=de \
  --label-smoothing=0.1 \
  --update-freq=1 \
  --optimizer adam \
  --adam-betas '(0.9, 0.98)' \
  --warmup-init-lr 1e-07 \
  --lr 0.0005 \
  --warmup-updates 4000 \
  --share-all-embeddings \
  --dropout 0.3 \
  --weight-decay 0.0 \
  --valid-subset=valid \
  --max-epoch=25 \
  --input_shapes 128x64 \
  --num_cores=8 \
  --metrics_debug \
  --log_steps=100

Verifica resultados de salida

Una vez que se completa el trabajo de entrenamiento, puedes encontrar los puntos de control del modelo en el siguiente directorio:

$HOME/checkpoints

Limpia

Realiza una limpieza para evitar incurrir en cargos innecesarios en tu cuenta después de usar los recursos que creaste:

  1. Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

    (vm) $ exit
    

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

  2. En Cloud Shell, usa la herramienta de línea de comandos de gcloud para borrar la instancia de Compute Engine.

    $  gcloud compute instances delete transformer-tutorial  --zone=us-central1-a
    
  3. Use la herramienta de línea de comandos de gcloud para borrar el recurso de Cloud TPU.

    $  gcloud compute tpus delete transformer-tutorial --zone=us-central1-a
    

Qué sigue

Prueba los siguientes colaboradores de PyTorch: