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许可证，完全开源免费使用