LitGPT: 高性能大语言模型的预训练、微调和部署工具

litgpt

LitGPT:高性能大语言模型的预训练、微调和部署工具

LitGPT是一个功能强大的命令行工具,专门用于对大语言模型(LLM)进行预训练、微调、评估和部署。它支持20多种流行的LLM架构,并采用了最先进的优化技术,使其成为AI研究人员和开发者的理想选择。

主要特点

LitGPT具有以下几个突出特点:

从头实现多种模型: LitGPT对20多种流行的LLM架构进行了从头实现,包括Llama 3、Code Llama、Mistral、Phi等。这种实现方式去除了抽象层,使代码更加简洁高效。
先进的优化技术: 采用Flash Attention v2、全分片数据并行(FSDP)等技术,支持多GPU训练,可选的CPU卸载,以及TPU和XLA支持。
低精度和量化: 支持FP16、BF16等低精度设置,以及4位浮点、8位整数等量化方法,以降低计算和内存需求。
参数高效微调: 支持LoRA、QLoRA、Adapter等参数高效的微调方法。
丰富的训练配方: 提供经过验证的YAML配置文件,可用于不同条件下的模型训练。
企业级应用: 采用Apache 2.0许可,适合企业无限制使用。代码易于调试,单文件实现方便修改。

快速上手

安装LitGPT非常简单:

pip install 'litgpt[all]'

然后就可以使用任何支持的LLM模型:

from litgpt import LLM

llm = LLM.load("microsoft/phi-2")
text = llm.generate("Fix the spelling: Every fall, the familly goes to the mountains.")
print(text)
# 输出: Every fall, the family goes to the mountains.

主要工作流程

LitGPT支持以下主要工作流程:

微调: 在预训练模型的基础上,使用特定任务的小型数据集进行进一步训练。
预训练: 从头开始训练模型,使其学习语言的基本结构和知识。
继续预训练: 在已有预训练模型的基础上,使用自定义数据继续训练。
评估: 测试模型在各种任务上的表现。
部署: 将训练好的模型部署到实际应用中。

以下是一个微调模型的简单示例:

# 下载预训练模型
litgpt download microsoft/phi-2

# 微调模型
litgpt finetune microsoft/phi-2 \
  --data JSON \
  --data.json_path my_custom_dataset.json \
  --data.val_split_fraction 0.1 \
  --out_dir out/custom-model

# 测试模型
litgpt chat out/custom-model/final

# 部署模型
litgpt serve out/custom-model/final