text-generation-inference 文档
训练 Medusa
并获得增强的文档体验
开始使用
训练 Medusa
本教程将向您展示如何根据您选择的数据集训练 Medusa 模型。有关 Medusa 工作原理和推测的更多信息，请查看推测文档。
训练 Medusa 模型的优势是什么？
训练 Medusa 头可以大大提高生成速度。Medusa 在 LLM 中添加了额外的“头”，以同时预测多个未来的 token。当使用 Medusa 增强模型时，原始模型保持不变，只有新的头在训练期间进行微调。
最重要的事情之一是拥有一个好的数据集（数据与生产中将使用的数据相似），因为当生成是领域内时，Medusa 的命中率要高得多。
如果您在与生产中将使用的数据集非常不同的数据集上训练 Medusa，那么模型将无法准确预测未来的 token，因此加速将是最小或不存在的。
自蒸馏（生成训练数据）
有许多方法可以准备训练数据，但最简单有效的方法之一是“自蒸馏”数据。这意味着您可以使用相同的模型来生成将用于训练模型的数据。
本质上，您使用与生产中将使用的输入类似的输入提示模型，模型将生成输出。
我们将使用此输出来帮助训练 Medusa 头，以预测序列中的 n+1、n+2、n+3 等 token。
训练
Medusa 的原始实现可在 https://github.com/FasterDecoding/Medusa 获得，我们将遵循与原始存储库中描述的非常相似的过程来训练模型。
开始
有两种训练模型的方法
torchrun是torch.distributed.launch的一个封装- 一个支持 Medusa 的
axlotl分叉版本
在本教程中，我们将使用 torchrun 来训练模型，因为它是训练模型最直接的方法，但如果您愿意，也可以遵循类似的步骤使用 axlotl 训练模型。
使用 torchrun 训练
mkdir medusa-training
cd medusa-training
pyenv install 3.10
pyenv local 3.10
uv venv -p 3.10
source .venv/bin/activate现在让我们克隆原始的 Medusa 仓库并安装库。
git clone https://github.com/FasterDecoding/Medusa.git
cd Medusa
pip install -e .接下来我们需要一些数据进行训练，我们可以使用 Hugging Face Hub 上提供的 ShareGPT_Vicuna_unfiltered 数据集。
apt install git-lfs
git lfs install
git clone https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered目前我们的目录结构如下：
. ├── assets ├── CITATION.cff ├── create_data.py ├── data_generation ├── deepspeed.json ├── last_run_prepared ├── LICENSE ├── llm_judge ├── medusa ├── medusa_llm.egg-info ├── mistral.json ├── notebooks ├── pyproject.toml ├── README.md ├── ROADMAP.md ├── scripts ├── ShareGPT_Vicuna_unfiltered │ ├── README.md │ ├── ShareGPT_2023.05.04v0_Wasteland_Edition.json │ └── ShareGPT_V4.3_unfiltered_cleaned_split.json ├── simple_gradio_interface.py ├── tiny-llama.json └── vicuna_7b_qlora_stage1
开始训练
现在让我们生成数据并开始训练模型。这个过程将需要一段时间，因为我们正在从模型生成数据。
首先确保您正在运行一个 TGI 实例，其中包含您想用于自蒸馏的模型。
model=HuggingFaceH4/zephyr-7b-beta
volume=/home/ubuntu/.cache/huggingface/hub/
docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:latest --model-id $model现在我们可以使用 create_data.py 脚本生成数据。
python create_data.py \
--input-filename ShareGPT_Vicuna_unfiltered/ShareGPT_V4.3_unfiltered_cleaned_split.json \
--output-filename zephyr_self_distill.json此时我们的终端应该看起来像这样
注意：上面的屏幕截图中，我们只使用了数据集的前500个示例来加速处理，您应该使用更大的数据集进行训练。
现在我们终于可以开始训练模型这个有趣的部分了！
使用 torchrun，我们可以轻松启动 medusa 训练脚本，并使用 zephyr_self_distill.json 配置文件。
注意：如果您刚刚完成了自蒸馏，模型可能仍在运行，请务必在开始训练之前停止它，以便所有资源都能用于训练。
WANDB_MODE=offline torchrun --nproc_per_node=4 medusa/train/train_legacy.py \
--model_name_or_path HuggingFaceH4/zephyr-7b-beta \
--data_path zephyr_self_distill.json \
--bf16 True \
--output_dir zephyr_out \
--num_train_epochs 5 \
--per_device_train_batch_size 4 \
--per_device_eval_batch_size 4 \
--gradient_accumulation_steps 4 \
--evaluation_strategy "no" \
--save_strategy "no" \
--learning_rate 1e-3 \
--weight_decay 0.0 \
--warmup_ratio 0.1 \
--lr_scheduler_type "cosine" \
--logging_steps 1 \
--tf32 True \
--model_max_length 2048 \
--lazy_preprocess True \
--medusa_num_heads 3 \
--medusa_num_layers 1 \
--deepspeed deepspeed.json
如果成功，您应该看到与以下类似的输出：
wandb: Run history:
wandb: train/epoch ▁▁▁▁▁▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇███
wandb: train/global_step ▁▁▁▁▁▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇███
wandb: train/learning_rate ▅███▇▇▆▅▅▄▃▂▂▁▁▁
wandb: train/loss ██▆▄▄▃▃▂▂▃▁▁▂▁▁▁
wandb: train/medusa0_loss ▆▆▇▆▆▅▄▅▃▃▃▃▂▂▂▂▂▃▂▂▂▁▁▁▂▁▁▁▁▁█▁▁▁▂▁▁▁▁▁
wandb: train/medusa0_top1 ▁▁▁▁▁▁▁▁▃▂▃▃▄▄▄▃▄▃▄▄▅▅▆▅▆▆▇▅▇▇▄▇█▇▅▇█▆▇▇
wandb: train/medusa1_loss ▇▇█▇▇▆▅▅▃▄▃▃▃▃▃▃▃▃▃▃▂▁▂▂▂▁▁▂▁▁▇▁▁▁▂▁▁▁▁▁
wandb: train/medusa1_top1 ▁▁▁▁▁▁▁▁▃▂▃▃▃▄▄▃▃▂▃▃▅▅▆▄█▆▇▅▇▇▅█▇▇▅▇█▆▆▇
wandb: train/medusa2_loss ▃▃▄▄▄▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁█▁▁▁▂▁▁▁▁▁
wandb: train/medusa2_top1 ▁▁▁▂▁▁▁▁▂▂▃▃▃▄▄▃▃▂▃▃▅▆▅▄█▆▆▅▆▆▄█▇▇▄▇█▆▆▇
wandb: train/total_flos ▁
wandb: train/train_loss ▁
wandb: train/train_runtime ▁
wandb: train/train_samples_per_second ▁
wandb: train/train_steps_per_second ▁
wandb:
wandb: Run summary:
wandb: train/epoch 2.0
wandb: train/global_step 16
wandb: train/learning_rate 0.0
wandb: train/loss 14.8906
wandb: train/medusa0_loss 4.25
wandb: train/medusa0_top1 0.28809
wandb: train/medusa1_loss 4.8125
wandb: train/medusa1_top1 0.22727
wandb: train/medusa2_loss 5.5
wandb: train/medusa2_top1 0.17293
wandb: train/total_flos 0.0
wandb: train/train_loss 23.98242
wandb: train/train_runtime 396.9266
wandb: train/train_samples_per_second 2.519
wandb: train/train_steps_per_second 0.04最后但最重要的是，别忘了将此模型推送到 Hugging Face Hub，以便您可以在项目中利用它。
python -m medusa.hf_utils \
--folder zephyr_out_medusa_mlp_zephyr-7b-beta_medusa_3_lr_0.001_layers_1 \
--repo drbh/zephyr_medusa_demo哇，我们成功训练了一个 Medusa 模型并将其推送到了 Hugging Face Hub！🎉
< > 在 GitHub 上更新