ColossalAI : des solutions efficaces pour la formation de modèles d'IA à grande échelle

Introduction générale

ColossalAI est une plateforme open source développée par HPC-AI Technologies pour fournir une solution efficace et rentable pour l'apprentissage et l'inférence de modèles d'IA à grande échelle. En prenant en charge de multiples stratégies parallèles, la gestion de la mémoire hétérogène et l'entraînement en précision mixte, ColossalAI est capable de réduire considérablement le temps et la consommation de ressources pour l'entraînement et l'inférence des modèles. Qu'il s'agisse de parallélisme de données, de parallélisme de tenseurs ou de parallélisme de pipeline, ColossalAI fournit des outils et des bibliothèques puissants pour aider les chercheurs et les développeurs à réaliser une formation et une inférence efficaces de modèles à grande échelle sur des clusters multi-GPU.

Liste des fonctions

• Prise en charge du parallélisme des données, du parallélisme tensoriel, du parallélisme en pipeline et d'autres stratégies parallèles
• Formation de précision mixte et optimiseur de redondance zéro (ZeRO)
• Gestion de la mémoire hétérogène pour une formation efficace des grands modèles
• Prise en charge de plusieurs modèles spécifiques à un domaine, tels que Open-Sora, Colossal-LLaMA, etc.
• Fournir des outils conviviaux pour la formation et l'inférence distribuées
• Intégration d'un noyau haute performance, d'un cache KV, d'une attention particulière à la pagination et d'un traitement séquentiel par lots
• Configuration aisée de l'entraînement en parallèle grâce à des fichiers de configuration
• Fournir des exemples et une documentation riche pour vous aider à démarrer rapidement
• Fournit plusieurs options d'installation pour les images Docker et la construction à partir du code source.

Utiliser l'aide

Guide d'installation

Installation à partir de PyPI

Vous pouvez facilement installer Colossal-AI avec la commande suivante :

pip install colossalai

Par défaut, les extensions PyTorch ne sont pas construites lors de l'installation. Si vous avez besoin de construire des extensions PyTorch, vous pouvez définir l'optionBUILD_EXT=1: :

BUILD_EXT=1 pip install colossalai

De plus, nous publions une version NIGHTLY chaque semaine, ce qui vous permet d'accéder aux dernières fonctionnalités inédites et aux corrections de bogues. L'installation se fait comme suit :

pip install colossalai-nightly

Installation à partir de la source

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

Les noyaux CUDA/C++ ne sont pas compilés par défaut. colossalAI les compilera au moment de l'exécution. Activez la fusion des noyaux CUDA si nécessaire :

BUILD_EXT=1 pip install .

Pour les utilisateurs de CUDA 10.2, vous pouvez télécharger manuellement la bibliothèque cub et la copier dans le répertoire approprié avant de l'installer.

Utilisation de Docker

Extraire des images de DockerHub

Vous pouvez obtenir les informations directement auprès de laPage DockerHubExtraire l'image Docker.

Construisez votre propre image

cd ColossalAI
docker build -t colossalai ./docker

Démarre le conteneur en mode interactif :

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

Fonction Opération Déroulement

parallélisme des données

Le parallélisme des données est le processus de division d'un ensemble de données en plusieurs sous-ensembles et d'entraînement du modèle en parallèle sur plusieurs GPUs.ColossalAI facilite l'entraînement en parallèle des données grâce à un profil simplifié de parallélisme des données :

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

parallélisme tensoriel

Le parallélisme tensoriel consiste à diviser le tenseur de paramètres d'un modèle en plusieurs sous-tenseurs et à les calculer en parallèle sur plusieurs GPU. ColossalAI propose des implémentations du parallélisme tensoriel 1D, 2D, 2,5D et 3D :

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

parallèle à la ligne d'assemblage

Le parallélisme de pipeline est la division d'un modèle en plusieurs étapes, chacune exécutée par un ou plusieurs GPU :

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

Entraînement de précision mixte

La formation en précision mixte réduit considérablement l'utilisation de la mémoire et accélère la formation en utilisant une combinaison de nombres à virgule flottante de 16 bits (FP16) et de nombres à virgule flottante de 32 bits (FP32) pendant la formation :

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

Optimiseur de redondance zéro (ZeRO)

L'optimiseur ZeRO réduit considérablement l'empreinte de la mémoire graphique en répartissant les états, les gradients et les paramètres de l'optimiseur sur plusieurs GPU :

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

Applications dans le monde réel

Open-Sora

Open-Sora est la solution complète de ColossalAI pour les modèles de génération vidéo, y compris les paramètres du modèle, les détails de l'entraînement et la capacité de générer des vidéos HD 720p de 16 secondes en un seul clic :

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

Pour plus d'informations, voirOpen-Sora.

Colossal-LLaMA

Colossal-LLaMA fournit une solution open-source pour les grands modèles de langage (LLM) spécifiques à un domaine, permettant d'obtenir des résultats comparables à ceux des grands modèles classiques avec une faible quantité d'argent pour la formation :

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

Pour plus d'informations, voirColossal-LLaMA.