Axolotl - 強大的AI模型後訓練工具

項目概述

Axolotl是一個專門設計用於簡化各種AI模型後訓練過程的開源工具。後訓練是指對預訓練模型進行的任何修改或額外訓練，包括全模型微調、參數高效調優（如LoRA和QLoRA）、監督微調（SFT）、指令調優和對齊技術。該工具支持多種模型架構和訓練配置，讓用戶能夠輕鬆開始使用這些先進的訓練技術。

核心特性

模型支持

• 多樣化模型架構: 支持訓練各種Hugging Face模型，包括LLaMA、Pythia、Falcon、MPT、Mistral、Mixtral等主流大語言模型
• 靈活訓練方法: 支持全量微調（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許可證，完全開源免費使用