Ir al contenido
viernes, 30 de octubre de 202623°Dólar Oficial $1.560Dólar Blue $1.565Dólar MEP $1.568Dólar Tarjeta $2.028
PicoTecno

Google publica en GitHub TorchTPU, un backend nativo de PyTorch para TPU

El proyecto permite ejecutar modelos de PyTorch en las unidades de procesamiento tensorial de Google. Incluye kernels personalizados, caché de compilación y ejemplos de entrenamiento distribuido e inferencia.

Redacción Pico Tecno

Google publica en GitHub TorchTPU, un backend nativo de PyTorch para TPU
(Exey Panteleev from Moscow, Russia / Wikimedia Commons · CC BY 2.0)

Según informó github.com, Google publicó en su plataforma un repositorio con TorchTPU, un backend para PyTorch que permite ejecutar modelos y operaciones de esta biblioteca sobre las TPU (unidades de procesamiento tensorial), los chips que la compañía diseña para cargas de trabajo de inteligencia artificial.

El proyecto incorpora kernels ATen personalizados para distintas operaciones de PyTorch. ATen es la biblioteca de operaciones sobre tensores que PyTorch usa por debajo, y los kernels son las implementaciones concretas de cada operación para un hardware determinado. También incluye una caché de compilación, orientada a optimizar la ejecución, y ejemplos para entrenamiento distribuido y para inferencia de modelos.

Cómo empezar a usarlo

La documentación para usuarios está disponible en un sitio de readthedocs. Los tutoriales interactivos están escritos en Marimo, una herramienta similar a Jupyter, y se ejecutan desde el navegador con el comando marimo edit. Según la fuente, para correrlos hay que hacerlo desde una máquina virtual con TPU, y la opción más sencilla sería una instancia v6e de Google Compute Engine.

TorchTPU sigue la política de versiones de Python de PyTorch y, por el momento, soporta Python 3.11 o superior. La versión pública más reciente, la 0.2.0, está disponible en PyPI y se puede instalar con pip, con un índice adicional que provee la versión de PyTorch para CPU, ya que la instalación por defecto trae la versión para CUDA.

También es posible compilar el paquete desde el código fuente con Bazel, el sistema de compilación que usa el proyecto. Se recomienda instalarlo a través de Bazelisk, un contenedor que descarga la versión adecuada de Bazel. Las compilaciones de Bazel de este repositorio se optimizan por defecto, y el repositorio indica que las compilaciones sin optimizar pueden degradar el rendimiento, por lo que no se usan para generar los paquetes oficiales. Para quienes no tengan credenciales de Remote Build Execution, existe una opción para desactivar esa configuración.

El repositorio también describe cómo configurar el editor VSCode con la extensión clangd para navegar el código C++. El proyecto acepta contribuciones bajo las pautas de su archivo CONTRIBUTING.md, que exige firmar un Contrato de Licencia de Contribuidor (CLA), y se distribuye bajo la licencia Apache 2.0.

Compartí esta nota

Seguí leyendo