ColossalAI: 효율적인 대규모 AI 모델 학습 솔루션 제공

일반 소개

ColossalAI는 대규모 AI 모델 훈련 및 추론을 위한 효율적이고 비용 효율적인 솔루션을 제공하기 위해 HPC-AI Technologies에서 개발한 오픈 소스 플랫폼입니다. 여러 병렬 전략, 이기종 메모리 관리, 혼합 정밀도 훈련을 지원함으로써 ColossalAI는 모델 훈련 및 추론에 소요되는 시간과 리소스 소비를 크게 줄일 수 있습니다. 데이터 병렬 처리, 텐서 병렬 처리, 파이프라인 병렬 처리 등 ColossalAI는 연구자와 개발자가 멀티 GPU 클러스터에서 대규모 모델을 효율적으로 훈련하고 추론할 수 있도록 강력한 도구와 라이브러리를 제공합니다.

 

기능 목록

• 데이터 병렬 처리, 텐서 병렬 처리, 파이프라인 병렬 처리 및 기타 병렬 전략 지원
• 혼합 정밀 훈련 및 제로 중복 옵티마이저(ZeRO)
• 대규모 모델의 효율적인 학습을 지원하는 이기종 메모리 관리
• Open-Sora, Colossal-LLaMA 등과 같은 여러 도메인별 모델을 지원합니다.
• 분산 교육 및 추론을 위한 사용자 친화적인 도구 제공
• 고성능 커널, KV 캐시, 페이징 주의 및 순차적 배치 처리의 통합
• 구성 파일을 통한 간편한 병렬 교육 구성
• 빠르게 시작할 수 있도록 풍부한 예제와 설명서를 제공합니다.
• Docker 이미지 및 소스 코드에서 빌드하기 위한 다양한 설치 옵션 제공

 

도움말 사용

설치 가이드

PyPI에서 설치

다음 명령어로 Colossal-AI를 쉽게 설치할 수 있습니다:

pip install colossalai

기본적으로 PyTorch 확장 프로그램은 설치 중에 빌드되지 않습니다. PyTorch 확장을 빌드해야 하는 경우, 설정할 수 있습니다.BUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

또한 매주 야간 버전을 출시하여 미출시된 최신 기능 및 버그 수정 사항을 확인할 수 있습니다. 설치 방법은 다음과 같습니다:

pip install colossalai-nightly

소스에서 설치

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

CUDA/C++ 커널은 기본적으로 컴파일되지 않습니다. colossalAI가 런타임에 빌드합니다. 필요한 경우 CUDA 커널 퓨전을 활성화합니다:

BUILD_EXT=1 pip install .

CUDA 10.2 사용자의 경우, cub 라이브러리를 수동으로 다운로드하여 설치하기 전에 해당 디렉터리에 복사할 수 있습니다.

Docker 사용

DockerHub에서 이미지 가져오기

다음에서 직접 정보를 얻을 수 있습니다.도커허브 페이지도커 이미지를 가져옵니다.

나만의 이미지 구축

cd ColossalAI
docker build -t colossalai ./docker

대화형 모드에서 컨테이너를 시작합니다:

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

기능 작동 흐름

데이터 병렬 처리

데이터 병렬화는 데이터 집합을 여러 하위 집합으로 나누고 여러 GPU에서 모델을 병렬로 훈련하는 과정으로, ColossalAI를 사용하면 간소화된 데이터 병렬화 프로필을 통해 사용자가 쉽게 데이터 병렬 훈련을 할 수 있습니다:

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

텐서 병렬 처리

텐서 병렬화는 모델의 파라미터 텐서를 여러 개의 서브 텐서로 나누어 여러 GPU에서 병렬로 계산하는 과정으로, ColossalAI는 1D, 2D, 2.5D 및 3D 텐서 병렬화의 구현을 제공합니다:

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

조립 라인과 평행하게 실행

파이프라인 병렬화는 모델을 여러 단계로 나누고 각각 하나 이상의 GPU에서 실행하는 것으로, ColossalAI는 간편한 파이프라인 병렬화 구성을 제공합니다:

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

혼합 정밀 교육

혼합 정밀도 훈련은 훈련 중에 16비트 부동 소수점 숫자(FP16)와 32비트 부동 소수점 숫자(FP32)를 조합하여 사용함으로써 메모리 사용량을 크게 줄이고 훈련 속도를 높입니다:

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

제로 중복성 옵티마이저(ZeRO)

ZeRO 옵티마이저는 옵티마이저 상태, 그라데이션 및 파라미터를 여러 GPU에 분산하여 그래픽 메모리 사용 공간을 크게 줄여줍니다:

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

실제 애플리케이션

오픈소라

Open-Sora는 모델 파라미터, 학습 세부 정보, 클릭 한 번으로 16초 720p HD 동영상을 생성할 수 있는 기능 등 동영상 생성 모델을 위한 ColossalAI의 완벽한 솔루션입니다:

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

자세한 내용은 다음을 참조하세요.오픈소라.

Colossal-LLaMA

Colossal-LLaMA는 적은 학습 비용으로 주류 대규모 모델과 비슷한 결과를 얻을 수 있는 도메인별 대규모 언어 모델(LLM)을 위한 오픈 소스 솔루션을 제공합니다:

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

자세한 내용은 다음을 참조하세요.Colossal-LLaMA.