💻 computer science

ShardTensor: Domain Parallelism for Scientific Machine Learning

Dieser Beitrag stellt ShardTensor vor, ein neuartiges Framework für Domänenparallelismus, das die räumliche Dimensionalität der Eingabedaten von Hardwarebeschränkungen entkoppelt, um skalierbares, hochgenaues Training und Inferenz von Scientific-Machine-Learning-Modellen auf Datensätzen mit extremer Auflösung zu ermöglichen.

Ursprüngliche Autoren: Corey Adams, Peter Harrington, Akshay Subramaniam, Mohammad Shoaib Abbas, Jaideep Pathak, Mike Pritchard, Sanjay Choudhry

Veröffentlicht 2026-05-13

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Corey Adams, Peter Harrington, Akshay Subramaniam, Mohammad Shoaib Abbas, Jaideep Pathak, Mike Pritchard, Sanjay Choudhry

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Das große Problem: Die „zu groß für die Tasche"-Box

Stellen Sie sich vor, Sie sind ein Wissenschaftler, der einen Hurrikan, ein Schwarzes Loch oder ein menschliches Gehirn simulieren möchte. Um genaue Ergebnisse zu erhalten, müssen Sie diese Dinge extrem detailliert betrachten – wie beim Heranzoomen auf ein Foto, bis Sie jeden einzelnen Pixel sehen können.

In der Welt der Künstlichen Intelligenz (KI) nennt man dies hochauflösende Daten.

Das Problem ist, dass KI-Modelle auf speziellen Computern namens GPUs (Graphics Processing Units) laufen. Diese GPUs haben einen begrenzten Speicherplatz, ähnlich wie ein kleiner Rucksack.

Das Problem: Wenn Wissenschaftler versuchen, ein massives, hochauflösendes Bild (wie einen 3D-Scan eines Sturms) in die KI einzuspeisen, ist die Datenmenge so riesig, dass sie einfach nicht in den Rucksack passt.
Die alte Lösung: Wissenschaftler mussten die Daten früher „herunterskalieren". Das ist wie das Herunterbrechen eines 4K-Films zu einem winzigen, unscharfen 144p-Vorschaubild, nur damit er in den Rucksack passt. Die KI kann zwar laufen, aber die Ergebnisse sind unscharf und ungenau.
Die andere alte Lösung: Sie könnten versuchen, die Daten auf viele Computer aufzuteilen (Data Parallelism), aber das funktioniert nur, wenn Sie viele separate „Datenstücke" haben (wie 100 verschiedene Stürme). Wenn Sie nur einen einzigen riesigen Sturm zu analysieren haben, stoßen die alten Methoden an eine Wand. Man kann einen Sturm nicht einfach in 100 Teile zerlegen, ohne die Mathematik zu zerstören.

Die Lösung: ShardTensor (Die „Team-Rucksack"-Strategie)

Die Autoren von NVIDIA haben ein neues Werkzeug namens ShardTensor vorgestellt.

Stellen Sie sich ShardTensor als eine magische Möglichkeit vor, eine einzelne, riesige Pizza (Ihre hochauflösenden Daten) in Scheiben zu schneiden und diese Scheiben einem Team von Köchen (GPUs), die in einem Kreis stehen, zu übergeben.

Wie es funktioniert: Anstatt zu versuchen, die ganze Pizza auf einen Teller zu legen, schneidet das System die Pizza räumlich auf (nach Fläche, nicht nach separaten Pizzen).
- Koch A hält die Scheibe oben links.
- Koch B hält die Scheibe oben rechts.
- Koch C hält die Scheibe unten links.
Die Magie: Wenn die Köche Zutaten mischen müssen (Mathematik betreiben), können sie die Ränder ihrer Scheiben an ihre Nachbarn weitergeben. Wenn Koch A wissen muss, was am äußersten Rand seiner Scheibe passiert, fragt er Koch B. Sie arbeiten zusammen, um das ganze Puzzle zu lösen, ohne jemals die ganze Pizza auf einen Teller legen zu müssen.

Dies nennt man Domain Parallelism. Es ermöglicht Wissenschaftlern, Daten zu verarbeiten, die größer sind als der Speicher eines einzelnen Computers, selbst wenn sie nur einen einzigen Datensatz zur Verfügung haben.

Warum das wichtig ist (Der Abschnitt „Warum sich die Mühe machen?")

Das Papier erklärt, dass in der wissenschaftlichen KI der größte Speicherfresser nicht das „Gehirn" der KI (die Modellgewichte) ist; es sind die Zwischenschritte (Aktivierungen).

Analogie: Stellen Sie sich vor, Sie lösen ein riesiges Matheproblem auf einer Tafel. Sie brauchen nicht nur Platz für die endgültige Antwort; Sie brauchen Platz für jede einzelne Kalkulation, die Sie unterwegs anstellen.
Das Ergebnis: Bei hochauflösenden Daten füllen diese „Kalkulationsnotizen" den Speicher sofort auf. ShardTensor verteilt diese Kalkulationen auf viele GPUs.
Der Vorteil:
1. Strong Scaling: Wenn Sie einen massiven Datensatz haben, lässt sich die Arbeit durch Hinzufügen weiterer GPUs viel schneller erledigen (wie das Hinzufügen weiterer Arbeiter auf einer Baustelle).
2. Weak Scaling: Wenn Sie einen Datensatz haben, der so riesig ist, dass er zuvor nicht ausführbar war, können Sie ihn nun ausführen, indem Sie mehr GPUs hinzufügen, um die Last zu teilen.

Reale Beispiele aus dem Papier

Die Autoren testeten dies an zwei spezifischen wissenschaftlichen Problemen, um zu beweisen, dass es funktioniert:

StormScope (Wettervorhersage):
- Die Herausforderung: Die Vorhersage einzelner Gewitter erfordert den Blick auf eine Karte der gesamten Vereinigten Staaten mit sehr hoher Detailgenauigkeit (3 km Auflösung).
- Das Problem: Der Speicher eines einzelnen Computers (80 GB) konnte die Daten für die gesamte US-Karte in dieser Detailgenauigkeit nicht halten. Es war wie der Versuch, die gesamte US-Karte in einem einzigen Ordner mitzuführen.
- Die Lösung: Mit ShardTensor teilten sie die US-Karte auf 32 GPUs auf. Die KI konnte nun das ganze Land in hoher Auflösung „sehen" und Stürme genau vorhersagen, ohne abzustürzen.
Transolver (Aerodynamik):
- Die Herausforderung: Die Simulation von Luftströmungen über ein Auto, um die Kraftstoffeffizienz zu verbessern.
- Das Ergebnis: Sie konnten die KI auf einem Gitter (einem 3D-Gitter) trainieren, das über 1,2 Millionen Punkte enthielt, die die Form des Autos darstellten. Dieses Maß an Detailgenauigkeit war zuvor auf einer einzigen Maschine nicht trainierbar.

Was das Papier nicht sagt (Einschränkungen)

Die Autoren sind offen über die Kompromisse:

Kommunikations-Overhead: Da die GPUs ständig miteinander sprechen müssen, um die Ränder ihrer Datenteile zu teilen, geht ein winziger Teil der Zeit für das „Reden" verloren.
Kleine Daten: Wenn die Daten klein sind, macht diese Gesprächszeit das System langsamer als die Nutzung eines einzelnen Computers. ShardTensor ist nur für riesige Daten geeignet.
Nicht für alles magisch: Es funktioniert am besten mit bestimmten Arten von mathematischen Operationen. Einige ältere oder sehr spezifische Operationen werden möglicherweise noch nicht unterstützt.

Zusammenfassung

ShardTensor ist ein neues Software-Tool, das Wissenschaftlern erlaubt, massive, hochauflösende wissenschaftliche Daten aufzuteilen und auf mehrere Computer zu verteilen. Dies ermöglicht es ihnen, KI-Modelle auf Daten zu trainieren, die so detailliert sind, dass sie zuvor nicht in den Speicher eines einzelnen Computers passten, was zu genaueren Wettervorhersagen, besseren Autodesigns und tieferen wissenschaftlichen Entdeckungen führt.

Technische Zusammenfassung: ShardTensor – Domänenparallelismus für Scientific Machine Learning

1. Problemstellung

Scientific Machine Learning (SciML) steht vor einem kritischen Engpass: Die Unfähigkeit, hochauflösende wissenschaftliche Daten aufgrund von GPU-Speichereinschränkungen zu trainieren und Inferenz darauf durchzuführen. Im Gegensatz zu Large Language Models (LLMs), bei denen der Speicher oft von der Parameteranzahl dominiert wird, sind SciML-Workloads häufig durch intermediate Activations (Zwischenschicht-Aktivierungen) eingeschränkt.

Die Arbeit identifiziert einen „computational curse of dimensionality" (Rechen-Fluch der Dimensionalität) in wissenschaftlichen Daten. Während sich bei einem LLM die Verdopplung der Textlänge die Eingabetokens ungefähr verdoppelt, erhöht sich bei einer Verdopplung der Auflösung von $N$ -dimensionalen wissenschaftlichen Daten die Datengröße um den Faktor $2^N$ . Für 3D- oder höherdimensionale Daten (z. B. Strömungsdynamik, medizinische Bildgebung, Klimamodellierung) wächst der Speicherbedarf zum Speichern der intermediate Activations während des Backward-Passes des Trainings exponentiell.

Bestehende Lösungen skalieren oft nicht oder verschlechtern die Genauigkeit:

Reduzierte Präzision: Bietet nur moderate Speichereinsparungen (Faktor 2) und kann bei wissenschaftlichen Surrogaten unter Problemen mit dem dynamischen Bereich leiden.
Räumliches Downsampling: Reduziert den Speicherbedarf, opfert jedoch die hochauflösende Genauigkeit, die für wissenschaftliche Durchbrüche erforderlich ist.
Modellreduktion: Verringert die Genauigkeit.
Activation Checkpointing/CPU-Offloading: Wirksam, führt jedoch zu erheblichem Rechenaufwand (Neuberechnung oder Host-GPU-Transfers).
Standardparallelismus: Datenparallelismus (Aufteilung des Batches) stößt bei Batch-Größe 1 an eine harte Grenze. Modellparallelismus (Aufteilung der Gewichte) ist für Parameter und nicht für Eingabedaten konzipiert und geht von statischen Tensorformen aus, die bei der Verarbeitung dynamischer Eingabedaten nicht gelten (z. B. Convolutionen, die globale Formen ändern).

Folglich gibt es kein generalisiertes Framework zur Parallelisierung von Eingabedaten unterhalb einer Batch-Größe von eins pro Gerät, was verhindert, dass SciML die nativen Auflösungen erreicht, die Wissenschaftler benötigen.

2. Methodik: ShardTensor

Die Autoren stellen ShardTensor vor, ein neues Paradigma des Domänenparallelismus, das die räumliche Dimensionalität von Eingabedaten von Hardwarebeschränkungen entkoppelt. Es ist eine Erweiterung der PyTorch-Abstraktion DTensor (Distributed Tensor), die speziell für die Handhabung von Eingabedaten und nicht von Modellgewichten entwickelt wurde.

Kern-Designprinzipien

Domänenparallelismus: Definiert als Parallelismus über Geräte hinweg für Eingabedaten, der über die Batch-Dimension hinausgeht. Wenn die Batch-Größe 1 erreicht, unterteilt ShardTensor die räumlichen Dimensionen (z. B. Höhe, Breite, Tiefe) zusätzlich über mehrere GPUs.
Handhabung dynamischer Formen: Im Gegensatz zu DTensor, das statische Verteilungsformen annimmt, die für Modellgewichte geeignet sind, unterstützt ShardTensor dynamische Formen. Es führt eine vierte Komponente von Metadaten ein: Sharding-Formen, die die lokale Chunk-Form jedes Tensors entlang der geteilten Mesh-Achse beschreiben. Dies ermöglicht eine beliebige Chunking von unstrukturierten oder nicht-uniformen Daten.
Imperatives Dispatch: Das Framework ist imperativ und nicht statisch. Es nutzt die PyTorch-Mechanismen __torch_dispatch__ und __torch_function__, um Operationen dynamisch zu routen. Dies ermöglicht eine schichtweise Domänenparallelismus-Verarbeitung ohne vorab kompilierte Berechnungsgraphen oder maßgeschneiderte Patches für jede Schicht.
Erweiterbarkeit: Benutzer können das System über hochlevelige funktionale Schnittstellen (@custom_op) oder low-level Dispatch-Logik (aten-Operationen) erweitern. Dies gewährleistet die Kompatibilität mit sich entwickelnden Modellarchitekturen und benutzerdefinierten Kernen.

Implementierungsdetails

ShardTensor wurde als Erweiterung von PyTorchs DTensor aufgebaut.

Dispatch-Architektur: Es fängt Operationen auf Python-Ebene ab. Wenn eine bestimmte Operation einen Domänen-Parallelismus-Handler besitzt, wird sie ausgeführt; andernfalls wird auf DTensor-Semantik oder das Standardverhalten von PyTorch zurückgegriffen.
Kommunikation: Das Framework verwaltet kollektive Operationen (z. B. All-Reduce, All-Gather) und „Halo"-Austausche (Abruf benachbarter Pixel von benachbarten Geräten), die für Operationen wie Convolution oder Normalisierung erforderlich sind, um die numerische Konsistenz aufrechtzuerhalten.
Integration: Es ist so konzipiert, dass es mit bestehenden Parallelismus-Strategien (Datenparallelismus und Modellparallelismus) kombinierbar ist, was 2D- oder höherdimensionale Parallelisierungsstrategien ermöglicht (z. B. Aufteilung des Batches auf einer Achse und räumlicher Dimensionen auf einer anderen).

3. Hauptbeiträge

ShardTensor-Framework: Ein generalisiertes, quelloffenes Framework (verfügbar über NVIDIA PhysicsNeMo), das Domänenparallelismus für Eingabedaten unterhalb der Batch-Größe eins ermöglicht.
Neuartige Abstraktion: Die Einführung von „Sharding-Formen" zur Handhabung dynamischer Eingabetensor-Verteilungen, was es von gewichtsorientierten Parallelismus-Tools wie DTensor unterscheidet.
Nicht-invasives Programmiermodell: Ein Design, das es Benutzern ermöglicht, Modelleingaben mit minimalen Codeänderungen zu umhüllen, ohne Schichten oder Modelle für Parallelismus neu schreiben zu müssen.
Komponierbarkeit: Nachgewiesene Fähigkeit, Domänenparallelismus mit Daten- und Modellparallelismus auf orthogonalen Mesh-Achsen zu kombinieren.

4. Ergebnisse und Benchmarks

Die Arbeit validiert ShardTensor durch synthetische Benchmarks und reale industrielle Anwendungen auf NVIDIA Blackwell (GB200) und H100 GPUs.

Ring Attention: Bei sehr kleinen Sequenzgrößen dominiert der Parallelismus-Overhead. Bei sehr großen Sequenzgrößen erreicht das System jedoch ein nahezu lineares Strong Scaling mit der GPU-Anzahl sowohl im Training als auch bei der Inferenz.
Vision Transformer (ViT):
- Strong Scaling: Bei hohen Auflösungen (z. B. $4096^2$ ) ist die Inferenz auf 16 GPUs 15-mal schneller als auf einer einzelnen GPU. Das Training auf $2048^2$ ist mit 8 GPUs 5-mal schneller.
- Speichereffizienz: Der Speicherverbrauch skaliert quadratisch für 2D-Daten und kubisch für 3D-Daten, was bestätigt, dass intermediate Activations dominieren. Strong Scaling mit ShardTensor reduziert den Speicher pro Gerät proportional und ermöglicht das Training von 3D-Modellen mit über 1 Milliarde Eingabepunkten auf Standardhardware.
Transolver (Automotive Aerodynamik):
- Trainiert auf dem DrivaerML-Datensatz mit bis zu 1,2 Millionen Punkten pro Domäne.
- Das System behielt numerische Stabilität und wettbewerbsfähige Genauigkeit (L2-Fehler für Druck und Geschwindigkeit) im Vergleich zu Single-GPU-Baselines bei, selbst als die Domänengröße um Faktoren von 2, 4 und 8 zunahm.
- Die gesamte Vorverarbeitungs-Pipeline wurde ebenfalls parallelisiert, um End-to-End-Skalierbarkeit sicherzustellen.
StormScope (Wettervorhersage):
- Ein kontinentales Modell (CONUS) mit 3 km Auflösung ( $1024 \times 1792$ Gitter).
- Der Spitzen-Speicherbedarf (114 GB) überstieg die Kapazität einer einzelnen H100 GPU (80 GB).
- ShardTensor ermöglichte stabiles Training durch Aufteilung der Activations über 2 GPUs innerhalb von 16 datenparallelen Gruppen. Das Modell konvergierte stabil, verfolgte die Loss-Trajektorie einer Single-GPU-Baseline mit niedrigerer Auflösung (6 km) und ermöglichte die Auflösung einzelner konvektiver Stürme, die zuvor ohne Einschränkung des Domänenbereichs nicht zugänglich waren.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass ShardTensor einen fundamentalen Engpass in SciML adressiert: die Unfähigkeit, Daten in der nativen Auflösung zu trainieren und Inferenz darauf durchzuführen, die Wissenschaftler benötigen.

Entblockung auflösungsbegrenzter Workloads: Es wandelt zuvor unmögliche Berechnungen (aufgrund von Speicherobergrenzen) in handhabbare um und ermöglicht hochauflösendes Training in Bereichen wie volumetrischer medizinischer Bildgebung, numerischer Strömungsmechanik und Klimamodellierung.
Komponierbarkeit: Es beseitigt den Kompromiss zwischen Datenvolumen, Modellgröße und Auflösung. Wissenschaftler können nun gleichzeitig entlang mehrerer Parallelismus-Achsen skalieren.
Zugänglichkeit: Durch die Bereitstellung eines nicht-invasiven Programmiermodells senkt es die Eintrittsbarriere für Praktiker, die keine Experten für verteilte Systeme sind, im Gegensatz zu früheren maßgeschneiderten, architekturspezifischen Parallelisierungsversuchen.

Anerkannte Einschränkungen

Die Autoren vermerken bescheiden, dass das imperativ schichtweise Dispatch-Modell Python-Level-Overhead einführt, was es für kleine Operationen oder Daten mit niedriger Auflösung weniger effizient macht. Die Leistung hängt zudem von der Bandbreite der Interconnects ab. Darüber hinaus haben nicht alle PyTorch-Operationen derzeit Domänen-Parallelismus-Dispatch-Pfade, und die vollständige Integration mit torch.compile für statische Graph-Optimierungen läuft noch.

Zusammenfassend stellt ShardTensor einen Schritt dar, um hohen Domänenparallelismus im Scientific Machine Learning so alltäglich zu machen wie Datenparallelismus heute, und ebnet den Weg für das Zeitalter der wissenschaftlichen Foundation Models.