Axolotl - 강력한 AI 모델 후속 훈련 도구

프로젝트 개요

Axolotl은 다양한 AI 모델 후속 훈련 과정을 간소화하기 위해 특별히 설계된 오픈 소스 도구입니다. 후속 훈련은 사전 훈련된 모델에 대한 모든 수정 또는 추가 훈련을 의미하며, 여기에는 전체 모델 미세 조정, 파라미터 효율적 조정(예: LoRA 및 QLoRA), 지도 미세 조정(SFT), 명령어 조정 및 정렬 기술이 포함됩니다. 이 도구는 다양한 모델 아키텍처와 훈련 구성을 지원하여 사용자가 이러한 고급 훈련 기술을 쉽게 시작할 수 있도록 합니다.

핵심 기능

모델 지원

• 다양한 모델 아키텍처: LLaMA, Pythia, Falcon, MPT, Mistral, Mixtral 등 주요 대규모 언어 모델을 포함하여 다양한 Hugging Face 모델 훈련 지원
• 유연한 훈련 방법: 전체 미세 조정(Full Fine-tuning), LoRA, QLoRA, ReLoRA, GPTQ 등 다양한 훈련 방식 지원

구성 관리

• YAML 구성 파일: 간단한 YAML 파일에 데이터 세트 전처리, 모델 훈련/미세 조정, 모델 추론 또는 평가 등 필요한 모든 구성 포함
• CLI 덮어쓰기: 명령줄 인수를 통해 구성 파일의 설정 덮어쓰기 지원
• 유연한 구성: 다양한 훈련 파라미터 및 모델 설정 사용자 정의 가능

데이터 처리 능력

• 다중 형식 데이터 세트: 로컬, HuggingFace 및 클라우드(S3, Azure, GCP, OCI) 데이터 세트 로드 지원
• 사용자 정의 형식: 사용자 정의 형식을 사용하거나 이미 토큰화된 데이터 세트를 직접 가져올 수 있습니다.
• 데이터 세트 전처리: 강력한 데이터 전처리 기능 내장

성능 최적화

• 첨단 최적화 기술: xformers, Flash Attention, Liger Kernel, 밧줄 축소 및 다중 패키징 기술 통합
• 다중 GPU 지원: FSDP 또는 DeepSpeed를 통해 단일 GPU 또는 다중 GPU 훈련 지원
• 효율적인 훈련: NVIDIA GPU(Ampere 이상, bf16 및 Flash Attention 지원) 및 AMD GPU에 최적화

배포 및 모니터링

• 클라우드 준비: Docker 이미지 및 PyPI 패키지 제공, 클라우드 플랫폼 및 로컬 하드웨어에서 사용 가능
• 결과 기록: 결과 및 체크포인트를 WandB, MLflow 또는 Comet에 기록 지원
• 모니터링 지원: 다양한 실험 추적 및 모니터링 도구 통합

기술 요구 사항

하드웨어 요구 사항

• NVIDIA GPU(Ampere 이상, bf16 및 Flash Attention용) 또는 AMD GPU
• 모델 훈련에 충분한 GPU 메모리

소프트웨어 요구 사항

• Python 3.11
• PyTorch ≥2.4.1
• 관련 종속성 패키지

설치 방법

빠른 설치

pip3 install -U packaging==23.2 setuptools==75.8.0 wheel ninja
pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]

# 예제 구성 파일 다운로드
axolotl fetch examples
axolotl fetch deepspeed_configs  # 선택 사항

소스에서 설치

git clone https://github.com/axolotl-ai-cloud/axolotl.git
cd axolotl
pip3 install -U packaging setuptools wheel ninja
pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'

Docker 방식

docker run --gpus '"all"' --rm -it axolotlai/axolotl:main-latest

사용 방법

기본 사용 흐름

1. 예제 구성 가져오기:

   axolotl fetch examples
   

2. 모델 훈련:

   axolotl train examples/llama-3/lora-1b.yml
   

3. 사용자 정의 구성: 필요에 따라 YAML 구성 파일의 파라미터 수정

구성 파일 구조

Axolotl은 YAML 구성 파일을 사용하여 모델 선택 및 파라미터, 데이터 세트 구성 및 전처리, 훈련 하이퍼파라미터, 최적화 프로그램 설정, 모니터링 및 로깅을 포함한 전체 훈련 흐름을 제어합니다.

지원 모델 호환성 매트릭스

모델	fp16/fp32	LoRA	QLoRA	GPTQ	Flash Attn	xformers
LLaMA	✅	✅	✅	✅	✅	✅
Mistral	✅	✅	✅	✅	✅	✅
Mixtral-MoE	✅	✅	✅	❓	❓	❓
Pythia	✅	✅	✅	❌	❌	❌
Falcon	✅	✅	✅	❌	❌	❌
Qwen	✅	✅	✅	❓	❓	❓
Gemma	✅	✅	✅	❓	❓	✅

✅: 지원 ❌: 미지원 ❓: 테스트되지 않음

응용 시나리오

연구 분야

• 대규모 언어 모델 미세 조정 연구
• 파라미터 효율적 훈련 방법 실험
• 모델 정렬 및 보안 연구

산업 응용

• 기업 수준 모델 맞춤화
• 특정 분야 모델 훈련
• 제품 기능 모델 최적화

교육 훈련

• AI/ML 과정 교육
• 실습 프로젝트 개발
• 기술 능력 향상

프로젝트 장점

1. 사용 편의성: 간단한 YAML 구성 파일을 통해 복잡한 훈련 흐름 제어
2. 유연성: 다양한 모델 아키텍처 및 훈련 방법 지원
3. 성능: 최신 최적화 기술 통합, 효율적인 훈련 경험 제공
4. 확장성: 단일 GPU에서 다중 노드까지 다양한 훈련 규모 지원
5. 오픈 소스: Apache 2.0 라이선스, 완전 오픈 소스 무료 사용