ColossalAI: Effiziente Lösungen für das Training von KI-Modellen in großem Maßstab bereitstellen

Allgemeine Einführung

ColossalAI ist eine Open-Source-Plattform, die von HPC-AI Technologies entwickelt wurde, um eine effiziente und kostengünstige Lösung für das Training und die Inferenz von KI-Modellen in großem Maßstab zu bieten. Durch die Unterstützung mehrerer paralleler Strategien, heterogener Speicherverwaltung und Mixed-Precision-Training ist ColossalAI in der Lage, den Zeit- und Ressourcenverbrauch von Modelltraining und Inferenz erheblich zu reduzieren. Ob Datenparallelität, Tensor-Parallelität oder Pipeline-Parallelität, ColossalAI bietet leistungsstarke Tools und Bibliotheken, die Forschern und Entwicklern helfen, effizientes Training und Inferenz von großen Modellen auf Multi-GPU-Clustern zu erreichen.

Funktionsliste

• Unterstützung von Datenparallelität, Tensorparallelität, Pipelineparallelität und anderen parallelen Strategien
• Gemischtes Präzisionstraining und Null-Redundanz-Optimierer (ZeRO)
• Heterogene Speicherverwaltung zur Unterstützung eines effizienten Trainings von großen Modellen
• Unterstützung für mehrere domänenspezifische Modelle wie Open-Sora, Colossal-LLaMA usw.
• Bereitstellung benutzerfreundlicher Tools für verteiltes Training und Inferenz
• Integration von Hochleistungskernel, KV-Cache, Paging Attention und sequentieller Stapelverarbeitung
• Einfache Konfiguration des parallelen Trainings durch Konfigurationsdateien
• Umfangreiche Beispiele und Dokumentation für einen schnellen Einstieg
• Bietet mehrere Installationsoptionen für Docker-Images und die Erstellung aus Quellcode

Hilfe verwenden

Installationsanleitung

Installation von PyPI

Sie können Colossal-AI einfach mit dem folgenden Befehl installieren:

pip install colossalai

Standardmäßig werden die PyTorch-Erweiterungen nicht während der Installation erstellt. Wenn Sie PyTorch-Erweiterungen bauen müssen, können Sie die OptionBUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

Darüber hinaus veröffentlichen wir jede Woche NIGHTTLY-Versionen, mit denen Sie Zugang zu den neuesten, noch nicht veröffentlichten Funktionen und Fehlerbehebungen haben. Die Installation erfolgt wie folgt:

pip install colossalai-nightly

Installation von der Quelle aus

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install .

CUDA/C++-Kernel werden standardmäßig nicht kompiliert. colossalAI wird sie zur Laufzeit erstellen. Aktivieren Sie die CUDA-Kernel-Fusion, falls erforderlich:

BUILD_EXT=1 pip install .

Benutzer von CUDA 10.2 können die cub-Bibliothek manuell herunterladen und in das entsprechende Verzeichnis kopieren, bevor sie sie installieren.

Docker verwenden

Abrufen von Images von DockerHub

Sie können die Informationen direkt von derDockerHub-SeiteZiehen Sie das Docker-Image.

Bauen Sie Ihr eigenes Image auf

cd ColossalAI
docker build -t colossalai ./docker

Startet den Container im interaktiven Modus:

docker run -ti --gpus all --rm --ipc=host colossalai bash

Funktion Betriebsablauf

Datenparallelität

Datenparallelität ist der Prozess der Aufteilung eines Datensatzes in mehrere Teilmengen und des parallelen Trainings des Modells auf mehreren GPUs. ColossalAI macht es den Benutzern leicht, mit einem vereinfachten Datenparallelitätsprofil zu trainieren:

from colossalai.nn.parallel import DataParallel
model = DataParallel(model)

Tensor-Parallelität

Tensor-Parallelität ist der Prozess der Aufteilung des Parameter-Tensors eines Modells in mehrere Sub-Tensoren und deren parallele Berechnung auf mehreren GPUs. ColossalAI bietet Implementierungen von 1D-, 2D-, 2,5D- und 3D-Tensor-Parallelität:

from colossalai.nn.parallel import TensorParallel
model = TensorParallel(model, parallel_mode='1D')

parallel zum Fließband verlaufend

Pipeline-Parallelität ist die Aufteilung eines Modells in mehrere Phasen, die jeweils von einem oder mehreren Grafikprozessoren ausgeführt werden. ColossalAI bietet eine einfache Konfiguration der Pipeline-Parallelität:

from colossalai.pipeline.parallel import PipelineParallel
model = PipelineParallel(model, num_stages=4)

Gemischtes Präzisionstraining

Mixed-Precision-Training reduziert den Speicherbedarf erheblich und beschleunigt das Training durch die Verwendung einer Kombination aus 16-Bit-Gleitkommazahlen (FP16) und 32-Bit-Gleitkommazahlen (FP32) während des Trainings:

from colossalai.amp import convert_to_amp
model, optimizer, criterion = convert_to_amp(model, optimizer, criterion)

Null-Redundanz-Optimierer (ZeRO)

Der ZeRO-Optimierer reduziert den Grafikspeicherbedarf erheblich, indem er Optimiererzustände, Gradienten und Parameter auf mehrere GPUs verteilt:

from colossalai.zero import ZeroOptimizer
optimizer = ZeroOptimizer(optimizer, model)

Anwendungen in der realen Welt

Open-Sora

Open-Sora ist ColossalAIs Komplettlösung für Videogenerierungsmodelle, einschließlich Modellparametern, Trainingsdetails und der Möglichkeit, 16-sekündige 720p-HD-Videos mit einem Klick zu generieren:

# 训练
python train.py
# 推理
python infer.py

Weitere Informationen finden Sie unterOpen-Sora.

Kolossal-LLaMA

Colossal-LLaMA bietet eine Open-Source-Lösung für domänenspezifische Large Language Models (LLMs), die mit geringem Trainingsaufwand vergleichbare Ergebnisse wie große Mainstream-Modelle erzielen können:

# 训练
python train_llama.py
# 推理
python infer_llama.py

Weitere Informationen finden Sie unterKolossal-LLaMA.