Cambiar versiones de software en tu Cloud TPU

Descripción general

La versión de software del framework que se ejecuta en la TPU debe coincidir con la versión que se ejecuta en la VM local. Esta versión del software ahora se puede cambiar en una Cloud TPU en ejecución, sin borrar y volver a crear la TPU. Esto también permite configurar Cloud TPU con versiones específicas de frameworks de software durante la noche. Aún se recomienda seleccionar una versión compatible de estos frameworks.

Uso

La forma recomendada de cambiar las versiones es usar la biblioteca de Python cloud-tpu-client.

Ejemplo de uso para TensorFlow.

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

Esto configura Cloud TPU para que coincida con la versión de TensorFlow que se ejecuta en tu VM local, incluidos los lanzamientos oficiales y las compilaciones con fecha de compilación nocturna.

El parámetro restart_type de la API de configure_tpu_version define el comportamiento de reinicio de la TPU cuando se cambian versiones. Las opciones son “siempre” (opción predeterminada) e “ifneeded”.

  • “siempre” se puede usar para corregir una TPU con, por ejemplo, el estado UNHEALTHY_TENSORFLOW o que muestra errores de memoria insuficiente (OOM) debido a recursos filtrados de una ejecución anterior. Cuando se configura esta opción, la TPU se reinicia incluso cuando no se instala una nueva versión del framework.

  • “ifneeded” puede ser útil porque no reinicia el entorno de ejecución si ya está en la versión correcta, por lo que no agregará ningún tiempo de inicio significativo a una secuencia de comandos de entrenamiento. Cuando se configura esta opción, la TPU solo se reinicia si no tiene instalada la versión del framework correcta.

La biblioteca se comunica directamente con Cloud TPU, por lo que este código debe ejecutarse en una VM en la misma red. Se recomienda ejecutarlo dentro del código para el resto de tu modelo.

Opciones de software adicionales

TensorFlow incluye una string tf.__version__, que es la forma más simple de configurar la versión correcta. Otras opciones de software incluyen lo siguiente:

  • PyTorch: pytorch-1.13, pytorch-nightly-dev20220930, pytorch-nightly
  • Jax: tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

Por ejemplo, para configurar una TPU a fin de que se ejecute con la compilación nocturna más reciente de PyTorch.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()