PyTorch es el framework detrás de la mayoría de los modelos de IA y, según algunas estimaciones, lo usa el 85 % de las aplicaciones de IA. Se diseñó pensando en las GPU de Nvidia, una de las razones por las que el dominio de Nvidia en el cómputo de IA es tan difícil de romper. Esta semana, los dos mayores proveedores de nube se movieron a la vez. El 19 de octubre Amazon lanzó TorchNeuron para sus chips Trainium, y el 20 de octubre, en la PyTorch Conference 2026, Google liberó TorchTPU para sus Tensor Processing Units, desarrollado junto al equipo de PyTorch de Meta.
Qué cambia en tu código
Ambos proyectos se conectan a PyTorch mediante su interfaz oficial para hardware nuevo, así que los tensores siguen siendo tensores de PyTorch normales. Para la mayoría de los modelos, pasar a una TPU es cuestión de elegir el dispositivo.
# pip install torch_tpu (on a Cloud TPU VM)
import torch
import torch_tpu # registers the "tpu" device
device = torch.device("tpu")
model = MyModel().to(device)
batch = batch.to(device)
loss = model(batch).loss
loss.backward()
# Optional: compile for extra speed
model = torch.compile(model, backend="tpu")import torch
device = torch.device("cuda")
model = MyModel().to(device)
batch = batch.to(device)
loss = model(batch).loss
loss.backward()
model = torch.compile(model)Cómo funciona TorchTPU
- Modo Strict (por defecto)
- Ejecuta cada operación como PyTorch normal, de forma asíncrona. El más fácil de depurar.
- Modo Debug
- Síncrono, con comprobación de límites, para localizar errores.
- DeferAndFuse
- Agrupa operaciones para reducir el tráfico de memoria; Google indica un 15 % más rápido en su prueba de ajuste fino, y un 50 % o más en algunos casos.
- torch.compile(backend="tpu")
- Captura todo el grafo y lo compila con OpenXLA a código máquina de TPU, con caché entre reinicios.
- Entrenamiento distribuido
- DDP, FSDP2 y DTensor mediante un grupo de procesos de TPU, para cargas idénticas o divergentes.
- Kernels propios
- Kernels Pallas, Tokamax y Helion cuando hace falta velocidad optimizada a mano.
Ejemplo de TorchTitan publicado por Google. Segundos por paso de entrenamiento.
TorchTPU, PyTorch/XLA y TorchNeuron comparados
TorchTPU
TPU de Google- Estado
- código abierto, v0.2.0, Apache 2.0
- Chips
- TPU v5e y posteriores, hasta Ironwood
- Se comporta como PyTorch normal, eager por defecto
- Soporte para vLLM, SGLang y Transformers (device_map="tpu")
La vía que recomienda Google para todo el trabajo nuevo en TPU.
PyTorch/XLA
TPU de Google, anterior- Estado
- en sustitución por TorchTPU
- Registra el grafo de forma diferida, así que hay que marcar qué compilar
- Solo admite el mismo trabajo en todos los chips
Planifica una migración si lo usas hoy.
TorchNeuron
Amazon Trainium- Estado
- beta pública (Neuron 2.33)
- Chips
- Trainium 2 y Trainium 3
- Eager más torch.compile con el compilador Neuron
- FSDP, DTensor, DDP y paralelismo de tensores; servicio con vLLM Neuron
Amazon también está aportando código a PyTorch para que dependa menos de Nvidia.
¿Merece la pena migrar?
- 1
Mira qué ejecutas de verdad
Ajustar un modelo de Hugging Face o servirlo con vLLM es el caso más fácil. Los kernels CUDA propios, el más difícil.
- 2
Pon precio al mismo trabajo en cada chip
Compara el coste por paso de entrenamiento o por millón de tokens servidos, no el precio por hora.
- 3
Empieza con una carga pequeña y real
Porta un trabajo de ajuste fino o de inferencia y compara velocidad y resultados.
- 4
Mantén el código portable
Usa código independiente del dispositivo y torch.compile, para que cambiar de chip sea solo un ajuste.
Pon a prueba lo que has aprendido
1/3¿Cómo se apunta el código PyTorch a una TPU con TorchTPU?
TorchTPU se conecta a la interfaz oficial de PyTorch para hardware nuevo.
2/3¿Qué pasa con PyTorch/XLA?
Google recomienda TorchTPU para todas las cargas nuevas en TPU.
3/3¿Para qué chips es TorchNeuron?
Está en beta pública dentro de Neuron 2.33.
Has acertado 0 de 3.
Más para ingenieros de Python y ML: novedades de scikit-learn 1.9, Python 3.15 y Apache Spark 4.2.
Fuentes: MarkTechPost sobre TorchTPU (20 de octubre de 2026); TechCrunch, «Google has joined Amazon in attacking Nvidia's stronghold with AI app developers» (20 de octubre de 2026); el README de TorchTPU. Las cifras de rendimiento son de Google; consulta la documentación oficial antes de instalar.
Alessandro Ecclesie AgazziEscrito por Alessandro Ecclesie Agazzi, analytics engineer en Londres. Actualizado el 21 de octubre de 2026. Sobre mí