Google ha anunciado el lanzamiento de TorchTPU, una nueva herramienta que permite a los desarrolladores utilizar de forma nativa y eficiente PyTorch en sus Tensor Processing Units (TPUs). TorchTPU aborda la creciente complejidad de construir modelos de IA a gran escala, que ahora requieren sistemas distribuidos con miles de aceleradores. El objetivo principal es facilitar la migración de proyectos PyTorch existentes a TPUs con cambios mínimos en el código, al tiempo que se maximiza el rendimiento del hardware.
La arquitectura de TorchTPU se centra en la usabilidad, la portabilidad y el rendimiento. Incorpora tres modos de ejecución (Debug, Strict y Fused Eager), siendo el modo Fused Eager el más destacado por su capacidad para aumentar el rendimiento en un 50-100% mediante la fusión dinámica de operaciones. Además, se integra con la interfaz torch.compile para la compilación de gráficos completos, aprovechando XLA para optimizar el uso de la Inter-Chip Interconnect (ICI). TorchTPU también soporta kernels personalizados escritos en Pallas y JAX, y APIs de PyTorch distribuidas como DDP y FSDPv2, permitiendo una experiencia de desarrollo similar a la de PyTorch en otros dispositivos. Google busca que los desarrolladores puedan aprovechar al máximo las capacidades de sus TPUs, incluso optimizando el diseño de modelos para alinear dimensiones como las de los 'attention heads' con las capacidades de los TensorCores.