GPT-2 蒸馏、压缩、剪枝实战

1.项目背景大语言模型虽然效果强但部署成本高。对于类似 GPT-2 这样的自回归语言模型常见的压缩思路主要有三类- 蒸馏让一个更小的 student 去模仿更大的 teacher- 剪枝删除一部分不重要的权重降低有效参数密度- 量化把权重和激活压缩到更低精度本项目聚焦前两种- 先把预训练 GPT-2 在本地语料上继续微调作为 teacher baseline- 再训练一个 6 层的 student GPT-2- 用 CE KL 做知识蒸馏- 再对蒸馏后的 student 做非结构化剪枝- 最后做短期恢复微调并用 PPL 进行对比实验语料使用了本地WikiText-103目前完整代码已开源在fsy233123/GPT2-Distilled: This is a local compression experiment project centered around GPT-2, aiming to fully implement the classic model compression route of distillation pruning, and strive to achieve the following as much as possible: - Locally executable - Clear structure - Supports single and multiple cards - Can directly produce comparable PPL resultshttps://github.com/fsy233123/GPT2-Distilled2.实验细节2.1 训练数据预处理首先进行将原始数据集预处理缓存到本地便于训练调用我设置的数据调用参数为block_size128。目前处理的流程是读取原始文本用 GPT-2 tokenizer 把文本转成 token ids把这些 token ids 按顺序拼接成一条连续 token 流每 block_size 个 token 切成一个 block在进行这步操作时我有一个疑问这种按block_size划分一批数据不会导致两个完全不同的句子的token划分到一个block中吗但实际上两个句子之间没有任何联系。实际上是会发生的很容易出现切块结果如... 句子A结束。 句子B开 ...但是这种策略是可以接受的因为大多数语言模型训练并不要求每个 block 都恰好是一句完整句子或者每个 block 都必须只包含一个语义独立样本对 GPT 来说更重要的是它持续看到真实文本流学到局部统计规律学到自然语言的 token 转移关系即使 block 内有句子边界模型也会学到换行标点段落结束文档结束符这些边界本身也是语言分布的一部分。比较正式的操作应该是eos_token或换行符或文档分隔符这样即使两个不相关文本被切到同一个 block模型也能知道“这里其实是一个边界”目前在我自己的项目中没有这种特殊的边界处理是一个简单的项目实战。2.2 Baseline我们加载已有的 GPT-2 预训练模型在本地 WikiText-103 数据上进行短期 fine-tuning然后在验证集上计算 perplexity。后续的 Student scratch、Student distilled、Distilled Pruned 模型都会和这个 baseline 结果进行对比。这里的 baseline代表的是“完整 GPT-2 模型在当前数据集上的适配后性能”。在 baseline 训练中我们采用标准的 causal language modeling loss也就是让 GPT-2 根据当前位置及其之前的上下文预测下一个 token。具体来说数据集中每个样本是一个长度为 block_size 的 token block例如当 block_size128 时一个 block 包含 128 个连续的 token_id。模型前向传播后会在每个位置输出一个词表维度的 logits 分布但训练时不会直接让第 t 个位置预测自己而是通过 shift 操作让第 t 个位置的输出预测第 t1 个 token。因此长度为 128 的 block 实际提供 127 个 next-token prediction 监督信号最后一个位置由于没有下一个 token不参与 loss 计算。在实现上代码中先令 labels input_ids.clone()然后在计算 loss 时执行 shift_logits logits[..., :-1, :] 和 shift_labels labels[..., 1:]。这样 logits[0] 对应预测 labels[1]logits[1] 对应预测 labels[2]依此类推。随后使用交叉熵损失 CrossEntropyLoss 计算每个预测位置上模型对真实下一个 token 的负对数似然并对 batch 内所有有效预测位置取平均。若 batch size 为 4、block size 为 128则单张 GPU 上一个 step 会产生 4 × 127 508 个 token-level loss最终训练 loss 是这些位置 loss 的平均值。2.3 Scratch在本项目中Student scratch 被设计为一个重要的对照实验用来观察“小模型在没有 Teacher 知识迁移的情况下仅依靠当前训练语料从零学习”的效果。具体来说我们基于 GPT-2 的配置构建了一个更小的 Student GPT-2 模型保留原 GPT-2 的词表大小、hidden size、attention head 数、位置编码长度等关键维度只将 Transformer Decoder Block 的层数从标准 GPT-2 small 的 12 层减少到 6 层。因此该模型在结构上仍然与 GPT-2 兼容但参数量明显下降。与蒸馏模型不同scratch 版本不会从 Teacher GPT-2 中拷贝权重也不会使用 Teacher 输出的 soft label 进行指导。它的 embedding、position embedding、Transformer block、LayerNorm 和 LM Head 都按照 GPT-2 默认初始化方式随机初始化然后只使用标准 causal language modeling 的交叉熵损失进行训练。也就是说模型需要完全依靠 WikiText-103 训练数据学习“根据前文预测下一个 token”的能力。2.4 Distilled在本项目中Teacher 模型使用的是 fine-tuned GPT-2 baselineStudent 模型则是一个 6 层 GPT-2。蒸馏训练希望 Student 不仅学习真实文本中的下一个 token 标签还学习 Teacher 对整个词表的预测分布。这样Student 可以从 Teacher 的 soft target 中获得比 hard label 更丰富的知识。整体流程可以概括为加载 fine-tuned GPT-2 baseline 作为 Teacher构建 6 层 GPT-2 Student用 Teacher 的部分层初始化 Student固定 Teacher 参数对同一个 batch 同时执行 Teacher 和 Student 前向传播计算 Student 对真实 label 的 CE loss计算 Student logits 和 Teacher logits 之间的 KL loss组合 CE 和 KL 得到蒸馏总 loss只更新 Student 参数验证 Student perplexity保存蒸馏后的 Student 模型因为 Teacher 有 12 层Student 有 6 层所以代码会从 Teacher 的 12 个 Transformer block 中均匀选择 6 层拷贝到 Student 中。例如大致会选择Teacher layer 0Teacher layer 2Teacher layer 4Teacher layer 7Teacher layer 9Teacher layer 11蒸馏中使用了 temperature--temperature 2.0当 T2 时预测分布会变得更平滑。原本非常尖锐的概率分布会被软化低概率 token 会获得更多可见的概率质量。在计算KL_loss乘上.这是知识蒸馏中的常见做法用来补偿 temperature 对梯度尺度的影响。否则 temperature 变大后KL loss 的梯度可能会被缩小。本项目中的蒸馏总 loss 是我们设置.2.5 Distilled Pruned在完成知识蒸馏后我们已经得到了一个 6 层的 Student GPT-2。相比原始 12 层 GPT-2它的参数量已经明显减少并且通过 Teacher 的指导保留了较好的语言建模能力。但模型压缩还可以继续推进除了减少层数还可以进一步把部分不重要的权重置为 0这就是剪枝。本项目中的 Distilled Pruned 阶段主要做三件事加载蒸馏后的Student模型对Attention和MLP中的投影权重进行全局非结构化剪枝对剪枝后的模型进行短期Recovery fine-tuning剪枝对象只剪 attention 和 MLP 投影层本项目并不是对 GPT-2 的所有参数都剪枝而是只剪 Transformer block 内部的主要线性投影权重。代码中收集可剪枝模块的逻辑是for block in model.transformer.h: for module_name in (attn.c_attn, attn.c_proj, mlp.c_fc, mlp.c_proj): module block.get_submodule(module_name)也就是说每一层 Transformer block 中会剪这些部分attn.c_attn: attention 中的 QKV 联合投影attn.c_proj: attention 输出投影mlp.c_fc: MLP 中的升维投影mlp.c_proj: MLP 中的降维投影对于 GPT-2 来说这些模块在 Hugging Face 实现中很多是 Conv1D 类型但本质上可以理解为线性投影层。项目中特意不剪token embedding,position embedding,LayerNorm,final LayerNorm这样设计比较稳妥。Embedding 直接关系到词表表示如果剪得太激进可能会明显破坏模型对 token 的基础表达能力。LayerNorm 参数量很小但对训练稳定性和分布校准很重要也不适合作为主要剪枝对象。剪枝方法Global Unstructured L1 Pruningprune.global_unstructured( [(module, weight) for module in prunable_modules], pruning_methodprune.L1Unstructured, amountargs.sparsity, )Unstructured 表示非结构化剪枝。也就是说它不是删除整个 attention head也不是删除整行、整列、整个 neuron而是直接在权重矩阵内部把一些单独的标量权重置为 0。例如原始权重可能是[0.82, -0.01, 0.37, 0.004, -0.55]剪枝后可能变成[0.82, 0.00, 0.37, 0.00, -0.55]L1Unstructured 表示根据权重绝对值大小判断重要性。绝对值越小的权重被认为越不重要越容易被剪掉。例如abs(-0.01) 0.01,abs(0.004) 0.004.这些权重就比 0.82 或 -0.55 更容易被置为 0。global_unstructured 表示全局剪枝。它不是每一层都固定剪掉 40%而是把所有可剪枝层的权重放在一起排序然后从全局范围内选出绝对值最小的一部分置零。这意味着有些层可能被剪得更多有些层可能被剪得更少具体取决于各层权重的绝对值分布。相比“每层单独剪 40%”全局剪枝通常更灵活也更符合“在整个模型范围内优先移除最不重要权重”的直觉。为什么剪枝后还要保存 mask虽然 prune.remove 会把被剪掉的权重变成真实的 0但如果后续继续 fine-tuning优化器仍然可能把这些 0 权重重新更新成非零。为了避免这种情况项目在 solidify_pruning 后额外保存了一个 maskmasks[id(module)] module.weight.detach().ne(0)这个 mask 记录哪些位置是保留权重哪些位置是剪掉权重。然后在 recovery fine-tuning 中每次 optimizer.step()后都会执行:enforce_pruning_masks(...)实际逻辑是:module.weight.mul_(mask)也就是说每次参数更新后都会把被剪掉的位置重新乘以 0确保它们保持为 0。这个设计使得剪枝后的稀疏结构在微调过程中不会被破坏。Recovery Fine-tuning剪枝后的性能恢复剪枝后脚本会对模型进行短期微调。这一步也叫 recovery fine-tuning目的不是重新训练模型而是让剩余的非零权重适应剪枝后的稀疏结构。训练方式和 baseline 类似使用 causal language modeling 的 CE loss.3.结果分析最终各模型结果如下Student distilled 的 PPL 为 24.9497相比 scratch 的 28.3608 有明显提升说明知识蒸馏有效。它通过 teacher-copy 初始化和 CE KL 联合损失不仅学习真实下一个 token还学习 Teacher 在整个词表上的 soft distribution因此比从零训练更容易接近完整 GPT-2 的语言建模能力。最值得注意的是 Distilled Pruned在可剪枝权重中达到 40.0% 稀疏率后PPL 反而降到 22.7057优于未剪枝 distilled 的 24.9497只比完整 GPT-2 baseline 高约 1.2682。这说明剪枝并没有破坏 Student反而可能起到了类似正则化的作用删除了一部分小幅度、低贡献或冗余权重后再经过 recovery fine-tuning模型泛化表现变好了。4.结论实验结果表明知识蒸馏能够显著改善小规模 Student GPT-2 的语言建模能力而剪枝后的短期微调可以有效恢复甚至提升模型在验证集上的表现。最终 Distilled Pruned 模型在参数量减少约 34%、可剪枝权重达到 40% 稀疏率的情况下验证集 PPL 仅比完整 GPT-2 baseline 高约 5.9%说明该压缩流程在当前 WikiText-103 实验设置下具有较好的有效性。