Stanford’s Marin Foundation Model: Das erste vollständig offene Modell, entwickelt mit JAX

In der aktuellen Ära der Künstlichen Intelligenz (KI) ist die Offenheit von Modellen ein entscheidender Faktor für die Innovationsgeschwindigkeit. Das Marin-Projekt der Stanford University zielt darauf ab, nicht nur Modelle zu teilen, sondern den gesamten Entwicklungsprozess zugänglich zu machen. Dies umfasst alles von den verwendeten Datensätzen über die Methodologien bis hin zu den Hyperparametern und Trainingsprotokollen.

Einführung in das Marin-Projekt

Das Center for Research on Foundation Models (CRFM) an der Stanford University hat das Marin-Projekt ins Leben gerufen, um eine offene Forschungsumgebung zu schaffen. Hierbei wird das Konzept der Offenheit erweitert, um den gesamten wissenschaftlichen Prozess hinter einem Modell zu umfassen. Das Ziel ist es, Transparenz zu fördern und eine vollständig reproduzierbare Ressource bereitzustellen, die es Forschern ermöglicht, die entwickelten Modelle zu überprüfen, darauf aufzubauen und ihnen zu vertrauen.

Die ersten Modelle: Marin-8B-Base und Marin-8B-Instruct

Die ersten Veröffentlichungen aus diesem offenen Labor sind die Modelle Marin-8B-Base und Marin-8B-Instruct. Diese Modelle, zusammen mit den zugehörigen Daten, dem Code und dem Tokenizer, werden unter der permissiven Apache 2.0 Lizenz veröffentlicht. Diese Verpflichtung zur vollständigen Reproduzierbarkeit stellt eine erhebliche technische Herausforderung dar, die eine Kontrolle über jede Quelle der Variabilität in einem massiv verteilten System erfordert.

Herausforderungen beim Bau offener Foundation-Modelle

Um erfolgreich wirklich offene, skalierbare und reproduzierbare Foundation-Modelle zu erstellen, musste das CRFM-Team mehrere technische Herausforderungen bewältigen. Die Wahl fiel auf JAX als Grundlage, da dessen Designprinzipien direkte Lösungen für diese Probleme bieten. Ein neues Framework namens Levanter wurde entwickelt, um die Leistungsfähigkeit von JAX zu nutzen.

Maximale Geschwindigkeit auf einem einzelnen Beschleuniger

Ein zentrales Problem ist, dass der Kerntrainingsloop Milliarden von Malen ausgeführt wird, was zu einem erheblichen Leistungsengpass führt, wenn eine interpretierte Sprache wie Python verwendet wird. Um dies zu umgehen, kapselt Levanter den gesamten mehrstufigen Trainingsschritt in eine einzige Funktion und verwendet den @jax.jit-Dekorator, um den Prozess in einen hochoptimierten Maschinenkern zu transformieren.

Komplexität des großflächigen Parallelismus verwalten

Das Training von Modellen auf dem neuesten Stand der Technik erfordert die Skalierung auf Tausende von Beschleunigerchips. JAX unterstützt nahtlos die Single-Program, Multiple-Data (SPMD) Parallelisierung, die die zugrunde liegende Datenpartitionierung und Kommunikation automatisiert.

Aufbau und Verwaltung widerstandsfähiger, kosteneffizienter Rechencluster

Für das großflächige Training ist flexibler Zugriff auf massive Rechencluster erforderlich. Levanter nutzt Google Cloud TPU Multislice, um viele kleinere TPU-Slices zu einem logischen Cluster zusammenzufassen und die Kosten zu verwalten.

Wissenschaftliches Vertrauen durch perfekte Reproduzierbarkeit fördern

Ein zentrales Ziel des Marin-Projekts ist es, verifizierbare Wissenschaft zu ermöglichen. Dies erfordert reproduzierbare Ergebnisse, selbst wenn das Training pausiert oder zwischen verschiedenen Maschinenkonfigurationen verschoben wird.

Das Marin-8B-Modell im Detail

Die Architektur des Marin-8B Modells ist eine Llama-Stil-Transformer-Architektur. Der Trainingsprozess war eine adaptive Reise, die über 12 Billionen Tokens umfasste und mehrere Phasen beinhaltete, die sich an neue Daten und Techniken anpassten.

Einladung zur Teilnahme an der Marin-Community

Das Marin-Projekt ist eine offene Einladung zur Teilnahme an der Entwicklung von Foundation-Modellen und zur Mitgestaltung des JAX-Ökosystems. Durch die Bereitstellung einer vollständig reproduzierbaren Ressource zielt das Projekt darauf ab, eine transparentere Zukunft für KI zu schaffen.

Quellen

Quellenliste:

• Quelle: STANFORD’S MARIN FOUNDATION MODEL: THE FIRST FULLY OPEN MODEL DEVELOPED USING JAX
• Marin Community
• Levanter GitHub Repository
• Marin Models on Hugging Face
• JAX Quickstart Guide