Entrena DLRM en Cloud TPU con PyTorch

En este instructivo, se muestra cómo entrenar DLRM de Facebook Research en Cloud TPU.

Objetivos

  • Crea y configura el entorno PyTorch
  • Ejecuta el trabajo de entrenamiento con datos falsos
  • Entrena en el conjunto de datos de Kaggle de Criteo (opcional)

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. Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

  2. En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

    Ir al selector de proyectos

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

  4. En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

    Ir al selector de proyectos

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

    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.

Configura una instancia de Compute Engine

  1. Abre una ventana de Cloud Shell.

    Abra Cloud Shell

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

    export PROJECT_ID=project-id

  3. Configura Google Cloud CLI para usar el proyecto en el que deseas crear 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.

  4. Desde Cloud Shell, inicia el recurso de Compute Engine necesario para este instructivo. Nota: Se recomienda usar un machine-type n1-highmem-96 si entrenas en el conjunto de datos de Kaggle de Criteo.

    gcloud compute instances create dlrm-tutorial \
--zone=us-central1-a \
--machine-type=n1-standard-64 \
--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 dlrm-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 dlrm-tutorial \
--zone=us-central1-a \
--network=default \
--version=pytorch-2.0  \
--accelerator-type=v3-8

  2. Identifica la dirección IP para el recurso Cloud TPU.

    (vm) $ gcloud compute tpus describe dlrm-tutorial --zone=us-central1-a

Crea y configura el entorno PyTorch

  1. Inicia un entorno conda.

    (vm) $ conda activate torch-xla-2.0

  2. Configura variables de entorno para el recurso de Cloud TPU.

    (vm) $ export TPU_IP_ADDRESS=ip-address

    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

Ejecuta el trabajo de entrenamiento con datos falsos

  1. Instala las dependencias.

    (vm) $ pip install onnx

  2. Ejecuta el modelo con datos aleatorios. Este proceso debería tardar entre 5 y 10 minutos.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
    --arch-embedding-size=1000000-1000000-1000000-1000000-1000000-1000000-1000000-1000000 \
    --arch-sparse-feature-size=64 \
    --arch-mlp-bot=512-512-64 \
    --arch-mlp-top=1024-1024-1024-1 \
    --arch-interaction-op=dot \
    --lr-num-warmup-steps=10 \
    --lr-decay-start-step=10 \
    --mini-batch-size=2048 \
    --num-batches=1000 \
    --data-generation='random' \
    --numpy-rand-seed=727 \
    --print-time \
    --print-freq=100 \
    --num-indices-per-lookup=100 \
    --use-tpu \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len=8 \
    --tpu-metrics-debug \
    --tpu-cores=8

Entrena en el conjunto de datos de Kaggle de Criteo (opcional)

Estos pasos son opcionales. Solo debes ejecutarlos si deseas entrenar en el conjunto de datos de Criteo Kaggle.

  1. Descarga el conjunto de datos.

    Descarga el conjunto de datos de Criteo Kaggle según las instrucciones que se indican aquí. Cuando se complete la descarga, copia el archivo dac.tar.gz en un directorio llamado ./criteo-kaggle/. Usa el comando tar -xzvf para extraer el contenido del archivo tar.gz en el directorio ./critero-kaggle.

     (vm) $ mkdir criteo-kaggle
 (vm) $ cd criteo-kaggle
 (vm) $ # Download dataset from above link here.
 (vm) $ tar -xzvf dac.tar.gz
 (vm) $ cd ..

  2. Procesa previamente el conjunto de datos.

    Inicia esta secuencia de comandos para procesar previamente el conjunto de datos de Criteo. Esta secuencia de comandos produce un archivo llamado kaggleAdDisplayChallenge_processed.npz y tarda más de 3 horas en procesar previamente el conjunto de datos.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_data_pytorch.py \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --mini-batch-size=128 \
    --memory-map \
    --test-mini-batch-size=16384 \
    --test-num-workers=4

  3. Verifica si el procesamiento previo se realizó correctamente.

    Deberías ver el archivo kaggleAdDisplayChallenge_processed.npz en el directorio criteo-kaggle.

  4. Ejecuta la secuencia de comandos de entrenamiento en el conjunto de datos previamente procesado de Criteo Kaggle.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
    --arch-sparse-feature-size=16 \
    --arch-mlp-bot="13-512-256-64-16" \
    --arch-mlp-top="512-256-1" \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --processed-data-file=criteo-kaggle/kaggleAdDisplayChallenge_processed.npz \
    --loss-function=bce \
    --round-targets=True \
    --learning-rate=0.1 \
    --mini-batch-size=128 \
    --print-freq=1024 \
    --print-time \
    --test-mini-batch-size=16384 \
    --test-num-workers=4 \
    --memory-map \
    --test-freq=101376 \
    --use-tpu \
    --num-indices-per-lookup=1 \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len 8 \
    --tpu-metrics-debug \
    --tpu-cores=8

    El entrenamiento debe completarse en 2 horas o más con una exactitud del 78.75% o más.

Realiza una limpieza

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 Google Cloud CLI para borrar la instancia de Compute Engine:

    $ gcloud compute instances delete dlrm-tutorial --zone=us-central1-a

  3. Usa Google Cloud CLI para borrar el recurso de Cloud TPU.

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

¿Qué sigue?

Prueba los siguientes colaboradores de PyTorch: