正文
为什么知识蒸馏有用?
知识蒸馏的理念是,相比从头开始训练或有监督微调,一个较小的模型可以通过使用教师模型的输出作为额外信号来获得更好的性能。例如,Llama 3.2轻量级1B和3B文本模型(https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices/)在剪枝后通过整合Llama 3.1 8B和70B的logits来恢复性能。此外,在指令跟随任务的微调中,LLM蒸馏的研究表明,知识蒸馏方法可以优于单独使用有监督微调(SFT)。
表1: 知识蒸馏方法与有监督微调的比较
下面是一个简化的例子,展示了知识蒸馏与有监督微调的区别。
torchtune中的知识蒸馏配方
使用torchtune,我们可以轻松地将知识蒸馏应用于Llama3以及其他LLM模型系列,这是通过使用torchtune的知识蒸馏配方(https://github.com/pytorch/torchtune/blob/4234b78b914af23384ce0348f564e2119d107a96/recipes/knowledge_distillation_single_device.py)实现的。这个配方的目标是通过从Llama3.1-8B蒸馏知识来在Alpaca指令跟随数据集上微调Llama3.2-1B。这个配方专注于训练后蒸馏,并假设教师和学生模型都已经预训练完成。
首先,我们需要下载模型权重。为了与其他torchtune微调配置保持一致,我们将使用Llama3.1-8B的指令调优模型作为教师模型,使用Llama3.2-1B作为学生模型。
tune download meta-llama/Meta-Llama-3.1-8B-Instruct --output-dir /tmp/Meta-Llama-3.1-8B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf_token
tune download meta-llama/Llama-3.2-1B-Instruct --output-dir /tmp/Llama-3.2-1B-Instruct --ignore-patterns "original/consolidated.00.pth" --hf_token
为了让教师模型的分布与Alpaca数据集相似,我们将使用LoRA对教师模型进行微调。基于我们在下一节展示的实验,我们发现当教师模型已经在目标数据集上进行了微调时,知识蒸馏的效果会更好。
tune run lora_finetune_single_device --config llama3_1/8B_lora_single_device
