Entrena modelos de PyTorch en Cloud TPU

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

VM de Cloud TPU

Con las VM de Cloud TPU, no hay VM del 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 en el host de TPU. Puedes ejecutar paquetes de instalación de código arbitrarios, ver registros, depurar código, etcétera.

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

image

Para obtener más información sobre cómo entrenar un modelo PyTorch en VM de TPU, consulta la guía de inicio rápido de PyTorch para VM de Cloud TPU.

Nodos de 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 las llamadas de gRPC. El nodo TPU ejecuta un servidor de TensorFlow. El cliente XRT (XLA runtime) se conecta al servidor de TensorFlow a través de la red.

image