TorchTPU is een project dat een PyTorch-backend biedt voor Google's Tensor Processing Units (TPU's). Het project stelt ontwikkelaars in staat om PyTorch-modellen te draaien op TPU-hardware door gebruik te maken van aangepaste ATen-kernels en een compilatiecache voor het optimaliseren van XLA-berekeningen.
Belangrijkste functionaliteiten:
- Ondersteuning voor gedistribueerde training via Data en Tensor Parallelism.
- Integratie met
torch.compile() voor grafiekmodus-uitvoering.
- Voorbeelden voor populaire modellen zoals Llama, Qwen en ResNet.
- Hulpmiddelen voor benchmarking en debugging.
Installatie en Ontwikkeling: Het project kan worden geïnstalleerd via pip (versie 0.2.0) of handmatig worden gebouwd vanuit broncode met behulp van Bazel. Voor C++-ontwikkelaars wordt een specifieke setup met VSCode en clangd aangeraden om codenavigatie te verbeteren.
Documentatie is beschikbaar via een statische website en interactieve Marimo-notebooks. Het project is open source en valt onder de Apache 2.0-licentie.
TorchTPU: PyTorch Backend voor Google TPU's
Functionaliteiten
- PyTorch ATen-kernel implementatie voor TPU.
- Compilatiecache voor geoptimaliseerde uitvoering van XLA-berekeningen.
- Ondersteuning voor gedistribueerde training (Data Parallelism, Tensor Parallelism).
- Voorbeelden voor diverse modellen (Llama, Qwen, ResNet, DLRM, minGPT, UNet).
- Integratie met
torch.compile() voor uitvoering in grafiekmodus.
- Hulpprogramma's voor debugging en benchmarking.
Gebruikershandleiding en Interactieve Notebooks
Een gebruikershandleiding is beschikbaar op https://torchtpu.readthedocs.io/. Dit is een statische versie van de interactieve notebooks, die te vinden zijn in de docs-map.
De notebooks maken gebruik van Marimo (vergelijkbaar met Jupyter). Deze kunnen worden uitgevoerd met het commando marimo edit file.py, waarna ze via de browser toegankelijk zijn. Hiervoor is een TPU VM vereist; een v6e GCE-instantie is hiervoor de eenvoudigste optie.
Installatie
Dezelfde instructies zijn beschikbaar in de bovengenoemde gebruikershandleiding. Het beleid voor de ondersteuning van Python-versies van TorchTPU komt overeen met de Release Compatibility Matrix van PyTorch.
Momenteel wordt Python 3.11+ ondersteund. Raadpleeg het Python Version Support Policy voor meer details.
Virtuele omgeving aanmaken
Maak eerst een Python-virtuele omgeving aan:
python -m venv .venv
source .venv/bin/activate
Of met uv:
uv venv .venv
source .venv/bin/activate
Installatie via pip
Installeer de nieuwste publieke release (0.2.0 op PyPI). Gebruikers van uv kunnen uv voor alle pip-commando's plaatsen:
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
Bouwen vanuit broncode
Wheels kunnen worden gebouwd via:
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel
Alle OSS Bazel-builds staan standaard op de geoptimaliseerde modus (-c opt) via .bazelrc. Niet-geoptimaliseerde builds (-c fastbuild, -c dbg) kunnen de prestaties verslechteren en worden niet gebruikt voor wheel-builds. Indien u C++-code ontwikkelt of debugt en expliciet een niet-geoptimaliseerde wheel wilt, gebruik dan:
bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel --//:allow_unoptimized_wheel=True
De vlag --config=wheelcommon is verplicht; zonder deze vlag mislukt de build. Deze configuratie stelt --//:wheelbuild=True in, wat de gedeelde XLA/MLIR-backend via de torchversion-reset transitie routeert, zodat pywrap dit verwerkt in één enkele libxlabase.so. Zonder deze vlag verzendt elke per-versie common zijn eigen kopie van de backend, waardoor de resulterende wheel bij het importeren van torch afbreekt vanwege dubbele statische registraties.
wheelcommon activeert ook de Remote Build Execution (RBE) remote cache en remote executor. Als u geen RBE-gegevens heeft, voeg dan --config=norbe toe om deze vlaggen te verwijderen:
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe
Let op dat de volgorde hier belangrijk is: --config=norbe moet ná --config=wheelcommon komen, omdat de laatste configuratie prioriteit heeft voor de remote-cache en remote-executor vlaggen die ze beide instellen.
Installeer de wheel via:
cd <path_to_repo>
python -m venv .venv; source .venv/bin/activate
# De extra-index-url zorgt voor de CPU-versie van torch in plaats van de standaard CUDA-build.
pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu
Let op: Als u installeert zonder --extra-index-url https://download.pytorch.org/whl/cpu, zal de CUDA-versie van torch worden geïnstalleerd. Gebruik de vlag zoals getoond of installeer de CPU-build apart.
Aan de slag
Enkele handige links om te beginnen:
- PyTorch: Officiële PyTorch-documentatie.
- Tutorials and examples: Makkelijk te begrijpen PyTorch-code met gebruik van TorchTPU.
- Public APIs: De TorchTPU-API's die u kunt gebruiken.
Afhankelijkheden
Het project maakt gebruik van Bazel voor afhankelijkheidsbeheer. De belangrijkste Python-afhankelijkheden zijn: (Nog te bepalen/TBD)
Bazel setup
Dit project gebruikt Bazel voor het bouwen en afhankelijkheidsbeheer. Om Bazel te installeren, installeren we in feite Bazelisk, een transparante wrapper die het downloaden en installeren van Bazel zelf afhandelt.
Er zijn verschillende manieren om Bazelisk te installeren:
Als u Go heeft geïnstalleerd:
go install github.com/bazelbuild/bazelisk@latest
Het binary direct downloaden:
wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64
Raadpleeg de Bazelisk README voor meer installatie-instructies. Voeg, ongeacht de installatiemethode, de locatie toe aan uw PATH.
Verifieer de installatie door bazel info uit te voeren in de root-directory van de repository:
bazel info
Dit zou informatie over de repository moeten printen zonder fouten.
VSCode clangd Setup
Voor C++-codenavigatie wordt aangeraden om VSCode te gebruiken met de clangd-extensie. Door te vertrouwen op een compilatiedatabase die is gegenereerd vanuit de Bazel action graph, kan clangd de editor verrijken met slimme functies, waaronder code-aanvullingen, 'go-to-definitions', etc.
Installeer de VSCode-extensies via de UI of via het volgende commando:
code --install-extension llvm-vs-code-extensions.vscode-clangd
# Verwijder Microsoft C++ extensies om interferentie te voorkomen.
code --uninstall-extension ms-vscode.cpptools
Genereer de initiële compile_commands.json:
tools/setup_clangd.py
De eerste run kan een tijd duren (> 30 minuten) om een volledige build te maken, maar volgende runs zullen incrementeel en snel zijn. Als u onlangs bazel build heeft uitgevoerd, bestaan de gegenereerde bestanden al en kunt u --no-build meegeven om het wachten op zware builds over te slaan (let op: in een nieuwe workspace kan het overslaan van de build ervoor zorgen dat clangd gegenereerde bestanden niet kan vinden). Zie de commentaren in het scriptbestand voor meer opties.
Alternatief kunt u het script aanroepen via nox (argumenten na -- worden doorgegeven):
nox -s refresh_compile_commands -- --no-build
nox -s refresh_compile_commands -- //src/torch_tpu/ops/...
Bijdragen
Bijdragen aan het PyTorch on TPU-project zijn welkom! Raadpleeg CONTRIBUTING.md voor richtlijnen over hoe u kunt bijdragen, inclusief het tekenen van onze Contributor License Agreement (CLA) en de processen voor codereviews.
Licentie
Dit project is gelicentieerd onder de Apache 2.0 License. Zie het LICENSE-bestand voor details.
TorchTPU: PyTorch Backend voor Google TPU's
Functionaliteiten
- PyTorch ATen-kernel implementatie voor TPU.
- Compilatiecache voor geoptimaliseerde uitvoering van XLA-berekeningen.
- Ondersteuning voor gedistribueerde training (Data Parallelism, Tensor Parallelism).
- Voorbeelden voor diverse modellen (Llama, Qwen, ResNet, DLRM, minGPT, UNet).
- Integratie met
torch.compile() voor uitvoering in grafiekmodus.
- Hulpprogramma's voor debugging en benchmarking.
Gebruikershandleiding en Interactieve Notebooks
Een gebruikershandleiding is beschikbaar op https://torchtpu.readthedocs.io/. Dit is een statische versie van de interactieve notebooks, die te vinden zijn in de docs-map.
De notebooks maken gebruik van Marimo (vergelijkbaar met Jupyter). Deze kunnen worden uitgevoerd met het commando marimo edit file.py, waarna ze via de browser toegankelijk zijn. Hiervoor is een TPU VM vereist; een v6e GCE-instantie is hiervoor de eenvoudigste optie.
Installatie
Dezelfde instructies zijn beschikbaar in de bovengenoemde gebruikershandleiding. Het beleid voor de ondersteuning van Python-versies van TorchTPU komt overeen met de Release Compatibility Matrix van PyTorch.
Momenteel wordt Python 3.11+ ondersteund. Raadpleeg het Python Version Support Policy voor meer details.
Virtuele omgeving aanmaken
Maak eerst een Python-virtuele omgeving aan:
python -m venv .venv
source .venv/bin/activate
Of met uv:
uv venv .venv
source .venv/bin/activate
Installatie via pip
Installeer de nieuwste publieke release (0.2.0 op PyPI). Gebruikers van uv kunnen uv voor alle pip-commando's plaatsen:
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
Bouwen vanuit broncode
Wheels kunnen worden gebouwd via:
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel
Alle OSS Bazel-builds staan standaard op de geoptimaliseerde modus (-c opt) via .bazelrc. Niet-geoptimaliseerde builds (-c fastbuild, -c dbg) kunnen de prestaties verslechteren en worden niet gebruikt voor wheel-builds. Indien u C++-code ontwikkelt of debugt en expliciet een niet-geoptimaliseerde wheel wilt, gebruik dan:
bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel --//:allow_unoptimized_wheel=True
De vlag --config=wheelcommon is verplicht; zonder deze vlag mislukt de build. Deze configuratie stelt --//:wheelbuild=True in, wat de gedeelde XLA/MLIR-backend via de torchversion-reset transitie routeert, zodat pywrap dit verwerkt in één enkele libxlabase.so. Zonder deze vlag verzendt elke per-versie common zijn eigen kopie van de backend, waardoor de resulterende wheel bij het importeren van torch afbreekt vanwege dubbele statische registraties.
wheelcommon activeert ook de Remote Build Execution (RBE) remote cache en remote executor. Als u geen RBE-gegevens heeft, voeg dan --config=norbe toe om deze vlaggen te verwijderen:
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe
Let op dat de volgorde hier belangrijk is: --config=norbe moet ná --config=wheelcommon komen, omdat de laatste configuratie prioriteit heeft voor de remote-cache en remote-executor vlaggen die ze beide instellen.
Installeer de wheel via:
cd <path_to_repo>
python -m venv .venv; source .venv/bin/activate
# De extra-index-url zorgt voor de CPU-versie van torch in plaats van de standaard CUDA-build.
pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu
Let op: Als u installeert zonder --extra-index-url https://download.pytorch.org/whl/cpu, zal de CUDA-versie van torch worden geïnstalleerd. Gebruik de vlag zoals getoond of installeer de CPU-build apart.
Aan de slag
Enkele handige links om te beginnen:
- PyTorch: Officiële PyTorch-documentatie.
- Tutorials and examples: Makkelijk te begrijpen PyTorch-code met gebruik van TorchTPU.
- Public APIs: De TorchTPU-API's die u kunt gebruiken.
Afhankelijkheden
Het project maakt gebruik van Bazel voor afhankelijkheidsbeheer. De belangrijkste Python-afhankelijkheden zijn: (Nog te bepalen/TBD)
Bazel setup
Dit project gebruikt Bazel voor het bouwen en afhankelijkheidsbeheer. Om Bazel te installeren, installeren we in feite Bazelisk, een transparante wrapper die het downloaden en installeren van Bazel zelf afhandelt.
Er zijn verschillende manieren om Bazelisk te installeren:
Als u Go heeft geïnstalleerd:
go install github.com/bazelbuild/bazelisk@latest
Het binary direct downloaden:
wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64
Raadpleeg de Bazelisk README voor meer installatie-instructies. Voeg, ongeacht de installatiemethode, de locatie toe aan uw PATH.
Verifieer de installatie door bazel info uit te voeren in de root-directory van de repository:
bazel info
Dit zou informatie over de repository moeten printen zonder fouten.
VSCode clangd Setup
Voor C++-codenavigatie wordt aangeraden om VSCode te gebruiken met de clangd-extensie. Door te vertrouwen op een compilatiedatabase die is gegenereerd vanuit de Bazel action graph, kan clangd de editor verrijken met slimme functies, waaronder code-aanvullingen, 'go-to-definitions', etc.
Installeer de VSCode-extensies via de UI of via het volgende commando:
code --install-extension llvm-vs-code-extensions.vscode-clangd
# Verwijder Microsoft C++ extensies om interferentie te voorkomen.
code --uninstall-extension ms-vscode.cpptools
Genereer de initiële compile_commands.json:
tools/setup_clangd.py
De eerste run kan een tijd duren (> 30 minuten) om een volledige build te maken, maar volgende runs zullen incrementeel en snel zijn. Als u onlangs bazel build heeft uitgevoerd, bestaan de gegenereerde bestanden al en kunt u --no-build meegeven om het wachten op zware builds over te slaan (let op: in een nieuwe workspace kan het overslaan van de build ervoor zorgen dat clangd gegenereerde bestanden niet kan vinden). Zie de commentaren in het scriptbestand voor meer opties.
Alternatief kunt u het script aanroepen via nox (argumenten na -- worden doorgegeven):
nox -s refresh_compile_commands -- --no-build
nox -s refresh_compile_commands -- //src/torch_tpu/ops/...
Bijdragen
Bijdragen aan het PyTorch on TPU-project zijn welkom! Raadpleeg CONTRIBUTING.md voor richtlijnen over hoe u kunt bijdragen, inclusief het tekenen van onze Contributor License Agreement (CLA) en de processen voor codereviews.
Licentie
Dit project is gelicentieerd onder de Apache 2.0 License. Zie het LICENSE-bestand voor details.
Groetjes,