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

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

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

Ejecuta el trabajo de entrenamiento con datos falsos

  1. Para obtener DLRM compatible con TPU, ejecuta el siguiente comando:

    (vm) $ git clone --recursive https://github.com/pytorch-tpu/examples.git

  2. Instala las dependencias.

    (vm) $ pip install onnx

  3. Ejecuta el modelo con datos aleatorios.

    (vm) $ python 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)

  1. Descarga el conjunto de datos.

    Descarga el conjunto de datos del conjunto de datos Criteo Kaggle según las instrucciones que se indican aquí. Una vez que termines de descargar el archivo dac.tar.gz, colócalo en el directorio ./criteo-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.

    Ejecuta el siguiente comando para procesar previamente el conjunto de datos de Criteo. Ten en cuenta que es prácticamente el mismo comando que el comando de entrenamiento real, pero lo estamos ejecutando en un solo núcleo, ya que queremos finalizar la ejecución del entrenamiento una vez que se haya terminado el procesamiento previo y que se haya producido el archivo kaggleAdDisplayChallenge_processed.npz. Esto llevará mucho tiempo (hora+) en procesarse.

    (vm) $ python 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 \
    --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 \
    --test-freq=101376 \
    --use-tpu \
    --num-indices-per-lookup=1 \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len 1 \
    --tpu-metrics-debug \
    --tpu-cores=1

    Una vez que veas registros como los que se muestran a continuación, el conjunto de datos ya se procesó previamente y puedes finalizar (Ctrl^C) la ejecución y ejecutar pkill python solo para asegurarte de que no haya procesos pendientes.

    Load 6548659/6548660  Split: 0  Label True: 0  Stored: 0
Saved /home/jysohn/test/train_day_0.npz!
Load 6548659/6548660  Split: 1  Label True: 1  Stored: 1
Saved /home/jysohn/test/train_day_1.npz!
Load 3742283/6548660  Split: 2  Label True: 0  Stored: 0
...
time/loss/accuracy (if enabled): 2020-08-26 13:17:18.139566
Finished training it 1024/306968 of epoch 0, -1.00 ms/it, loss 0.518172, accuracy 75.720 %, 131072 samples, @ 2020-08-26 13:18:04.515570

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

    Deberías ver que kaggleAdDisplayChallenge_processed.npz se produjo dentro del directorio criteo-kaggle.

  4. Ejecuta el entrenamiento en el conjunto de datos preprocesado de Kaggle de Criteo.

    (vm) $ python 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 \
    --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 3 horas o más con una exactitud del 78.75% o más.

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 la herramienta de línea de comandos de gcloud para borrar la instancia de Compute Engine.

    $ gcloud compute instances delete dlrm-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 dlrm-tutorial --zone=us-central1-a

Qué sigue

Prueba los siguientes colaboradores de PyTorch: