Integrasi vLLM

Dokumen ini akan memandu Anda melalui proses penggunaan vLLM dengan TRL untuk generasi yang lebih cepat pada metode online seperti GRPO dan Online DPO. Kami pertama-tama merangkum tl;dr tentang cara menggunakan vLLM dengan TRL, kemudian kami akan membahas detail tentang cara kerjanya di balik layar.

TRL saat ini hanya mendukung versi vLLM dari 0.23.0 hingga 0.33.x. Pastikan Anda menginstal versi dalam rentang ini untuk menghindari masalah kompatibilitas.

Trainer berikut saat ini mendukung generasi dengan vLLM:

• GRPOTrainer
• RLOOTrainer
• experimental.online_dpo.OnlineDPOTrainer

🚀 Bagaimana saya bisa menggunakan vLLM dengan TRL untuk mempercepat pelatihan?

💡 Catatan: Sumber daya yang diperlukan untuk contoh spesifik ini: satu node dengan 8 GPU.

Saat menggunakan vLLM dengan TRL, server vLLM dan trainer harus dijalankan pada perangkat CUDA yang terpisah untuk mencegah konflik. Untuk panduan tentang cara mengonfigurasi ini dengan benar, lihat Modes of using vLLM during training.

Pertama, instal vLLM menggunakan perintah berikut:

pip install "trl[vllm]"

Lalu jalankan server pada GPU tertentu (misalnya, GPU 0-3):

CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-7B --tensor-parallel-size 4 \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

Setelah server berjalan, Anda dapat menggunakannya untuk menghasilkan kelengkapan (completions) untuk pelatihan. Pada contoh di bawah, kami menggunakan berbagai trainer yang didukung dengan menggunakan server vLLM untuk generasi. Argumen --tensor-parallel-size dan --data-parallel-size mengontrol bagaimana model dan data dipecah (sharded) di antara GPU.

Pada contoh ini, kami membagi satu model ke 4 GPU dengan tensor parallelism. Kemudian, jalankan skrip pelatihan di GPU yang berbeda (misalnya, GPU 4-7) dengan meneruskan use_vllm=True dalam argumen pelatihan sebagai berikut:

Contoh skrip train.py sederhana:

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=GRPOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

from datasets import load_dataset
from trl.experimental.online_dpo import OnlineDPOConfig, OnlineDPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = OnlineDPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=OnlineDPOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = RLOOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=RLOOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

Dan perintah train pada GPU yang terpisah dari server:

CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py

Mengapa menggunakan vLLM?

Metode online menghasilkan kelengkapan selama pelatihan, dan menghasilkan kelengkapan tersebut dengan generate milik model sendiri merupakan titik penyumbat. vLLM menyajikan kelengkapan tersebut jauh lebih cepat, berkat teknik seperti PagedAttention.

Cara TRL menggunakan server 🔍

Trainer meminta completion pada endpoint /v1/completions yang kompatibel dengan OpenAI, dengan mengirimkan ID token prompt. Prompt multimodal mengambil jalur yang berbeda: server memproses gambar secara mandiri, dan fitur yang dihasilkan dipasangkan dengan ID token yang sama pada /inference/v1/generate, karena tidak ada endpoint yang kompatibel dengan OpenAI yang menerima ID token dan gambar sekaligus.

Server hanya menghasilkan. Setelah setiap langkah optimizer, trainer mengalirkan bobot yang diperbarui ke dalamnya melalui NCCL, mengumumkannya dengan /start_weight_update dan /update_weights, lalu mengonfirmasinya dengan /finish_weight_update.

Penggunaan lanjutan

🍷 Opsi kustomisasi lebih lanjut dengan vLLM?

Anda dapat menyesuaikan konfigurasi server dengan meneruskan argumen vllm serve apa pun, misalnya --tensor-parallel-size, --data-parallel-size, --max-model-len, --enable-prefix-caching, --enforce-eager, --kv-cache-dtype atau --trust-remote-code. Jalankan vllm serve --help untuk daftar lengkap.

Hanya yang berikut ini yang diperlukan oleh TRL:

trl vllm-serve sudah tidak didukung lagi: perintah ini sekarang hanya membangun perintah ini dan menjalankan server vLLM. Perintah ini mencetak perintah vllm serve yang tepat yang dijalankan, sehingga Anda dapat menyalinnya dan menghapus pembungkusnya.

💆🏻‍♀️ Setup distribusi mana yang terbaik?

Skalakan generasi dengan --tensor-parallel-size. Parallelisme data tidak lagi membantu model padat: sejak vLLM PR #30739 (dirilis di 0.14.0), penskalaan parallelisme data offline untuk model non-MoE tidak didukung.

vLLM dengan Backend Transformers

vLLM dapat menggunakan backend Transformers untuk implementasi model, yang berfungsi untuk LLM maupun VLM. Untuk mengaktifkannya, atur vllm_model_impl="transformers" dalam konfigurasi Anda atau lewati melalui argumen baris perintah.

Untuk detail lebih lanjut, lihat vLLM Transformers Backend.

Contoh:

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-VL-3B-Instruct \
    --tensor-parallel-size 1 --port 8000 --enforce-eager --model-impl transformers \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1 \
    --enable-scale-out  # vLLM 0.30.0 and later, drop on earlier versions

Mode Penggunaan vLLM Selama Pelatihan

TRL mendukung dua mode untuk mengintegrasikan vLLM selama pelatihan: mode colocate (default) dan mode server.

Mode Colocate

Dalam mode colocate, vLLM berjalan di dalam proses trainer dan berbagi memori GPU dengan model pelatihan. Hal ini menghindari peluncuran server terpisah dan dapat meningkatkan utilisasi GPU, tetapi dapat menyebabkan perebutan memori pada GPU pelatihan. Ini adalah mode default.

Contoh konfigurasi:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

Mode Server

Dalam mode server, vLLM berjalan sebagai proses terpisah pada GPU khusus dan berkomunikasi dengan trainer melalui HTTP. Konfigurasi ini ideal jika Anda memiliki GPU yang khusus untuk inferensi.

Contoh konfigurasi:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

Periksa dokumentasi trainer yang Anda gunakan untuk detail spesifik mengenai penggunaan dan parameter vLLM.

Untuk mengurangi penggunaan memori GPU saat menjalankan vLLM, pertimbangkan untuk mengaktifkan mode tidur vLLM.