ColossalAI: fornecendo soluções eficientes de treinamento de modelos de IA em grande escala

Introdução geral

O ColossalAI é uma plataforma de código aberto desenvolvida pela HPC-AI Technologies para oferecer uma solução eficiente e econômica para o treinamento e a inferência de modelos de IA em grande escala. Ao oferecer suporte a várias estratégias paralelas, gerenciamento de memória heterogênea e treinamento de precisão mista, o ColossalAI é capaz de reduzir significativamente o tempo e o consumo de recursos para treinamento e inferência de modelos. Seja no paralelismo de dados, no paralelismo de tensor ou no paralelismo de pipeline, o ColossalAI oferece ferramentas e bibliotecas avançadas para ajudar pesquisadores e desenvolvedores a obter treinamento e inferência eficientes de modelos em grande escala em clusters com várias GPUs.

Lista de funções

• Suporte a paralelismo de dados, paralelismo de tensor, paralelismo em pipeline e outras estratégias paralelas
• Treinamento de precisão mista e otimizador de redundância zero (ZeRO)
• Gerenciamento de memória heterogêneo para dar suporte ao treinamento eficiente de modelos grandes
• Suporte a vários modelos específicos de domínio, como Open-Sora, Colossal-LLaMA, etc.
• Fornecimento de ferramentas fáceis de usar para treinamento e inferência distribuídos
• Integração de kernel de alto desempenho, cache KV, atenção à paginação e processamento sequencial em lote
• Fácil configuração do treinamento paralelo por meio de arquivos de configuração
• Fornecer exemplos e documentação ricos para ajudá-lo a começar rapidamente
• Oferece várias opções de instalação para imagens do Docker e criação a partir do código-fonte

Usando a Ajuda

Guia de instalação

Instalação a partir do PyPI

Você pode instalar facilmente o Colossal-AI com o seguinte comando:

pip install colossalai

Por padrão, as extensões do PyTorch não são criadas durante a instalação. Se você precisar compilar as extensões do PyTorch, poderá definir a opçãoBUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

Além disso, lançamos versões NIGHTLY todas as semanas, permitindo que você acesse os recursos inéditos e as correções de bugs mais recentes. A instalação é a seguinte:

pip install colossalai-nightly

Instalação a partir da fonte

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

Os kernels CUDA/C++ não são compilados por padrão. O colossalAI os compilará em tempo de execução. Ative a fusão do kernel CUDA, se necessário:

BUILD_EXT=1 pip install .

Para usuários do CUDA 10.2, você pode fazer o download manual da biblioteca cub e copiá-la para o diretório apropriado antes de instalá-la.

Usando o Docker

Extração de imagens do DockerHub

Você pode obter as informações diretamente doPágina do DockerHubExtraia a imagem do Docker.

Crie sua própria imagem

cd ColossalAI
docker build -t colossalai . /docker

Inicia o contêiner no modo interativo:

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

Função Fluxo de operação

paralelismo de dados

O paralelismo de dados é o processo de dividir um conjunto de dados em vários subconjuntos e treinar o modelo em paralelo em várias GPUs. A ColossalAI facilita para os usuários o treinamento em paralelo de dados com um perfil simplificado de paralelismo de dados:

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

paralelismo de tensores

O paralelismo de tensores é o processo de dividir o tensor de parâmetros de um modelo em vários subtensores e computá-los em paralelo em várias GPUs. A ColossalAI oferece implementações de paralelismo de tensores 1D, 2D, 2,5D e 3D:

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

correndo paralelamente à linha de montagem

O paralelismo de pipeline é a divisão de um modelo em vários estágios, cada um executado por uma ou mais GPUs. O ColossalAI oferece uma configuração fácil do paralelismo de pipeline:

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

Treinamento de precisão misto

O treinamento com precisão mista reduz significativamente o uso da memória e acelera o treinamento usando uma combinação de números de ponto flutuante de 16 bits (FP16) e números de ponto flutuante de 32 bits (FP32) durante o treinamento:

from colossalai.amp import convert_to_amp
modelo, otimizador, critério = convert_to_amp(modelo, otimizador, critério)

Otimizador de redundância zero (ZeRO)

O otimizador ZeRO reduz significativamente o espaço ocupado pela memória gráfica ao distribuir estados, gradientes e parâmetros do otimizador em várias GPUs:

from colossalai.zero import ZeroOptimizer
otimizador = ZeroOptimizer(otimizador, modelo)

Aplicativos do mundo real

Open-Sora

O Open-Sora é a solução completa da ColossalAI para modelos de geração de vídeo, incluindo parâmetros de modelo, detalhes de treinamento e a capacidade de gerar vídeos HD 720p de 16 segundos com um clique:

Treinamento #
python train.py
Inferência do #
python infer.py

Para obter mais informações, consulteOpen-Sora.

Colossal-LLaMA

O Colossal-LLaMA oferece uma solução de código aberto para modelos de linguagem grandes (LLMs) específicos de domínio que podem obter resultados comparáveis aos modelos grandes convencionais com uma pequena quantidade de dinheiro para treinamento:

Treinamento #
python train_llama.py
# Inferência
python infer_llama.py

Para obter mais informações, consulteColossal-LLaMA.