Entrena modelos de PyTorch en Cloud TPU

En este documento, se describen los detalles de la ejecución de modelos de PyTorch en Cloud TPU. La forma en la que entrenas los modelos de PyTorch depende de la arquitectura de Cloud TPU que uses. Para obtener más información, consulta Arquitectura del sistema.

VM de Cloud TPU

Con las VM de Cloud TPU, no hay una VM de usuario y tu modelo y código se ejecutan en una VM en la máquina anfitrión de TPU. Puedes establecer una conexión SSH directamente al host de TPU. Puedes ejecutar paquetes de instalación de código arbitrarios, ver registros, depurar código, etc.

PyTorch/XLA crea un servidor local de TensorFlow cada vez que se ejecuta un nuevo programa de PyTorch/XLA. El cliente de XRT (tiempo de ejecución de XLA) se conecta al servidor local de TensorFlow.

imagen

Para obtener más información sobre el entrenamiento de un modelo de PyTorch en VM de TPU, consulta la Guía de inicio rápido de PyTorch para VM de Cloud TPU.

Nodos TPU heredados

Cuando usas nodos de TPU heredados, necesitas una VM de usuario para cada TPU. Las VM de usuario se comunican con los nodos TPU a través de llamadas de gRPC. El nodo TPU ejecuta un servidor de TensorFlow. El cliente XRT (entorno de ejecución de XLA) se conecta al servidor de TensorFlow mediante la red.

imagen