1. 项目概述Sheared LLaMA——一种高效的大语言模型预训练加速方法如果你正在为训练一个性能优秀的小规模语言模型比如1.3B或2.7B参数而发愁觉得从头开始预训练成本太高、周期太长那么来自普林斯顿NLP团队的Sheared LLaMA项目或许能给你提供一个全新的、极具性价比的思路。这个项目的核心思想非常直接与其从零开始训练一个小模型不如从一个已经训练好的、性能强大的大模型比如LLaMA-2-7B出发通过一种结构化剪枝技术精准地“裁剪”掉模型中不那么重要的部分得到一个更小、更紧凑的架构然后再用少量数据对这个裁剪后的模型进行持续预训练使其恢复甚至超越同等规模从头训练模型的性能。简单来说Sheared LLaMA不是“造一个小引擎”而是“从一个优秀的大引擎上拆解出核心部件重新组装成一个更高效的小引擎”。根据论文中的对比通过剪枝LLaMA-2-7B得到的2.7B模型其性能可以达到与从头预训练的OpenLLaMA-2.7B相当的水平而所需的计算成本仅为后者的约3%。这个数字对于资源有限的研究者、开发者和小型团队来说吸引力是巨大的。这个项目开源了完整的代码库涵盖了从数据准备、模型剪枝到持续预训练的全流程。接下来我将结合自己的实践和理解为你深入拆解这套方案的原理、实操步骤以及其中的关键细节。2. 核心原理为什么结构化剪枝持续预训练如此有效要理解Sheared LLaMA我们需要先打破一个常见的思维定式模型越小训练越快越便宜。这当然没错但忽略了“知识迁移”的成本。一个7B参数的大模型已经耗费了海量计算资源2T tokens学习了丰富的语言知识和世界知识。Sheared LLaMA的核心洞见在于大模型中蕴含的知识密度分布是不均匀的有些参数和结构对最终性能的贡献远大于其他部分。2.1 结构化剪枝与传统的非结构化剪枝模型剪枝并非新概念但传统方法多是非结构化剪枝即直接寻找并置零网络中不重要的单个权重。这种方法虽然能减少参数总量但产生的稀疏矩阵模式不规则在通用的硬件如GPU上无法获得实际的推理加速需要特殊的稀疏计算库支持实用性大打折扣。Sheared LLaMA采用的是结构化剪枝。它不是在单个权重级别操作而是在更高的结构维度上进行裁剪例如注意力头剪枝移除Transformer中某个注意力层的部分注意力头。前馈网络中间维度剪枝减少前馈网络中隐藏层的维度。隐藏维度剪枝缩减每层Transformer的嵌入维度。层剪枝直接移除整个Transformer层。这种做法的好处是裁剪后得到的模型仍然是一个稠密、规整的模型架构可以直接被PyTorch、Hugging Face等标准框架加载和高效运行立即获得内存和计算上的收益。2.2 基于L0正则化的可微分剪枝那么如何决定剪掉哪些结构呢Sheared LLaMA使用了一种基于L0正则化的可微分剪枝方法。其核心思想是为模型中每个可剪枝的结构如一个注意力头引入一个连续的门控变量gating variable和一个拉格朗日乘子Lagrangian multiplier。门控变量可以理解为该结构的“重要性分数”是一个可训练的参数范围在0到1之间。在训练过程中通过优化不重要的结构的门控值会趋向于0。拉格朗日乘子用于控制模型的整体稀疏度即裁剪比例确保最终模型的大小符合我们预设的目标如从7B剪到2.7B。在训练时模型的前向传播会同时考虑原始权重和门控变量。损失函数不仅包含常规的语言建模损失让模型预测下一个词还包含一个L0正则项鼓励门控变量趋近于0或1即“开”或“关”。通过这种可微分的方式模型在学习任务的同时也“学会”了哪些部分对自己是冗余的。训练完成后我们将门控值接近0的结构及其对应权重直接移除就得到了一个更小的、结构化的目标模型。2.3 持续预训练恢复“手术”后的性能直接剪枝后的模型就像动了一次大手术虽然架构变小了但内部参数的协调性被破坏了性能通常会有一个显著的下降。因此持续预训练是关键的第二阶段。我们用新的、高质量的数据以较小的学习率对这个剪枝后的模型进行继续训练。这个过程有两个目的知识恢复与巩固让模型在新的、更紧凑的架构下重新调整和巩固从大模型中继承来的知识。适应新数据分布如果持续预训练使用的数据与原始大模型的训练数据有差异模型还能学习到新的知识。这个阶段的成本远低于从头预训练因为模型已经具备了良好的知识初始化只需要微调以适应新架构。3. 环境搭建与数据准备实操理解了原理我们来看如何动手实践。Sheared LLaMA的代码库基于MosaicML的Composer库构建这是一个专为大规模语言模型训练优化的框架。整个剪枝逻辑是通过Composer的回调函数实现的对训练流程的侵入性很小设计很优雅。3.1 系统环境与依赖安装首先确保你的环境有足够的GPU资源建议A100 80GB。项目对PyTorch和Flash Attention有特定版本要求。步骤一安装PyTorch与Flash Attention# 安装指定版本的PyTorchCUDA 11.8 pip install torch2.0.1cu118 torchvision0.15.2cu118 torchaudio2.0.2 --index-url https://download.pytorch.org/whl/cu118 # 安装Flash Attention 1.x版本注意暂不支持2.0 pip install flash-attn1.0.3.post注意Flash Attention 2目前不被支持强行使用可能需要手动修改模型文件容易出错。请务必使用1.x版本。步骤二安装项目依赖克隆项目代码后进入目录安装剩余依赖。cd LLM-Shearing pip install -r requirement.txt步骤三以可编辑模式安装llmshearing包这样做的好处是你可以在开发过程中直接修改本地代码而无需反复重新安装。pip install -e .3.2 数据准备使用MosaicML StreamingSheared LLaMA使用MosaicML的Streaming数据集库来处理海量预训练数据。它的优势在于可以将超大规模数据集以流式方式加载无需全部下载到本地磁盘特别适合多机多卡训练。项目的数据处理脚本在llmshearing/data/目录下。你需要将自己的原始文本数据如.jsonl格式每行一个包含text字段的JSON对象转换为Streaming支持的.mds格式。核心流程包括编写数据转换脚本参考示例将你的文本文件分片、序列化。创建索引Streaming会为数据创建索引支持随机访问。配置YAML在训练配置中指定本地.mds文件的路径。一个简化的数据目录结构可能如下所示your_data_dir/ ├── train/ │ ├── shard.00000.mds │ ├── shard.00001.mds │ └── index.json └── eval/ ├── shard.00000.mds └── index.json在论文中作者主要使用了RedPajama v1数据集。你也可以尝试其他高质量数据集如Dolma、RedPajama-v2等数据的质量直接影响最终模型的性能。3.3 基础模型准备Hugging Face模型转Composer格式Composer训练需要特定格式的模型权重文件。我们需要将Hugging Face上的预训练模型如meta-llama/Llama-2-7b-hf进行转换。# 定义变量 HF_MODEL_NAMEmeta-llama/Llama-2-7b-hf OUTPUT_PATH./models/Llama-2-7b-composer/state_dict.pt # 创建输出目录 mkdir -p $(dirname $OUTPUT_PATH) # 执行转换脚本 python3 -m llmshearing.utils.composer_to_hf save_hf_to_composer $HF_MODEL_NAME $OUTPUT_PATH转换完成后强烈建议测试一下转换是否正确确保Hugging Face模型和Composer模型输出一致MODEL_SIZE7B python3 -m llmshearing.utils.test_composer_hf_eq $HF_MODEL_NAME $OUTPUT_PATH $MODEL_SIZE实操心得这个转换脚本目前主要适配LLaMA/LLaMA2架构。如果你想对Mistral、CodeLlama等其他模型进行剪枝需要仔细检查并可能修改composer_to_hf.py中的键名映射逻辑。一个常见的坑是模型配置中的维度名称可能不同比如hidden_sizevsd_model。4. 核心流程详解剪枝与持续预训练一切准备就绪后我们就可以开始核心的剪枝和训练流程了。项目提供了清晰的脚本示例。4.1 剪枝阶段配置与执行剪枝的主要脚本是llmshearing/scripts/pruning.sh。你需要准备一个YAML配置文件来定义所有参数。关键配置项可以分为四大类1. 数据配置 (data_local)data_local: /path/to/your/streaming/data train_loader: dataset: split: null # 动态加载时无需指定由回调函数控制 eval_loader: dataset: split: eval # 评估集使用固定的‘eval’分片2. 基础训练配置这些配置与标准的Composer训练类似。max_duration: 3200ba # 剪枝阶段训练3200个batch save_interval: 3200ba # 保存检查点 t_warmup: 320ba # 学习率热身步数占总步数10% optimizer: lr: 1.0e-4 # 主模型参数的学习率 max_seq_len: 4096 # 序列长度 device_train_microbatch_size: 4 # 每个GPU的batch大小 global_train_batch_size: 32 # 全局batch大小3. 剪枝专用配置 (model.l0_module)这是Sheared LLaMA的核心。model: l0_module: lagrangian_warmup_steps: 640 # 剪枝率热身步数20% * 3200 pruning_modules: [head, intermediate, hidden, layer] # 剪哪些结构 eval_target_model: false # 训练时评估当前带掩码的模型而非目标模型 target_model: d_model: 2048 # 目标模型隐藏层维度 (对应1.3B) n_heads: 16 # 目标模型注意力头数 n_layers: 24 # 目标模型层数 intermediate_size: 5504 # 目标模型FFN中间层维度from_model和to_model参数通常在命令行或脚本中指定指向定义源模型和目标模型架构的配置文件。optimizer.lag_lr是用于优化门控变量和拉格朗日乘子的学习率默认1.0通常不需要改动。4. 动态批次加载配置 (callbacks.data_loading)这是另一个亮点它允许在训练过程中动态调整不同数据源领域的采样比例。callbacks: data_loading: dynamic: true set_names: [domain1, domain2, domain3] # 你的数据子集名称 proportion: [0.3, 0.5, 0.2] # 初始采样比例总和为1 update_type: doremi # 使用Doremi算法动态调整比例 target_loss: 2.5 # 预设的目标验证损失值动态加载的原理是每隔一定的评估间隔(eval_interval)计算当前模型在各个数据域上的损失。如果某个域的损失远高于目标损失说明模型在这个域上表现不佳下一阶段就增加该域数据的采样比例反之则减少。这能帮助模型更均衡地学习多领域知识。配置完成后启动剪枝训练composer train.py \ /path/to/your/pruning_config.yaml \ --from_model /path/to/source_model_config.yaml \ --to_model /path/to/target_model_config.yaml \ --load_path /path/to/converted/composer/checkpoint.pt4.2 剪枝后处理模型转换与“瘦身”剪枝训练完成后保存的检查点包含了完整源模型的参数和一系列门控掩码。我们需要进行后处理得到真正的、紧凑的目标模型。MODEL_PATH/path/to/saved/checkpoint/latest-rank0.pt python3 -m llmshearing.utils.post_pruning_processing prune_and_save_model $MODEL_PATH这个脚本会做两件事应用掩码根据门控变量的值接近0的视为被剪枝物理地移除对应的权重和结构如整行的权重矩阵、整层的参数。重命名键值将权重字典的键名重新组织使其符合一个连续的、标准的目标模型架构定义以便后续加载。处理后的模型会保存在$(dirname $MODEL_PATH)/pruned-latest-rank0.pt。这个文件才是我们需要的、剪枝后的模型权重。4.3 持续预训练阶段得到剪枝后的模型后我们需要对其进行持续预训练以恢复性能。这个阶段的配置与剪枝阶段类似但更接近标准的语言模型预训练。关键变化移除剪枝配置YAML配置文件中不再需要model.l0_module部分。调整训练参数通常会增加训练步数(max_duration: 48000ba)、增大批次大小(global_train_batch_size: 256)、调整学习率热身(t_warmup: 1440ba约3%)。加载剪枝后的模型--load_path指向上一步生成的pruned-latest-rank0.pt文件。执行脚本参考llmshearing/scripts/continue_pretraining.sh。4.4 最终模型转换Composer格式转Hugging Face格式为了便于下游使用、推理和微调我们通常需要将Composer训练好的模型转换回Hugging Face Transformers格式。MODEL_PATH/path/to/your/composer/checkpoint.pt OUTPUT_PATH/path/to/output/hf_model MODEL_CLASSLlamaForCausalLM python3 -m llmshearing.utils.composer_to_hf save_composer_to_hf $MODEL_PATH $OUTPUT_PATH \ model_class${MODEL_CLASS} \ hidden_size2048 \ num_attention_heads16 \ num_hidden_layers24 \ intermediate_size5504 \ num_key_value_heads16 \ _name_or_pathSheared-Llama-1.3B注意这里的模型架构参数hidden_size,num_attention_heads等必须与你剪枝目标模型的配置严格一致而不是源模型的配置。转换成功后你就可以像使用任何其他Hugging Face模型一样使用from_pretrained(OUTPUT_PATH)来加载它了。5. 实战经验、常见问题与性能调优在实际操作中你可能会遇到一些预料之外的情况。这里分享一些从实验和社区反馈中总结的经验。5.1 硬件与吞吐量考量论文中提供了在A100 80GB GPU上的吞吐量参考这对于规划你的训练资源很重要阶段所用GPU数量单卡吞吐量 (tokens/sec)总吞吐量 (tokens/sec)剪枝 7B 模型8184414750持续预训练 3B 模型16495779306持续预训练 1.3B 模型168684138945经验分享剪枝阶段内存消耗大因为需要保存完整大模型的参数和额外的门控变量显存占用比普通训练高。8xA100 80GB剪裁7B模型是一个合理的起点。持续预训练效率高一旦模型被剪小持续预训练的吞吐量会大幅提升因为计算和通信量都减少了。批量大小选择剪枝阶段由于稳定性考虑用了较小的微批次(microbatch_size4)。在持续预训练时可以尝试增大以提升吞吐但要注意如果遇到CUDA OOM可以启用梯度累积来模拟更大的全局批次。5.2 超参数选择与调优论文作者也提到由于计算限制他们没有进行大规模的超参数搜索。这意味着现有的配置可能有优化空间。学习率1e-4是一个比较保守的起点。对于持续预训练如果损失下降很慢可以尝试稍微增大如2e-4如果训练不稳定损失出现NaN则需减小。剪枝热身步数lagrangian_warmup_steps设置为总步数的20%是一个经验值。这个阶段允许门控变量缓慢地从初始值变化避免过于激进的剪枝导致模型崩溃。如果发现模型在剪枝后期性能急剧下降可以尝试增加这个比例。目标损失动态加载中的target_loss是一个需要预先估算的值。一个实用的方法是先用初始比例混合数据训练一个很小的模型或训练几步看其验证损失是多少以此作为目标损失的粗略参考。5.3 常见问题排查训练中途崩溃NaN Loss检查学习率这是最常见的原因。尝试降低optimizer.lr和optimizer.lag_lr。检查数据确保数据中没有异常字符或空样本。Streaming数据集在构建时可能因为某些样本编码问题导致读取错误。梯度爆炸可以尝试添加梯度裁剪 (grad_clip_norm)。剪枝后模型性能远低于预期确认剪枝目标架构检查to_model的配置文件是否正确确保d_model,n_layers等参数是你想要的。检查后处理确保prune_and_save_model脚本成功运行并且生成的模型大小符合预期。可以用torch.load检查权重字典的键和形状。持续预训练数据不足或质量差剪枝后的模型需要高质量数据来恢复。确保你的持续预训练数据是多样且干净的。动态加载不工作或比例不更新确认配置确保callbacks.data_loading.dynamic设置为true并且set_names和proportion数组长度匹配。检查评估集动态更新依赖于每个域上的验证损失。确保你的评估集(evalsplit)包含了所有域的数据并且加载正确。查看日志Composer的日志会输出每个域上的损失值。检查这些值是否被正常计算和记录。转换到Hugging Face格式后加载失败架构参数不匹配这是最可能的原因。仔细核对save_composer_to_hf命令中输入的hidden_size,num_hidden_layers等参数必须与剪枝后模型的实际架构完全一致。一个有用的调试方法是先用这些参数在Hugging Face端初始化一个空的模型看看它的参数形状是否与你的权重字典匹配。5.4 扩展与未来方向Sheared LLaMA的框架是通用的不限于LLaMA模型。你可以尝试更强的基模型使用Mistral-7B、Gemma等更强大的7B模型作为剪枝起点有望得到性能更优的小模型。领域特定模型对CodeLlama、医学或法律领域的大模型进行剪枝快速获得一个该领域的高效小模型。探索不同剪枝粒度目前的pruning_modules选项比较全面。你也可以尝试只剪枝某一部分如只剪注意力头层观察对性能和效率的影响。要实现这些扩展关键步骤是确保模型与掩码的兼容性。项目中的llmshearing/utils/test_pruning.py脚本就是用来测试你为自定义模型实现的prune_params函数是否正确。你需要保证给模型参数加上掩码后进行前向传播与直接加载剪枝后的模型进行前向传播两者的输出在数值上是相等的允许极小的浮点误差。我个人在实验中发现这套流程的稳定性相当不错一旦环境配置正确复现论文结果并不困难。最大的挑战往往来自于数据管道的构建和计算资源的规划。对于资源有限的团队从一个大而强的开源模型出发通过结构化剪枝“蒸馏”出一个小而精的专用模型是一条非常务实且高效的技术路径。它让你能够站在巨人的肩膀上用有限的算力聚焦于模型架构的优化和领域数据的适配而不是重复消耗在通用的基础预训练上。