CoLLiE (Collaborative Tuning of Large Language Models in an Efficient Way)，一个帮助您从零开始训练大模型的完整工具箱。
- [2023/12] 🎉 CoLLiE被EMNLP System Demonstrations接收：CoLLiE: Collaborative Training of Large Language Models in an Efficient Way
- [2023/08] 评测结果新增显存占用与模型大小的关系和吞吐量。
- [2023/07] 发布Python包
collie-lm。您可以在PyPI中查看更多细节！
CoLLiE是一个可以帮助您从零开始训练大模型的完整工具箱，它提供了数据预处理、模型微调、模型保存以及训练过程各项指标监测等功能。CoLLiE集成了现有的并行策略、高效参数微调方法和高效优化器，以加快训练的速度，提高训练的质量，降低训练的开销。CoLLiE支持主流的多种模型（如MOSS, InternLM, LLaMA, ChatGLM等），您可以轻松在不同的模型之间切换。此外，CoLLiE提供了丰富的文档，使初学者可以快速入门。同时，CoLLiE还提供了高度可定制化的功能和灵活的配置选项，使有经验的用户能够根据自己的需求进行个性化定制。无论您是初学者还是有经验的专业人士，CoLLiE都可以为您提供满足需求的解决方案。
CoLLiE 基于 DeepSpeed 和 PyTorch，为大型语言模型提供协作式和高效的调优方法。 它主要包括以下四个特点：
- 并行策略
- 数据并行 (DP)
- 流水线并行 (PP)
- 张量并行 (TP)
- 零冗余优化器 (ZeRO)
- 高效微调
- 设计优雅
- 用户友好
完整特性
使用张量并行测试了批量大小为 1，序列长度为 2048，梯度累计步数为 2 下显存占用情况，结果如下：
在 A100 和 RTX-3090 上测试了不同批量大小下使用 Adam 优化器的吞吐量，结果如下：
在安装前，你需要确保：
- PyTorch >= 1.13
- CUDA >= 11.6
- Linux OS
你可以简单地通过PyPI安装，命令如下：
pip install collie-lmgit clone https://github.com/OpenLMLab/collie
python setup.py install下面将提供一个使用CoLLiE训练Moss的样例，同时使用LOMO优化器，并且开启ZeRO3来降低显存消耗。
那么，请按照下面的步骤开启你的大模型训练之旅吧~
from transformers import AutoTokenizer
from collie.config import CollieConfig
from collie.data import CollieDatasetForTraining
from collie.data import CollieDataLoader
from collie.optim.lomo import Lomo
from collie.controller.trainer import Trainer
from collie.controller.evaluator import EvaluatorForPerplexity, EvaluatorForGeneration
from collie.models.moss_moon import Moss003MoonForCausalLM
from collie.utils.monitor import StepTimeMonitor, TGSMonitor, MemoryMonitor, LossMonitor, EvalMonitor
from collie.metrics import DecodeMetric, PPLMetric
from collie.module import GPTLMLoss
from collie.utils.data_provider import GradioProvider选择预训练模型为MOSS
pretrained_model = "fnlp/moss-moon-003-sft"
config = CollieConfig.from_pretrained(pretrained_model, trust_remote_code=True)
# 张量并行
config.tp_size = 2
# 数据并行
config.dp_size = 1
# 流水线并行
config.pp_size = 1
# 训练的epoch数量
config.train_epochs = 1
# 每{100}个step进行一次eval
config.eval_per_n_steps = 100
# 每{1}个epoch进行一次eval
config.eval_per_n_epochs = 1
# 每个GPU的batch_size设置为{16}
config.train_micro_batch_size = 16
# 每次eval的batch_size为{1}
config.eval_batch_size = 1
# 设置DeepSpeed配置
config.ds_config = {
# 开启FP16
"fp16": {
"enabled": True
},
"zero_allow_untested_optimizer": True,
"zero_force_ds_cpu_optimizer": False,
# 开启ZeRO-3
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": False
}
},
"monitor_config": {
"enabled": True,
"tag": "adan",
"csv_monitor": {
"enabled": True,
"output_path": "./ds_logs/"
}
}
}tokenizer = AutoTokenizer.from_pretrained("fnlp/moss-moon-003-sft", trust_remote_code=True)这里自定义一个数据集，数据格式可以提供两种形式，具体请参考文档。
train_dataset = [
{
'input': 'Collie is a python package for ',
'output': 'finetuning large language models.'
} for _ in range(10000)
]
train_dataset = CollieDatasetForTraining(train_dataset, tokenizer)
eval_dataset = train_dataset[:32]model = Moss003MoonForCausalLM.from_pretrained(pretrained_model, config=config)optimizer = Lomo(
model,
lr = 0.001,
clip_grad_norm = 5.0
)monitors = [
# 每个step用时监测
StepTimeMonitor(config),
# TGS（每秒生成token数量监测）
TGSMonitor(config),
# 显存使用情况监测
MemoryMonitor(config),
# Loss值监测
LossMonitor(config),
# Eval结果监测
EvalMonitor(config)
]这里添加两个Evaluator，分别用于计算PPL(困惑度：Perplexity)和保存Decode结果。
evaluator_ppl = EvaluatorForPerplexity(
model = model,
config = config,
dataset = eval_dataset,
monitors = [
EvalMonitor(config)
],
metrics = {
'ppl': PPLMetric()
}
)
evaluator_decode = EvaluatorForGeneration(
model = model,
config = config,
tokenizer = tokenizer,
dataset = eval_dataset,
monitors = [
EvalMonitor(config)
],
metrics = {
'decode': DecodeMetric()
}
)trainer = Trainer(
model = model,
config = config,
loss_fn = GPTLMLoss(-100),
optimizer = optimizer,
train_dataset = train_dataset,
monitors = monitors,
evaluators = [evaluator_ppl, evaluator_decode],
)
# 开始训练/验证
trainer.train()Command CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --rdzv_backend=c10d --rdzv_endpoint=localhost:29402 --nnodes=1 --nproc_per_node=4 finetune_moss_for_training.py如果你的命令行出现如下的进度条，那么恭喜你，你已经成功开始训练你的大模型！
完整代码请参考examples/finetune_moss_for_training.py。
CoLLiE提供了许多即插即用的插件，下面将介绍Monitor检测器和异步DataProvider，更多插件等待探索和开发...
在CollieConfig.ds_config中添加monitor配置，并在Trainer中启用即可在训练过程中打开监测器。
"monitor_config": {
# 开启检测器
"enabled": True,
# 保存的文件名前缀
"tag": "adan",
# 保存文件格式:csv
"csv_monitor": {
"enabled": True,
# 保存文件夹
"output_path": "./ds_logs/"
}
}启用检测器后，你将在ds_logs文件夹中获取相关的文件，如：
你只需要在Trainer中添加：data_provider，即可在训练过程中打开一个异步DataProvider，方便及时Human Eval！
trainer = Trainer(
model = model,
config = config,
loss_fn = GPTLMLoss(-100),
optimizer = optimizer,
train_dataset = train_dataset,
monitors = monitors,
evaluators = [evaluator_ppl, evaluator_decode],
# 添加
data_provider = GradioProvider(tokenizer)
)
CoLLiE提供了完整的 教程。 更多的示例也可在 示例 中查看。