Sheared LLaMA：结构化剪枝与持续预训练实现高效大模型压缩-开发者社区

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

# 安装指定版本的PyTorch（CUDA 11.8） pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118 # 安装Flash Attention 1.x版本（注意：暂不支持2.0） pip install flash-attn==1.0.3.post

注意：Flash Attention 2目前不被支持，强行使用可能需要手动修改模型文件，容易出错。请务必使用1.x版本。

步骤二：安装项目依赖克隆项目代码后，进入目录安装剩余依赖。

cd LLM-Shearing pip install -r requirement.txt

步骤三：以可编辑模式安装llmshearing包这样做的好处是，你可以在开发过程中直接修改本地代码，而无需反复重新安装。

pip install -e .

3.2 数据准备：使用MosaicML Streaming

Sheared 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_NAME="meta-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_SIZE="7B" 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.pt

4.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_CLASS=LlamaForCausalLM python3 -m llmshearing.utils.composer_to_hf save_composer_to_hf $MODEL_PATH $OUTPUT_PATH \ model_class=${MODEL_CLASS} \ hidden_size=2048 \ num_attention_heads=16 \ num_hidden_layers=24 \ intermediate_size=5504 \ num_key_value_heads=16 \ _name_or_path="Sheared-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 模型	8	1844	14750
持续预训练 3B 模型	16	4957	79306
持续预训练 1.3B 模型	16	8684	138945

经验分享：

剪枝阶段内存消耗大：因为需要保存完整大模型的参数和额外的门控变量，显存占用比普通训练高。8xA100 80GB剪裁7B模型是一个合理的起点。
持续预训练效率高：一旦模型被剪小，持续预训练的吞吐量会大幅提升，因为计算和通信量都减少了。
批量大小选择：剪枝阶段由于稳定性考虑，用了较小的微批次(microbatch_size=4)。在持续预训练时，可以尝试增大以提升吞吐，但要注意如果遇到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函数是否正确。你需要保证，给模型参数加上掩码后进行前向传播，与直接加载剪枝后的模型进行前向传播，两者的输出在数值上是相等的（允许极小的浮点误差）。

我个人在实验中发现，这套流程的稳定性相当不错，一旦环境配置正确，复现论文结果并不困难。最大的挑战往往来自于数据管道的构建和计算资源的规划。对于资源有限的团队，从一个大而强的开源模型出发，通过结构化剪枝“蒸馏”出一个小而精的专用模型，是一条非常务实且高效的技术路径。它让你能够站在巨人的肩膀上，用有限的算力聚焦于模型架构的优化和领域数据的适配，而不是重复消耗在通用的基础预训练上。