Axolotl: 一款强大的AI模型微调工具

Axolotl简介

Axolotl是一款强大的AI模型微调工具,为多种AI模型的微调提供了统一的仓库和接口。它的名字源自墨西哥钝口螈,这种两栖动物以其强大的再生能力而闻名,恰如其分地体现了该工具在AI模型"再生"(即微调)方面的卓越能力。

墨西哥钝口螈

Axolotl的主要特点包括:

支持多种Hugging Face模型,如llama、pythia、falcon、mpt等
支持全参数微调、LoRA、QLoRA、ReLoRA和GPTQ等多种微调方法
可通过简单的YAML文件或命令行参数自定义配置
支持加载不同格式的数据集,可使用自定义格式或预先分词的数据集
集成了xformer、flash attention、liger kernel、rope scaling和multipacking等先进技术
支持单GPU或多GPU训练,兼容FSDP和Deepspeed
可在本地或云端轻松使用Docker运行
支持将结果和检查点记录到wandb或mlflow

快速开始

要开始使用Axolotl,只需几个简单的步骤:

git clone https://github.com/axolotl-ai-cloud/axolotl
cd axolotl

安装依赖:

pip3 install packaging ninja
pip3 install -e '.[flash-attn,deepspeed]'

预处理数据集(可选但推荐):

CUDA_VISIBLE_DEVICES="" python -m axolotl.cli.preprocess examples/openllama-3b/lora.yml

开始微调:

accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml

推理:

accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
    --lora_model_dir="./outputs/lora-out"

使用Gradio进行交互式推理:

accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
    --lora_model_dir="./outputs/lora-out" --gradio

高级设置

环境配置

Axolotl提供了多种环境配置选项,以适应不同用户的需求:

Docker: 提供了预配置的Docker镜像,可以快速启动训练环境。
Conda/Pip虚拟环境: 适合喜欢自定义环境的用户。
云GPU: 支持在各种云GPU提供商上运行,如Latitude.sh、JarvisLabs.ai和RunPod等。
裸金属云GPU: 提供了在LambdaLabs和GCP等平台上的详细设置指南。
Windows和Mac: 分别提供了WSL和特定的安装说明。
Google Colab: 提供了专门的Notebook示例。
SkyPilot和dstack: 支持通过这些工具在公共云上快速部署。

数据集准备

Axolotl支持多种数据集格式,推荐使用JSONL格式。数据集的具体模式取决于任务类型和您选择的提示模板。除了JSONL,您还可以使用Hugging Face数据集,其中包含与JSONL字段对应的列。

配置文件

配置文件是Axolotl的核心,允许用户精细控制训练过程。主要配置项包括:

模型选择: 指定基础模型的本地路径或Hugging Face仓库。
数据集配置: 支持多种数据集源,包括Hugging Face仓库、本地文件、S3/GCS存储等。
加载选项: 控制模型加载方式,如4-bit/8-bit量化、混合精度训练等。
LoRA配置: 设置LoRA参数,如rank、alpha值、dropout率等。

Axolotl支持的模型

训练过程

Axolotl提供了灵活的训练选项:

单GPU训练: 使用简单的命令即可开始训练。
多GPU训练: 支持DeepSpeed和FSDP两种分布式训练方案。
- DeepSpeed: 提供了ZeRO stage 1、2、3的预配置JSON文件。
- FSDP: 支持全分片和自动包装等高级功能。
- FSDP + QLoRA: 结合量化和低秩适应,进一步优化内存使用。
权重与偏置(W&B)日志记录: 支持将训练日志上传到W&B平台,方便监控和分析。
特殊令牌处理: 可以自定义特殊令牌,以适应特定的任务需求。
Liger Kernel: 集成了这一高效的注意力机制实现。

推理与模型合并

训练完成后,Axolotl提供了方便的推理接口:

命令行推理: 可以快速测试模型性能。
Gradio界面: 提供交互式Web界面,方便演示和实验。
LoRA模型合并: 支持将训练好的LoRA权重合并回基础模型。

常见问题与调试

使用Axolotl过程中可能遇到的常见问题包括:

推理与训练的分词不匹配: 提供了解决方案和最佳实践。
内存相关问题: 给出了针对CUDA out of memory等错误的处理建议。
多GPU训练问题: 详细说明了使用DeepSpeed和FSDP时可能遇到的问题及解决方法。

Axolotl还提供了详细的调试指南,帮助用户解决训练过程中遇到的各种问题。

社区与贡献

Axolotl拥有活跃的社区支持:

社区展示: 展示了使用Axolotl训练的优秀模型。
贡献指南: 欢迎社区成员参与项目开发,提供了详细的贡献流程。
赞助商: 列出了支持Axolotl发展的赞助商,分为钻石、金、银、铜四个级别。

Axolotl社区支持

结语

Axolotl作为一个强大而灵活的AI模型微调工具,为研究人员和开发者提供了一个统一的平台,简化了模型微调的复杂性。无论是对于初学者还是经验丰富的AI从业者,Axolotl都提供了丰富的功能和清晰的文档支持,使得AI模型的定制化变得更加容易和高效。随着AI技术的不断发展,Axolotl将继续发挥其"再生"能力,助力更多创新模型的诞生和优化。