Framework de machine learning y matemáticas para Go, con entrenamiento acelerado por XLA, backend puro en Go y carga de modelos de HuggingFace.
1,7 mil84GoApache-2.0Actualizado hace 8 días
google-pytorch/torch_tpu
Backend de PyTorch que permite ejecutar modelos y operaciones de PyTorch en las unidades de procesamiento tensorial (TPU) de Google.
- Estrellas
- 590
- Forks
- 33
- Issues abiertos
- 282
- Último commit
- hace 2 horas
¿Qué es?
TorchTPU proporciona un backend de PyTorch para las TPU de Google, lo que permite ejecutar código de PyTorch existente en dispositivos TPU. Incluye kernels ATen personalizados para diversas operaciones de PyTorch y una caché de compilación que reutiliza cálculos XLA compilados para acelerar las ejecuciones repetidas.
El proyecto incluye entrenamiento distribuido con paralelismo de datos y de tensores, integración con torch.compile() para ejecutar en modo gráfico y utilidades de depuración y benchmarking. El código de ejemplo incluye modelos como Llama, Qwen, ResNet, DLRM, minGPT y UNet, y también ofrece una guía de usuario con notebooks interactivos de Marimo. Los notebooks están pensados para ejecutarse en una TPU VM.
La instalación se realiza mediante pip (torch_tpu en PyPI) o compilando un wheel desde el código fuente con Bazel. Se admite Python 3.11 o posterior.
Funciones clave
- 01Implementaciones personalizadas de kernels ATen para TPU.
- 02Caché de compilación para reutilizar cálculos XLA compilados.
- 03Paralelismo de datos y de tensores para entrenamiento distribuido.
- 04Integración con torch.compile() para ejecutar en modo gráfico.
- 05Código de ejemplo para modelos como Llama, Qwen, ResNet y UNet.
- 06Utilidades de depuración y benchmarking.
Casos de uso
- Ejecutar código de entrenamiento de PyTorch existente en las TPU de Google Cloud.
- Entrenar modelos grandes de forma distribuida en varios dispositivos TPU.
- Servir inferencias de modelos de PyTorch en hardware TPU.
- Estudiar implementaciones de ejemplo de modelos habituales en TPU.
Inicio rápido
- 1
Crea un entorno virtual de Python
python -m venv .venv - 2
Instala la última versión pública desde PyPI usando el índice de PyTorch para CPU
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
Antes de usarlo
- Está diseñado para hardware Google TPU, por lo que solo resulta útil si tienes acceso a TPU, como una TPU VM.
- Requiere Python 3.11 o posterior.
- Para compilar desde el código fuente se necesita Bazel y, para generar el wheel completo, credenciales de Remote Build Execution o la opción no_rbe.
- La licencia en los metadatos figura como NOASSERTION, aunque el README indica Apache 2.0.
Temas
Preguntas frecuentes
- ¿Qué es torch_tpu?
- Backend de PyTorch que permite ejecutar modelos y operaciones de PyTorch en las unidades de procesamiento tensorial (TPU) de Google. TorchTPU proporciona un backend de PyTorch para las TPU de Google, lo que permite ejecutar código de PyTorch existente en dispositivos TPU. Incluye kernels ATen personalizados para diversas operaciones de PyTorch y una caché de compilación que reutiliza cálculos XLA compilados para acelerar las ejecuciones repetidas.
- ¿torch_tpu es gratuito? ¿Qué licencia tiene?
- torch_tpu es open source bajo la licencia Other. Revisa los términos de la licencia antes de usarlo en productos comerciales.
- ¿torch_tpu sigue mantenido?
- Mantenimiento activo. El último commit de torch_tpu fue el 22 oct 2026.
- ¿Para qué puedo usar torch_tpu?
- Ejecutar código de entrenamiento de PyTorch existente en las TPU de Google Cloud. Entrenar modelos grandes de forma distribuida en varios dispositivos TPU. Servir inferencias de modelos de PyTorch en hardware TPU. Estudiar implementaciones de ejemplo de modelos habituales en TPU.
- ¿Cómo se instala torch_tpu?
- Crea un entorno virtual de Python: python -m venv .venv Instala la última versión pública desde PyPI usando el índice de PyTorch para CPU: pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
- ¿Dónde consigo torch_tpu?
- El código fuente y las instrucciones de instalación están en GitHub: https://github.com/google-pytorch/torch_tpu
Resumen generado por IA a partir del README del proyecto. Las estadísticas de GitHub se actualizan a diario.
Proyectos relacionados
Biblioteca open source de LLM basada en JAX y una implementación de referencia para preentrenar y postentrenar modelos grandes en TPUs y GPUs de Google Cloud.
2,4 mil628PythonApache-2.0Actualizado hace 19 horas
Ray es un framework unificado para escalar aplicaciones de Python e IA desde una laptop hasta un clúster, con un runtime distribuido y bibliotecas de IA.
44,1 mil8,2 milPythonApache-2.0Actualizado anteayer
Una biblioteca de optimización de deep learning que facilita y hace más eficiente el entrenamiento distribuido y la inferencia de modelos grandes.
43,2 mil5 milPythonApache-2.0Actualizado hace 3 días
Sistema de entrenamiento e inferencia distribuidos, basado en PyTorch, para modelos grandes de IA, con paralelismo, optimización de memoria y entrenamiento heterogéneo.
IA y Machine Learning#deep-learning#distributed-training#model-parallelism#pipeline-parallelism#pytorch
41,4 mil4,5 milPythonApache-2.0Actualizado hace 24 horas
Biblioteca de PyTorch con arquitecturas de modelos de imagen, pesos preentrenados y scripts de entrenamiento, validación y exportación para tareas de visión artificial.
IA y Machine Learning#pytorch#computer-vision#image-classification#pretrained-models#vision-transformer
37,2 mil5,2 milPythonApache-2.0Actualizado ayer