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 上更新