Kubeflow Trainer v2.4 führt workload-bewusstes Scheduling für TrainJob ein

Warum verteiltes Training Gruppen-Scheduling benötigt

Das standardmäßige Kubernetes-Scheduling verfolgt einen gierigen Ansatz und versucht, Pods einzeln einzuplanen. Laut Kubeflow führt dieses Modell beim verteilten Training zu Problemen, da Worker möglicherweise gemeinsam starten und nahe beieinander platziert werden müssen.

In einem Beispiel des Projekts benötigt ein Trainingsjob 64 GPUs, findet aber möglicherweise Kapazität für 63 GPUs und bleibt untätig, während er auf die letzte GPU wartet. Ein weiteres Beispiel betrifft einen Dateninitialisierer und 64 GPU-Worker. Wenn die Worker das verfügbare Kontingent verbrauchen, kann der Initialisierer nicht ausgeführt werden. Gleichzeitig können die Worker ohne die Daten nicht starten. Dadurch entsteht ein Abhängigkeits-Deadlock.

Kubeflow Trainer v2.4 nutzt Bausteine der Kubernetes Workload API, darunter Workload und PodGroup, um Gang-Scheduling und topologiebewusstes Scheduling zu unterstützen. Der TrainJob-Controller erstellt diese Ressourcen, damit der Scheduler Pods als einheitliche Gruppe zulassen kann. Nach Angaben des Projekts kann dies verhindern, dass teure GPUs ungenutzt bleiben, und die Kommunikationslatenz verringern.

Drei Ebenen ordnen Workloads der Infrastruktur zu

Trainer v2.4 führt eine dreistufige Scheduling-Hierarchie ein. Auf Ebene 1, der TrainJob-Ebene, repräsentiert eine PodGroup den gesamten Trainingsjob. Das physische Ziel ist eine Cloud-Zone, die über topology.kubernetes.io/zone festgelegt wird. Dies ist für homogene Workloads wie PyTorch DDP gedacht, bei denen alle Worker als eine Einheit eingeplant, platziert und präemptiert werden sollen.

Auf Ebene 2, der ReplicatedJob-Ebene, kann jeder konfigurierte Eintrag eine eigene PodGroup erhalten. Damit lassen sich unterschiedliche Richtlinien innerhalb einer Pipeline unterstützen. Ein Dateninitialisierer kann beispielsweise eine einfache Richtlinie verwenden und einzeln ausgeführt werden, während ein MPI-Launcher und GPU-Worker eine Gang-Richtlinie nutzen und innerhalb derselben Zone koordiniert werden.

Auf Ebene 3, der Job- oder Pro-Replica-Ebene, wird die Platzierung auf Racks oder Top-of-Rack-Switches über topology.kubernetes.io/rack angesteuert. Kubeflow beschreibt diese Ebene als Unterstützung für maximalen Netzwerkdurchsatz und die Isolation von Ausfällen. In v2.4 darf der anvisierte ReplicatedJob höchstens eine Replica haben, da PodGroups pro Replica Kubernetes v1.38 voraussetzen.

Plattformadministratoren konfigurieren die Richtlinie

Die Scheduling-Konfiguration wird zur Laufzeit von Plattformadministratoren festgelegt. Data Scientists können anschließend in ihrem TrainJob einen Wert wie numNodes angeben, ohne die zugrunde liegenden Workload- und PodGroup-Ressourcen manuell erstellen zu müssen.

Für ein verteiltes PyTorch-Beispiel zeigt die Quelle einen TrainJob mit numNodes: 100 und einer Anforderung von vier NVIDIA-GPUs pro Knoten. Der TrainJob-Controller wandelt die angeforderte Knotenzahl in ein Minimum für das Gang-Scheduling um und erstellt einen Workload sowie eine automatisch generierte PodGroup, die der daraus resultierenden Pod-Spezifikation zugeordnet ist.

Laut Kubeflow verbirgt dieses Design die Scheduling-Ressourcen auf niedriger Ebene vor den Anwendern und gibt Plattformingenieuren zugleich Kontrolle über Platzierung und Scheduling-Verhalten. Die Version verbindet damit die TrainJob-Konfiguration mit der physischen Topologie von KI-Rechenzentren. Ziel ist es, verteilte Workloads zu koordinieren und die vom Projekt beschriebenen Ressourcenverschwendungen und Abhängigkeits-Deadlocks zu vermeiden.