Google ha publicado TorchTPU como proyecto de código abierto bajo licencia Apache 2.0, y Amazon ha lanzado TorchNeuron, su backend nativo de PyTorch para Trainium. Las dos noticias apuntan a lo mismo: que cambiar de acelerador deje de exigir reescribir el modelo. Pero “cambiar el dispositivo” esconde bastante trabajo, y es ahí donde se decide si la migración sale barata o no.
Esta guía explica qué ocurre por debajo, cómo localizar el código que te ata a un fabricante y cómo comprobar una migración con pruebas y números propios. Al final tendrás un script de auditoría, una función de paridad numérica y una fórmula para comparar costes.
Qué es un backend de dispositivo en PyTorch
Un tensor vive en un dispositivo. Cuando escribes .to("cuda"), PyTorch envía los datos a la GPU y cada operación se despacha a la implementación de ese dispositivo. Un backend nuevo es, en esencia, un conjunto de implementaciones de esas operaciones integrado con el resto del ecosistema: autograd, torch.compile, distribuidos y el profiler.
AWS describe TorchNeuron como una integración con el mecanismo PrivateUse1 de PyTorch, el punto de extensión previsto para backends de dispositivo que no vienen de serie. Según su documentación, el cambio en código existente es el nombre del dispositivo: de cuda a neuron. Al importar torch el backend se carga solo; hace falta importar un paquete aparte únicamente para sus APIs propias, como los kernels NKI.
Según la cobertura de TorchTPU, el cambio principal también es el destino del dispositivo, y algunos medios citan la sustitución de .to("cuda") por .to("tpu"). No he podido confirmar ese nombre en documentación oficial, así que compruébalo en el repositorio antes de copiarlo.
La lección práctica: cambiar el nombre del dispositivo solo funciona si todo lo demás también lo soporta. Lo que falla en una migración casi nunca es la línea del device.
Auditar: qué código te ata a un fabricante
Hay cuatro fuentes típicas de acoplamiento:
- Dispositivo fijado en el código.
.cuda(),.to("cuda"),device="cuda",torch.device("cuda:0"). - APIs específicas de CUDA.
torch.cuda.synchronize(),torch.cuda.max_memory_allocated()o cualquier llamada atorch.cuda. - Kernels escritos a mano. CUDA nativo (archivos
.cu,__global__) o Triton. No se portan con un cambio de dispositivo. - Dependencias compiladas. Librerías con extensiones de CUDA. Revisa si tienen equivalente para el backend de destino.
Este script localiza los tres primeros puntos en un proyecto:
import re
import sys
from pathlib import Path
PATRONES = {
r"\.cuda\(": "llamada .cuda()",
r"""(device\s*=|\.to\()\s*["']cuda""": "dispositivo fijado a cuda",
r"""torch\.device\(\s*["']cuda""": "torch.device con cuda",
r"torch\.cuda\.": "API específica de torch.cuda",
r"import\s+triton": "kernel Triton",
r"__global__": "kernel CUDA nativo",
}
EXTENSIONES = {".py", ".cu", ".cuh"}
def auditar(raiz):
for ruta in Path(raiz).rglob("*"):
if ruta.suffix not in EXTENSIONES or ".venv" in ruta.parts or not ruta.is_file():
continue
lineas = ruta.read_text(encoding="utf-8", errors="ignore").splitlines()
for n, linea in enumerate(lineas, 1):
for patron, motivo in PATRONES.items():
if re.search(patron, linea):
print(f"{ruta}:{n}: {motivo}: {linea.strip()}")
if __name__ == "__main__":
auditar(sys.argv[1] if len(sys.argv) > 1 else ".")Ejecútalo con python3 audit.py ./mi-proyecto. El resultado es una lista de trabajo, no un veredicto: un .cuda() en un script de exportación pesa mucho menos que uno dentro del forward del modelo.
Migrar por fases, no de golpe
Paso 1: centraliza el dispositivo. Un único punto de decisión, controlado por configuración. Si tu código ya funciona así, este paso no cuesta nada.
import os
import torch
DEVICE = torch.device(os.environ.get(
"TRAIN_DEVICE",
"cuda" if torch.cuda.is_available() else "cpu",
))
model = model.to(DEVICE)
batch = {k: v.to(DEVICE) for k, v in batch.items()}Paso 2: prueba de paridad. Misma entrada, dos dispositivos, y comparación de la salida. Las diferencias pequeñas son normales en coma flotante, y la tolerancia depende de la precisión (fp32, bf16). Fíjala con datos de tu modelo, no a ojo.
import copy
import torch
def paridad(modelo, entrada, dispositivo, atol=1e-4, rtol=1e-3):
modelo = modelo.eval()
with torch.no_grad():
referencia = modelo(entrada)
candidato = copy.deepcopy(modelo).to(dispositivo)
salida = candidato(entrada.to(dispositivo)).cpu()
diferencia = (referencia - salida).abs().max().item()
ok = torch.allclose(referencia, salida, atol=atol, rtol=rtol)
return ok, diferenciaLlámala con el modelo y la entrada en CPU como referencia, por ejemplo paridad(modelo, entrada, "tpu"). Probé la función en CPU con PyTorch 2.x y un modelo pequeño: devuelve (True, 0.0). Si la diferencia máxima crece con la longitud de secuencia o con el batch, revisa la precisión antes de culpar al backend.
Paso 3: mide con tu carga. Los benchmarks de catálogo se hacen en condiciones que casi nunca son las tuyas. Mide tokens por segundo con tu longitud de secuencia, tu batch y tu precisión, y calcula el coste:
coste por millón de tokens = (precio por hora × 1.000.000) / (tokens por segundo × 3.600)
Ejemplo: 4 € por hora con 2.500 tokens por segundo da 4 × 1.000.000 / (2.500 × 3.600) ≈ 0,44 € por millón de tokens. La fórmula solo incluye cómputo. Si comparas de verdad, añade almacenamiento, red y el tiempo de ingeniería que cuesta la migración.
Lo que todavía no sabemos
- Estado de TorchNeuron. Las fuentes no coinciden. La documentación de Neuron 2.33 lo describe como beta pública en trn2 y trn3; una versión anterior de la documentación hablaba de un programa beta cerrado. Revisa las notas de versión.
- Fin de soporte de PyTorch/XLA. Según AWS, el soporte de inferencia con PyTorch/XLA en Trn2 y Trn3 termina en Neuron 2.34. Si tu proyecto depende de eso, es una fecha que cuenta.
- Un dato de rendimiento. Se ha citado un 52 % de model FLOPs utilization al entrenar Qwen 3 8B con hasta 8.192 chips TPU v7x. Es una cifra reportada, no verificada de forma independiente.
- Relación con PyTorch/XLA. Una fuente indica que TorchTPU sustituye a PyTorch/XLA y se recomienda para cargas nuevas. Si ya usas PyTorch/XLA, confirma en la documentación oficial qué te afecta.
Checklist antes de decidir
| Pregunta | Señal de alerta |
|---|---|
| ¿El dispositivo está centralizado? | .cuda() o .to("cuda") repetido en varios módulos |
| ¿Hay kernels a mano? | Archivos .cu, Triton o extensiones compiladas sin equivalente |
| ¿La paridad numérica pasa con tu tolerancia? | Diferencias que crecen con la longitud de secuencia |
| ¿El coste por millón de tokens mejora con tu carga? | La mejora solo aparece en el benchmark del proveedor |
| ¿Tu versión tiene fecha de fin de soporte? | Soporte que termina antes de tu próximo ciclo de trabajo |
[EXPERIENCIA: ¿has migrado alguna carga entre aceleradores? ¿Qué consumió más tiempo: el código, las pruebas de paridad o el coste?]
Que Google y Amazon compitan en el software reduce una barrera real: quedarse con un proveedor porque reescribir cuesta demasiado. No elimina el trabajo de migrar. Si haces webs o integraciones de IA, la pregunta útil no es qué chip usa tu proveedor, sino cuánto te costaría cambiar sin reescribir el proyecto. Eso se puede medir hoy con el script de arriba.
Fuentes
- TechCrunch, 20 de octubre de 2026: Google se une a Amazon en el ataque a Nvidia con desarrolladores de apps de IA
- MarkTechPost: Google open-sources TorchTPU. Licencia, versión del paquete y relación con PyTorch/XLA.
- Introducing TorchNeuron: Native PyTorch backend for AWS Trainium, en el blog de PyTorch
- Documentación de PyTorch en AWS Neuron. Estado de la beta y fechas de soporte.
