用ColossalAI完成一次完整的预训练

太难了，累懵了，全是坑...

      最近没更新，其实有机会（怎么个机会不细说了）可以玩玩两台新出炉的H100，而且是8卡400G IB的，这两台估计已经超过了库里南的价格了, 极其的豪华...

      因为我正好没看《乡村爱情15》，我买了个youku会员，可以边看《乡村爱情15》边拿H100跑一跑训练，看看具体怎么个实力。

      但是时间有限啊，主要是成本，那好钢要用在刀刃上，所以我在考虑用什么来测试它, 我没玩过megatron，deepspeed又弄过了，所以我挑战了一下Colossal，然后就被一顿教育...

点，可是这代码质量和维护项目的速度真的是...

      具体的细节我后面会再写一篇文章写，也不是单独针对Colossal的，因为像分布式训练的一些基础，我觉得非常有必要深入浅出的多写几篇文章才能说明白，估计TP，PP，DDP，Gemini，Zero，1D,  2D, 2.5D，3D,每个概念都够水一篇文章的了

      有道是光说不练，假把式，所以今天呢，先从一个demo做起，让大家先看一下它是怎么个操作流程和玩法（我改动的一些原始代码就算了，第一太多，第二这一个礼拜改的太多，我都快记不住哪是哪了

，我也几乎每个发现的问题都报了bug，希望后面能fix） 

言归正传，项目地址：hpcaitech/ColossalAI at v0.3.4 (github.com)

       下载了之后项目的目录结构是这样的

       为了简单测试，我就选择了Llama2的模型，在上面的红框部分，有两个Llama2的Demo，今天我以application中的demo为案例演示一下

首先安装需要的库

cd Colossal-LLaMA-2

pip install -r requirements.txt

在flash-attention的目录下（这个要去git自己下载，然后到了项目目录下，才能安装以下几个组件，pip install flash-attn是不行的）

# At the root folder

cd csrc/xentropy && pip install .

# At the root folder

cd csrc/layer_norm && pip install .

# At the root folder

cd csrc/rotary && pip install .

pretrain第一步：

扩充词表，崭新的分词器，这一步的意义是引入原模型无法识别的词汇，如比较复杂的中文词组。

词表的格式是是这样的：

{"piece": "你好"}

{"piece": "人工智能"}

export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION='python'

python colossal_llama2/tokenizer/init_tokenizer.py \

    --source_tokenizer_dir "/aml2/llama2" \

    --target_tokenizer_dir "/aml2/llam2/newtoken" \

    --expand_tokens_file "/aml2/llama2/jsl/3.jsonl"

pretrain第二步：

加载预训练模型：

我没有那么多词表，弄起来也麻烦所以我直接用Alpaca的词表来当构建的新词表

这一步的目的是把新词表和原始模型做整合

python colossal_llama2/model/init_model.py \

    --source_model_and_tokenizer_path "/aml2/llama2" \

    --target_tokenizer_path "/aml/chinese-alpaca-2-7b" \

    --target_model_path "/aml2/llama2/model"

Alpaca词表搞得有点大了，1分钟500个input，有20000多增量，太慢了(失误)...

pretrain第三步：

整理数据格式，符合训练要求。

大概可以整理成这样

{"source": "", "target": "Lionel Andrés Messi(Spanish pronunciation: [ljoˈnel anˈdɾes ˈmesi] (i); born 24 June 1987), also known as Leo Messi, is an Argentine professional footballer who plays as a forward for and captains both Major League Soccer club Inter Miami and the Argentina national team.", "category": "sports"}

{"source": "猜谜语：一身卷卷细毛，吃的青青野草，过了数九寒冬，无私献出白毛。（打一动物）", "target": "白羊", "category": "riddle"}

当然它的dataset格式不一定符合现有数据的要求，可以改代码部分（我就改了），也可以用python调整一下原始数据的catalog

一切就绪就开启训练模式

开启训练模式

具体的脚本大概长这样

基本就还是那些配置

• 基础模型（就是参数）
• parallel的模式 Zero，3D之类的
• 精度
• 梯度accumulate，checking之类的

这样就可以训练属于你自己的Llama预训练模型了

总结：Colossal真的是feature满满，但是不够成熟，你得有耐心和手法去折腾，另外对环境要求还是比较古老，如果这些不能接受，还是选Deepspeed吧！

Tips：不支持Torch2.0及以上，我把代码迁移到H100上以后....  

  warnings.warn(incompatible_device_warn.format(device_name, capability, " ".join(arch_list), device_name))

/aml/colo/lib/python3.10/site-packages/colossalai/kernel/cuda_native/mha/flash_attn_2.py:28: UserWarning: please install flash_attn from https://github.com/HazyResearch/flash-attention

  warnings.warn("please install flash_attn from https://github.com/HazyResearch/flash-attention")

/aml/colo/lib/python3.10/site-packages/torch/cuda/__init__.py:155: UserWarning: 

NVIDIA H100 80GB HBM3 with CUDA capability sm_90 is not compatible with the current PyTorch installation.

The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_61 sm_70 sm_75 sm_80 sm_86 compute_37.

If you want to use the NVIDIA H100 80GB HBM3 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

先是Torch 根本就发挥不出H100 高级算子的实力，后来NCCL也来捣乱，索性就放弃了...

下一篇可以写SFT