Pints.ai -- 9天内训练出强大的小型语言模型

1.5-Pints

9天内预训练模型的方法，可与苹果OpenELM和微软Phi等人工智能助手相媲美。

本仓库包含由Pints.AI开发的1.5-Pints和0.12-Pint的模型架构、训练脚本和实用工具。通过提供模型的代码库和架构访问，本计划旨在促进Pint的复制、实验和进一步的开源开发。

加入我们的Discord: https://discord.gg/eGTRzDdH

论文与引用

@misc{tan202415pintstechnicalreportpretraining,
      title={1.5-Pints技术报告：数天而非数月的预训练 -- 您的语言模型依赖于高质量数据}, 
      author={Calvin Tan and Jerome Wang},
      year={2024},
      eprint={2408.03506},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2408.03506}, 
}

安装

安装conda

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
sh Miniconda3-latest-Linux-x86_64.sh

为确保conda命令可用，请执行：

source ~/.bashrc

如果仍然遇到conda: command cannot be found，可以找到安装位置并执行：

注意：此路径假设您使用了默认安装设置。否则，请找到您的安装位置。

source ~/miniconda3/etc/profile.d/conda.sh

克隆此仓库

git clone https://github.com/Pints-App/Pints-Train.git && \
cd Pints-Train

创建conda环境

conda create --prefix ./.conda python=3.10 && \
conda activate ./.conda

注意：请使用Python 3.10。截至目前（2024年2月23日），3.12版本会导致许多问题，而3.11版本尚未经过测试。

安装CUDA工具包

conda install nvidia/label/cuda-12.1.1::cuda-toolkit

安装依赖

pip install -r requirements.txt && \
pip install flash-attn --no-build-isolation && \
pip install -r pretrain/requirements.txt

注意：dropout_layer_norm的pip安装可能需要约30分钟的编译时间，具体取决于机器性能。

如果尚未安装git-lfs，请安装

curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get install git-lfs

下载pints-ai Expository Prose预训练数据集

cd /path/to/dataset_dir
git clone https://huggingface.co/datasets/pints-ai/Expository-Prose-V1

准备数据集

python -m prepare_dataset.standard_parquet \
--source_path /path/to/dataset_dir  \
--train_val_split_ratio 0.9 \
--max_cores 60 \
--destination_path /path/to/output_dir

请参考prepare_dataset文件夹获取数据集准备脚本。

如果在高核心机器上不会出现内存溢出，则不需要设置max_cores。

训练模型

Lightning

fabric run model \
--accelerator=cuda \
--devices=8 \
pretrain/main.py \
--data_dir data/output \
--out_dir ../1.5-pints \
--gpus 8 \
--global_batch_size 512 \
--learning_rate 4e-4 \
--micro_batch_size 8 \
--max_step 96180 \
--warmup_steps 2000 \
--weight_decay 0.1 \
--beta1 0.9 \
--beta2 0.95 \
--grad_clip 1.0  \
--min_lr 4e-5 \
--model_name 1.5-Pints-2k \
--wandb_name <run_name> \
--wandb_project <project_name> \
--tokenizer_dir tokenizer/pints

注意1：--devices和--gpus必须相同。有关所有可调整的参数，请参见pretrain.py的setup参数。

注意2：使用--model_name选择架构（层数/维度/头数）配置。这必须在lit_gpt/config.py中定义。

注意3：选择micro_batch_size以优化GPU内存使用。到目前为止，一旦开始训练，即使在验证期间也能保持稳定。micro_batch_size不需要是batch_size的整除数。batch_size由global_batch_size / devices得出。

注意4：修改pretrain/main.py中的TRAIN_DATA_CONFIG以决定用于训练的数据集。确保事先已准备好数据集。

Wandb

如果要求提供wandb API密钥，可以登录并从以下地址获取：https://wandb.ai/authorize

微调模型

cd finetune && \
pip install -r requirements.txt

转换预训练权重

开始微调之前，需要转换预训练权重：

python convert/convert_pretrained_checkpoint.py --checkpoint_dir path/to/checkpoint --output_dir path/to/output

lightning run model \
--accelerator=cuda \
--devices=8 \
finetune/full.py \
--checkpoint_dir <path to lit_model.pth> \
--out_dir ~/1.5-pints-2k/ep2/step-00045000/finetuned \
--model_name 1.5-Pints-2k \
--gpus 8 \
--train.save_interval 6000 \
--train.global_batch_size 512 \
--train.micro_batch_size 8 \
--train.lr_warmup_steps 1125 \
--train.epoch 5 \
--train.learning_rate 2e-5 \
--train.max_seq_length 2048 \
--train.beta1 0.9 \
--train.beta2 0.95 \
--train.weight_decay 0.1 \
--logger_name wandb \
--tokenizer_dir tokenizer/pints \
--known_data_max_seq_length 2048 \
--wandb_project <project name>

运行直接偏好优化（DPO）

DPO用于微调后的使用。有关执行过程，请参见此处。

评估模型

请参见此处

使用模型

将lit模型转换为HuggingFace模型（pytorch和safetensors）

python convert_lit_to_hf.py \
--checkpoint_name lit_model.pth \
--directory ../models/1.5-pints \
--model_name 1.5-Pints-2k \
--output_config=True \
--safetensors=True \
--delete_pytorch_model=True

注意：我们发现使用safetensors文件效果更好。因此，建议使用它而不是pytorch_model.bin。

使用huggingface

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("/model/path")
model = AutoModelForCausalLM.from_pretrained("/model/path")

prompt = '''<|im_start|>user
Do not go gentle into that good night.<|im_end|>
<|im_start|>assistant
'''

tokenized_input = tokenizer.encode(prompt)
tokenized_output = model.generate(tokenized_input)
print(tokenizer.decode(tokenized_output))