TorchRec: PyTorch-Erweiterung für verteilte Empfehlungsmodelle

Embedding-Tabellen über viele GPUs verteilen

TorchRec beschreibt sich als PyTorch-Domain-Library, die gemeinsame Sparsity- und Parallelitäts-Primitive für große Recommender-Systeme bereitstellt. Empfehlungsmodelle verbringen einen grossen Teil ihrer Parameter in Embedding-Tabellen, und genau diese Tabellen sollen über viele GPUs verteilt trainier- und inferierbar sein. Die README behauptet, dass TorchRec produktive RecSys-Modelle bei Meta antreibt; das ist eine Selbstauskunft des Herstellers, kein von aussen gepruefter Betrieb.

Die Domänenbibliothek ist in Python geschrieben, das Projekt liegt unter meta-pytorch mit der Homepage pytorch.org/torchrec. Für einen schnellen Einstieg verweist die README auf ein offizielles Tutorial und die API-Dokumentation, nicht auf lokale Beispiele im Klon.

Was die README als Einsatzbelege nennt

Der Abschnitt External Presence dient als Vertrauensanker. Die neueste Version von Metas DLRM, dem Deep Learning Recommendation Model, sei mit TorchRec gebaut. Dazu kommen ein Konferenzpapier zur Disaggregated Multi-Tower-Technik (arXiv 2403.00877), das Machine-Learning-Repository the-algorithm-ml von Twitter sowie ein Abschnitt in der Databricks-Dokumentation zum Training von Empfehlungsmodellen. Ein weiteres Papier behandelt Embedding Offloading hin zu Modellen mit 100 TB Umfang.

Diese Liste zeigt, wohin TorchRec gedacht ist: in Arbeitslasten, bei denen Embedding-Tabellen die Grossteil des Modells ausmachen und sharding-technische Fragen zur Tagesordnung gehören. Für ein kleines Ranking-Modell auf einer einzelnen GPU ist davon nichts relevant.

Sharding-Strategien und ein automatischer Planer

Die README beschreibt fünf Sharding-Varianten für Embedding-Tabellen: data-parallel, table-wise, row-wise, table-wise-row-wise, column-wise und table-wise-column-wise. Die Mischformen erlauben, einzelne Tabellen nach ihrem Zugriffsmuster unterschiedlich aufzuteilen. Dazu gibt es einen Planner, der laut README automatisch optimierte Sharding-Pläne generieren kann.

Diese Auswahl hat direkte Kostenfolgen. Row-wise Sharding verteilt eine Tabelle horizontal, table-wise legt ganze Tabellen auf einzelne Geräte, column-wise teilt die Embedding-Dimension. Welche Strategie am schnellsten läuft, hängt von der Verteilung der Zugriffe ab, und genau darin liegt der Nutzen des Planers. Die README gibt allerdings keine Faustformel, wann welche Variante gewinnt.

Pipelined Training und gekapselte Bausteine

TorchRec überlagert laut README das Laden der Daten, den Transfer auf die GPU, die Gerätekommunikation über input_dist und die Berechnung von Forward und Backward. Dadurch entstehen Pipelines, in denen diese Phasen parallel laufen statt hintereinander. Für RecSys-Kernel setzt das Projekt auf FBGEMM, eine separate Bibliothek unter pytorch.

Zusätzlich sind gemeinsame Module für Empfehlungssysteme enthalten sowie Datensätze: die Criteo-Click-Logs und MovieLens. Das README nennt als Ende-zu-Ende-Beispiel ein DLRM-Event-Prediction-Modell, das auf dem Criteo-Datensatz trainiert wird. Damit ist der Einstieg belegt, aber auch klar, dass TorchRec keine fertige Anwendung mitliefert, sondern Bausteine, die man selbst zu einer Pipeline zusammenbaut.

Installation aus dem Source-Tag

Für die meisten Fälle empfiehlt die README, dem Getting-Started-Pfad der Dokumentation zu folgen und nicht aus dem Source zu bauen. Wer das neueste aus dem Klon will, installiert zunächst PyTorch als Nightly-Build, etwa für CUDA 12.6 mit pip install torch --index-url https://download.pytorch.org/whl/nightly/cu126; für 12.8 und 12.9 wechselt das Ende der URL auf cu128 beziehungsweise cu129, für CPU auf cpu.

Danach folgen git clone --recursive, fbgemm-gpu über den gleichen Nightly-Index und pip install -r requirements.txt. Abschließend python setup.py install develop. Den Test übernimmt torchx, das README gibt torchx run -s local_cwd dist.ddp -j 1x2 --gpu 2 --script test_installation.py vor. Wer CUDA 12.6 oder älter fährt, sollte die Index-URL vorher prüfen, denn die Vorgaben decken nur die genannten Versionen ab.

Quantisierung und der Weg zur C++-Inference

TorchRec unterstützt laut README Quantisierung für reduziert-genaues Training und Inferenz. Dazu gehört auch das Optimieren eines TorchRec-Modells für die C++-Inference. Diese Angaben bleiben auf Konzeptebene, konkrete Befehle für die Quantisierung liefert die README nicht.

Für anspruchsvollere Beispiele verweist sie auf das DLRM-Repository unter facebookresearch, konkret auf torchrec_dlrm/dlrm_main.py. Der Weg von einem trainierbaren Modell zu einem exportierten Artefakt ist damit angedeutet, aber die eigentliche Exportkette muss man in der Dokumentation von FBGEMM und PyTorch nachschlagen.

BSD-3-Clause und ein Rhythmus von Release-Kandidaten

Die README schließt mit einer knappen Lizenzangabe: TorchRec ist BSD-lizenziert, der Text liegt im LICENSE-File. Für die Weiterverwendung bedeutet die BSD-3-Clause wenig Auflagen, wobei Zuschreibungsklauseln zu beachten sind; eine Rechtsberatung ersetzt dieser Satz nicht.

Auffällig ist der Release-Rhythmus. Die GitHub-Versionen tragen Namen wie v1.8.0-rc1, v1.7.0-rc1 und v1.6.0-rc1, also Release-Kandidaten statt stabiler Tags. Wer TorchRec in einen Dienst einbaut, sollte diese Namensgebung mit einplanen und Versionen explizit fixieren, denn der Zeitpunkt, an dem ein Tag als final markiert wird, ist aus den Metadaten allein nicht abzulesen.

Redaktionelles Fazit

TorchRec passt für Teams, die mit PyTorch große Empfehlungsmodelle mit verteilten Embedding-Tabellen trainieren und dazu die dokumentierten Sharding-Strategien brauchen. Wer nur ein einfaches Modell lokal trainiert, kann sich die Installation aus dem Source-Tag mit FBGEMM und PyTorch-Nightly sparen. Vor dem Einsatz ist die DLRM-Beispielkette aus dem README isoliert nachzuziehen und der Arbeitsablauf mit torchx zu prüfen, die BSD-3-Clause-Lizenz erlaubt die Weiterverwendung.

Offizielle Quellen

1. Official documentation
2. Official README
3. Project repository
4. Release notes