訓練 Medusa

本教學將向您展示如何使用您選擇的資料集來訓練 Medusa 模型。請參閱推測文件 (speculation documentation)，以獲取有關 Medusa 運作方式及推測技術的更多資訊。

訓練 Medusa 模型有什麼好處？

訓練 Medusa 頭部 (heads) 可以大幅提升生成速度。Medusa 為大型語言模型 (LLM) 增加了額外的「頭部」，用以同時預測多個未來的 token。在使用 Medusa 增強模型時，原始模型保持不變，僅在訓練過程中對新的頭部進行微調。

最重要的事情之一是擁有一個優質的資料集（與生產環境中所使用的資料相似），因為當生成的內容屬於領域內 (in-domain) 時，Medusa 的命中率會高得多。

如果您在與生產環境差異極大的資料集上訓練 Medusa，模型將無法準確預測未來的 token，因此加速效果將微乎其微甚至不存在。

自我蒸餾 (為訓練生成資料)

有許多準備訓練資料的方法，但最簡單且最有效的方法之一是「自我蒸餾」(self-distill) 資料。這意味著您可以使用同一個模型來生成您稍後將用於訓練模型的資料。

基本上，您使用與生產環境中相似的輸入提示 (prompt) 模型，模型就會產生輸出。

我們將利用此輸出幫助訓練 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，我們可以輕鬆啟動帶有 zephyr_self_distill.json 設定檔的 medusa 訓練腳本。

注意：如果您剛完成自我蒸餾，模型可能仍在執行中，請務必先將其停止，以便將所有資源用於訓練。

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