Framework de machine learning e matemática para Go, com treinamento acelerado por XLA, backend em Go puro e carregamento de modelos do HuggingFace.
1,7 mil84GoApache-2.0Atualizado há 8 dias
google-pytorch/torch_tpu
Um backend de PyTorch que permite executar modelos e operações do PyTorch nas Tensor Processing Units (TPUs) do Google.
- Estrelas
- 590
- Forks
- 33
- Issues abertas
- 282
- Último commit
- há 2 horas
O que é?
O TorchTPU oferece um backend de PyTorch para as TPUs do Google, permitindo executar código PyTorch existente em dispositivos TPU. O projeto inclui kernels ATen personalizados para diversas operações do PyTorch, além de um cache de compilação que reutiliza computações XLA já compiladas para acelerar execuções repetidas.
O projeto oferece suporte a treinamento distribuído com paralelismo de dados e de tensores, integração com torch.compile() para execução em modo de grafo e utilitários de depuração e benchmarking. Os exemplos incluem modelos como Llama, Qwen, ResNet, DLRM, minGPT e UNet, e o projeto também disponibiliza um guia do usuário com notebooks interativos do Marimo. Os notebooks foram feitos para rodar em uma TPU VM.
A instalação pode ser feita via pip (torch_tpu no PyPI) ou com a compilação de um wheel a partir do código-fonte usando Bazel. Há suporte ao Python 3.11 ou superior.
Recursos principais
- 01Implementações de kernels ATen personalizados para TPU.
- 02Cache de compilação para reutilizar computações XLA já compiladas.
- 03Paralelismo de dados e de tensores para treinamento distribuído.
- 04Integração com torch.compile() para execução em modo de grafo.
- 05Exemplos de código com modelos como Llama, Qwen, ResNet e UNet.
- 06Utilitários de depuração e benchmarking.
Casos de uso
- Executar código de treinamento PyTorch existente nas TPUs do Google Cloud.
- Treinar modelos grandes de forma distribuída em vários dispositivos TPU.
- Servir inferências de modelos PyTorch em hardware TPU.
- Estudar exemplos de implementações de modelos comuns em TPU.
Início rápido
- 1
Crie um ambiente virtual Python
python -m venv .venv - 2
Instale a versão pública mais recente do PyPI usando o índice do PyTorch para CPU
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
Antes de usar
- É voltado para hardware Google TPU, portanto só é útil se você tiver acesso a TPUs, como uma TPU VM.
- Requer Python 3.11 ou superior.
- A compilação a partir do código-fonte requer Bazel e, para gerar o wheel completo, credenciais do Remote Build Execution ou a flag no_rbe.
- A licença nos metadados aparece como NOASSERTION, embora o README indique Apache 2.0.
Tópicos
Perguntas frequentes
- O que é torch_tpu?
- Um backend de PyTorch que permite executar modelos e operações do PyTorch nas Tensor Processing Units (TPUs) do Google. O TorchTPU oferece um backend de PyTorch para as TPUs do Google, permitindo executar código PyTorch existente em dispositivos TPU. O projeto inclui kernels ATen personalizados para diversas operações do PyTorch, além de um cache de compilação que reutiliza computações XLA já compiladas para acelerar execuções repetidas.
- torch_tpu é gratuito? Qual é a licença?
- torch_tpu é open source sob a licença Other. Confira os termos da licença antes de usá-lo em produtos comerciais.
- torch_tpu ainda é mantido?
- Manutenção ativa. O último commit de torch_tpu foi em 22 de out. de 2026.
- Para que posso usar torch_tpu?
- Executar código de treinamento PyTorch existente nas TPUs do Google Cloud. Treinar modelos grandes de forma distribuída em vários dispositivos TPU. Servir inferências de modelos PyTorch em hardware TPU. Estudar exemplos de implementações de modelos comuns em TPU.
- Como instalar torch_tpu?
- Crie um ambiente virtual Python: python -m venv .venv Instale a versão pública mais recente do PyPI usando o índice do PyTorch para CPU: pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
- Onde consigo torch_tpu?
- O código-fonte e as instruções de instalação estão no GitHub: https://github.com/google-pytorch/torch_tpu
Resumo gerado por IA a partir do README do projeto. As estatísticas do GitHub são atualizadas diariamente.
Projetos relacionados
Biblioteca open source de LLM baseada em JAX e implementação de referência para pré-treinar e pós-treinar modelos grandes em TPUs e GPUs do Google Cloud.
2,4 mil628PythonApache-2.0Atualizado há 19 horas
Ray é um framework unificado para escalar aplicações Python e de IA, do laptop ao cluster, com runtime distribuído e bibliotecas de IA.
44,1 mil8,2 milPythonApache-2.0Atualizado anteontem
Biblioteca de otimização de deep learning que facilita e torna mais eficiente o treinamento distribuído e a inferência de modelos grandes.
43,2 mil5 milPythonApache-2.0Atualizado há 3 dias
Sistema baseado em PyTorch para treinamento e inferência distribuídos de grandes modelos de IA, com paralelismo, otimização de memória e treinamento heterogêneo.
IA e Machine Learning#deep-learning#distributed-training#model-parallelism#pipeline-parallelism#pytorch
41,4 mil4,5 milPythonApache-2.0Atualizado há 24 horas
Biblioteca PyTorch com arquiteturas de modelos de imagem, pesos pré-treinados e scripts de treinamento, validação e exportação para tarefas de visão computacional.
IA e Machine Learning#pytorch#computer-vision#image-classification#pretrained-models#vision-transformer
37,2 mil5,2 milPythonApache-2.0Atualizado ontem