前言
最近好像各家都拿出来了自己的大模型,有Finetune开源模型,也有从头开始预训练的模型. 我所在的课题组也Pretrain了一个小一点的LLM(1.5B和500M两个大小). 刚开始训练的时候相关的教程和资料还都比较少,后来相关的博客和Codebase也慢慢增加了一些,让我们收益匪浅. 现在模型训练已经基本上到了比较稳定的时候,也在此分享下自己的一些小小经验.
首先呈上我们训练的硬件的相关参数.
- 我们集群内的卡是
NVIDIA A100 PCIe 80GB,无NVLink和NVSwitch,集群间互联靠基于Infiniband的 RDMA. - 每个节点通过NFS挂载存储节点来实现共享存储.
- CPU 为8358P 64线程,内存1T.
通过上面的介绍也可以看出来,我们的节点基本上是比较乞丐版的配置,特别是没有NVLink卡间互联导致基本告别Tensor Parallelism和Model Parallelism的训练策略,从而只能选择单张卡能放下的模型(3B以下),这很大程度影响了我们下面对模型参数量和结构的决策.
LLM的模型结构重要吗?
有位大佬曾经这样形容: “LLM的效果中数据集占60%,算力(训练Token数)占30%,模型结构占10%”. 我个人的观点是目前Transformer架构的参数量可能有很大的冗余,因此一些架构上的小修改可能对最终的效果没有很大的影响. 但是除开LLM Benchmark的分数,我们也关注推理和训练成本的优化,而一些架构改进可能能更高效的进行训练和推理.下面是我们做的几个尝试(踩的几个坑)
- VocabSize : 我们的实践来看,VocabSize的增加是少数真的能提高效果的Tricks. 主流的Tokenizer对于不存在词表中的未知字词采用BPE byte fallback 的方式.这种方式会把字词拆分成UTF-8的表示并分别转换为Token. 这给后续的LLM理解毫无疑问增加了很大难度,因此小词表的LLM 如LLaMA 1等在多语言任务上表现比较差.此外, 很多词表会包含常见词组的组合,这有助于模型更好更完整理解词组的含义(当然也降低了上下文的Token数目). 我们尝试过3w2(LLaMA), 5W(Huggingface), 12w(OpenBMB和LLaMA3.1) , 15w(Qwen 2)等词表.大致的趋势是在偏重语言理解和Commonsense Reasoning的任务(Arc,Winogrand等)上,分数会随词表数目数目增加. 当然大词表会带来显存占用的上涨,导致训练batch降低.
- FFN : 测试发现增大HiddenSize和IntermediateSize,把模型做宽能提升推理性能(特别是在移动端).但是同参数量下可能层数就要相应减少.这是个需要考量Tradeoff.
- GQA & MLA: 这两种都是面向减小KVCache方向而实现的架构改进.毕竟在Memory Bound的Decoding阶段,如果能减小访存开销,就能有效提升推理速度.更激进的架构如MLA甚至可以把访存的瓶颈提升到GPU L2 Cache甚至寄存器Cache.但是我们最终只采用了GQA,一方面是尝试训练了一个MLA的模型发现效果有比较明显的损失,另一方面感觉MLA的很多配套设施不太完善,官方的初版代码甚至实现有错误,还是选择了较为保守的GQA.
以上是我们尝试过有益的一部分模型结构,因为我不是算法方向的同学,可能给出的相关信息有点疏漏,仅供参考. 我们其实还根据相关论文或开源模型尝试过一些较为奇怪的模型:
- 每一层的QKV Head Num和宽度是变化的,呈前小后大的形状(OpenELM). 官方给出的原因是考虑到每一层的包含是信息量是不一样的,这样有利于将参数留给更重要的层.我个人感觉有点过度设计了,而且因为每层的Hidden Size都不一样,给VLLM等推理框架的兼容造成了很大难度,KVCache的池化共享也被限制在了每层之内.感觉没有必要,所以没有采用.
- 每两层Transformer Block的权重是共享的,使用同一个Weights Tensor(MobiLLaMA).确实能节省参数量,但是没有观察到特别好的效果.而且最重要的是和框架代码有冲突,DeepSpeed的Zero2并行代码assert了每个权重只能被反向传播一次.
训练的过程
开始之前先Profile
在模型训练之前,对代码和设备进行Profile其实是很有必要的.如果是从零开始组建的集群,可以跑下Nvidia官方的Benchamrk,如nvbandwidth,cuda-samples和nccl-tests来校验配置是否正确. 此外,在启动训练之前,应当分别在单机多卡和多级多卡上使用Torch Profilor启动一次测量,分析出来代码中的耗时情况.对于大部分多机并行的模型训练而言, NCCL的通信时长近似甚至多于正反向传播时长是很正常的.流水线和框架层面一些Overlap操作可以尽可能掩盖通信开销,但是GPU的气泡空转仍然不可避免.
当你发现一次Step中通信占据了2/3以上,那么继续增加节点数目可能就不太是个好主意了.所以NVLink等加速通信的配置还是能买就要买.或者尽可能把训练框架更新到最新的版本,应用TreeAllReduce或者RingAllReduce等尽可能节省通信量的技术.
代码编写
最近也有很多LLM开放了自己的训练代码,也有一些基于DeepSpeed或者Megatron的分布式训练框架,比如https://github.com/InternLM/InternEvo/ 或者Apple 的https://github.com/apple/corenet(很烂,建议避雷) ,或者也可以选择用Huggingface的Trainer和Accelerate自己封装一套训练代码. 此外我们也定制了几个callback来将训练指标上传到wandb和将训练checkpoint上传到OSS.
是否要提前Tokenize数据?
取决于训练框架的设计,如果Training Dataloader能比较好的与训练过程异步并行,那么Tokenize的开销会被卡间通信Overlap.具体的情形可以通过一次Profile来详细查看.
是否要开启激活值重计算?
根据我们的实验,激活值重计算大概会增加20%的时间开销,但是激活值释放可以换来更大的Batchsize.当显存不足或者需要增加Batchsize时,激活值重计算要优于开启Zero3.
数据集处理
我们一开始选择了手写一个读取Tokenize好的数据的IterableDataloader.事实证明这个选择不太明智.
- 使用迭代器来滚数据虽然看上去实现简单,但是当训练挂掉Resume的时候,我们不得不重新开始滚到数据一直到指定的TrainingState.这个过程非常漫长(几小时到1天),而且训练进程很可能在这过程中出现NCCL或者网络互联故障. 因此我强烈建议实现一个可寻址的(即实现__getitem__魔法方法的)数据集.能大大提高Resume和训练的效率. 如果你懒得实现,参考https://github.com/InternLM/InternEvo/blob/develop/internlm/data/tokenized/packed_dataset.py
- 除了之前提到的Tokenized Dataset会占用更多存储空间之外,遵循SIMT并行模型的训练框架会在一个节点上启动卡数对应的子进程同时读取数据,在我们某几台节点上这样会导致NFS的IO负担过重,从而导使用致mmap方式映射到内存的数据集文件在读取时出现Segfault,直接带崩整个训练.如果你的节点内存比较富余,可以尝试将只读的数据实现拷贝到挂载内存盘上,能有效提高IO性能和稳定性.
训练中断的原因Top5
在LLaMA3 的技术报告中,Meta给出了大规模训练中断的归因分析. 不同与Meta,我们训练中断的原因可能更多是代码写错了(特别是Eval部分或Callback的代码),存储空间占满了(要及时清理log和ckpt),NCCL Timeout(不在同一个交换机下面的几组节点偶尔会出现这种问题,但是常见在训练开始的时候)以及训崩了.
当你发现Train Loss和Eval Loss(如果有)都开始上升而且似乎没有下降的势头并持续了好几百Steps.那你的模型可能训崩了,走入了错误的局部优解.这时候你可以选择杀掉训练,回滚到还没有走下坡路的Checkpoint,并且跳过接下来的几百Steps对应的数据,祈祷也许有好转. 当然如果你没来得及做这个事情,你也可以洗洗睡觉.有时候Loss也会慢慢涨回来XD,毕竟模型不会总是记住所有的数据.