Distilasi Asinkron

Trainer ini membutuhkan vllm, transformers>=5.2.0 dan kernels (pelatihan memuat FlashAttention dari Hub melalui modul ini). Untuk pelatihan terdistribusi, hanya FSDP2 yang didukung (DeepSpeed ZeRO tidak).

Saat ini, vllm dan transformers memiliki batasan dependensi yang saling bertentangan. Untuk mengatasinya, instal vLLM terlebih dahulu, lalu paksa instal transformers, kemudian tambahkan kernels melalui extra TRL, yang menetapkan versi yang diterima oleh transformers:

Disalin

pip install vllm
pip install 'transformers>=5.2.0' --no-deps
pip install 'trl[kernels]'

Ikhtisar

experimental.async_distillation.AsyncDistillationTrainer adalah padanan asinkron dari DistillationTrainer, yang dirancang seperti experimental.async_grpo.AsyncGRPOTrainer: pekerja rollout latar belakang menghasilkan kelengkapan on-policy milik sendiri dari model siswa dan memberikan skor terhadap guru, sementara pelatihan berlangsung secara bersamaan alih-alih bergantian antara generasi dan pembaruan gradien. Berbeda dengan trainer sinkron, guru tidak pernah dimuat secara lokal — hanya URL server vLLM yang diperlukan, sehingga guru dapat berjalan di perangkat keras yang sepenuhnya terpisah dari siswa dan trainer, atau bahkan menjadi model yang jauh lebih besar daripada yang biasanya muat berdampingan dengan siswa.

compute_loss meminimalkan Jensen-Shannon Divergence yang digeneralisasi antara distribusi token per posisi dari student dan teacher (beta=0.0 adalah forward KL, beta=1.0 adalah reverse KL, nilai di antaranya melakukan interpolasi), yaitu tujuan yang sama yang digunakan oleh DistillationTrainer dan ServerDistillationTrainer. Pelatihan selalu on-policy: student menghasilkan setiap completion yang dilatihinya.

Distilasi on-policy multi-guru (MOPD)

MOPD bukan bagian dari makalah 2306.13649 yang menjadi dasar tujuan inti trainer ini (yang menggambarkan pengaturan sinkron dengan satu guru); ini adalah metode terpisah, yang dijelaskan dalam MOPD: Multi-Teacher On-Policy Distillation for Capability Integration in LLM Post-Training. Di sana, MOPD adalah tahap ketiga dari tiga tahap — SFT umum, kemudian pelatihan RL per-domain independen untuk satu ahli per domain, lalu MOPD menggabungkan ahli-ahli beku tersebut menjadi satu siswa. AsyncDistillationTrainer hanya mengimplementasikan tahap ketiga, yaitu tahap fusi: guru-guru ahli per-domain harus sudah ada (misalnya dilatih secara terpisah dengan GRPOTrainer/RLOOTrainer) dan disajikan melalui HTTP sebelum Anda mengarahkan teacher_server_urls ke mereka. Tahap 3 dari makalah itu sendiri menggunakan reverse KL (beta=1.0), bukan default trainer ini beta=0.0 (forward KL).

teacher_server_urls menerima lebih dari satu entri. Dengan satu entri, setiap sampel dinilai oleh satu guru tersebut (distilasi on-policy biasa). Dengan beberapa entri, kolom teacher_id dari setiap sampel pelatihan memilih guru mana yang akan menilainya — misalnya, merutekan prompt matematika ke guru spesialis matematika dan prompt kode ke guru spesialis kode, masing-masing disajikan secara independen. Setiap sampel dikirim ke satu guru yang cocok, tidak pernah dirata-rata atau di-ensemble antar guru. Sampel dengan teacher_id yang hilang atau tidak terpetakan akan memicu error alih-alih secara diam-diam kembali ke guru yang salah. Lihat examples/async_distillation_mopd/async_distillation_mopd.py untuk contoh yang dapat dijalankan yang menggabungkan ahli matematika, kode, dan mengikuti instruksi yang dirilis bersama Open-MOPD ke checkpoint SFT bersama mereka, dan examples/async_distillation_mopd/launcher.py untuk menjalankannya di 4 GPU dengan satu perintah Hugging Face Jobs.

Setiap guru harus membagikan tokenizer milik siswa. Kelengkapan dikirim ke para guru sebagai id token mentah, dan id kandidat yang dilaporkan kembali oleh seorang guru mengindeks langsung ke dalam kosakata siswa sendiri di compute_loss. Guru dari keluarga model yang sama dengan siswa (seperti dalam contoh MOPD, di mana siswa SmolLM3 difusi dari para ahli SmolLM3 yang dilatih dari checkpoint yang sama) memenuhi syarat ini; guru dengan kosakata yang berbeda akan melatih siswa terhadap token yang salah, secara diam-diam kecuali kosakatanya lebih besar daripada kosakata siswa.

Checkpoint dan resume: ignore_data_skip secara default bernilai True; loop skip-and-replay dari Trainer dasar tidak berlaku untuk antrean rollout live. Sebagai gantinya, indeks prompt pertama yang belum dilatih disimpan ke rollout_state.json bersama setiap checkpoint dan dipulihkan saat resume, sehingga worker melakukan fast-forward ke prompt tersebut tanpa memutar ulang sampel. Ini adalah posisi terlatih, bukan posisi generator: worker berjalan di depan pelatihan sebesar kedalaman antrean rollout, dan sampel yang dibuffer tersebut hilang saat run berakhir, sehingga resume dari posisi generator akan melewati prompt yang telah dihasilkan tetapi tidak pernah dilatih. Dataset streaming (IterableDataset) tidak dapat diposisikan ulang; worker-nya dimulai ulang dari prompt 0 saat resume. Di bawah FSDP dengan bf16=True (default) atau fp16=True, model.safetensors ditulis dalam dtype tersebut meskipun master weights adalah fp32, sehingga checkpoint yang di-push berukuran sama dengan model yang disajikan, bukan dua kali lipat; master weights fp32 dan state optimizer tetap masuk ke file resume (pytorch_model_fsdp.bin, optimizer.bin), yang merupakan file yang dibaca saat resume, sehingga proses resume tidak terpengaruh. Tanpa FSDP, checkpoint juga merupakan file resume, sehingga mempertahankan dtype master.

Bagaimana perbedaannya dengan DistillationTrainer

Di DistillationTrainer, model guru dimuat secara lokal: generasi, forward pass guru, dan pembaruan gradien semuanya terjadi secara berurutan dalam proses yang sama. experimental.async_distillation.AsyncDistillationTrainer memisahkan kekhawatiran ini dengan cara yang sama seperti experimental.async_grpo.AsyncGRPOTrainer memisahkan rollout GRPO dari pembaruannya:

• Worker rollout (proses latar belakang) — menghasilkan kelengkapan dari server vLLM siswa, mengirimkan urutan lengkap kembali ke /v1/completions dari guru yang telah dipecahkan dengan prompt_logprobs untuk penilaian yang dipaksa oleh guru (tidak ada token baru yang dihasilkan oleh guru), dan mendorong sampel yang siap dilatih ke dalam antrean.
• Loop pelatihan (proses utama) — mengambil sampel dari antrean, menghitung kerugian generalized-JSD, dan memperbarui bobot siswa.

Karena guru dinilai melalui HTTP alih-alih forward pass lokal, hanya potongan top-k yang jarang dari distribusinya yang dikirimkan (teacher_top_k), bukan seluruh kosakata — lihat dokumentasi beta dan teacher_top_k dari experimental.async_distillation.AsyncDistillationConfig untuk mengetahui kandidat mana yang protokol wire menjamin logprob guru untuknya pada setiap rezim beta.

Setiap weight_sync_steps langkah pelatihan, bobot siswa yang diperbarui ditransfer ke server vLLM-nya melalui NCCL. Sama seperti experimental.async_grpo.AsyncGRPOTrainer, proses generasi berjalan di depan pelatihan, sehingga sampel mungkin mencerminkan kebijakan yang sedikit usang; max_staleness mengontrol berapa banyak pembaruan bobot yang boleh tertinggal oleh sebuah sampel sebelum sampel tersebut dibuang.

Mulai cepat

# train_async_distillation.py
from datasets import load_dataset
from trl.experimental.async_distillation import AsyncDistillationTrainer

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

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

Server guru, server vLLM siswa, dan trainer harus dijalankan di GPU yang terpisah. Keduanya adalah instance vllm serve biasa, tetapi membutuhkan flag yang berbeda: guru hanya dinilai (--logprobs-mode processed_logprobs agar teacher_temperature mencapai logprobs yang dikembalikan, --max-logprobs -1 agar teacher_top_k dapat melebihi batas default vLLM sebesar 20), sedangkan siswa hanya digunakan untuk menghasilkan teks dan membutuhkan transfer bobot NCCL diaktifkan agar trainer dapat mendorong bobot yang diperbarui ke dalamnya. Penilaian mengirimkan seluruh urutan siswa dan meminta satu token lagi dari guru, sehingga --max-model-len guru (atau default-nya) harus lebih besar dari milik siswa:

# Terminal 1: teacher server on GPU 0 (static, never updated)
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen2.5-1.5B-Instruct \
    --port 8001 \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

# Terminal 2: student's vLLM server on GPU 1 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}'

# Terminal 3: training on GPU 2
CUDA_VISIBLE_DEVICES=2 accelerate launch train_async_distillation.py

LoRA

Lewati peft_config untuk melatih adapter LoRA pada student alih-alih model penuh:

from peft import LoraConfig

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)

Setiap sinkronisasi bobot harus mengirimkan model student yang diperbarui ke servernya. Secara default, trainer menggabungkan adapter ke dalam model dasar dan mengirimkan bobot gabungan penuh, yang ukurannya sekitar 15 GB untuk model 7B pada setiap sinkronisasi. Jalankan server vLLM student dengan --enable-lora dan server tersebut hanya akan mengirimkan adapter, sekitar 160 MB pada r=32:

CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}' \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 membuka endpoint yang digunakan trainer untuk mengirim setiap adapter baru. Tambahkan VLLM_WORKER_MULTIPROC_METHOD=spawn jika Anda menyajikan tensor-parallel, dan pertahankan --weight-transfer-config dalam kondisi apa pun: trainer hanya memilih mode sinkronisasi saat dimulai, pada titik tersebut server sudah berjalan, dan merged sync adalah fallback. Server teacher tidak terpengaruh: ia tidak pernah diperbarui.

--max-lora-rank harus salah satu dari 1, 8, 16, 32, 64, 128, 256, 320, 512. Parameter ini mengatur rank tertinggi yang dapat dilayani oleh server, bukan rank yang akan dilayani, sehingga adapter r=4 tetap berfungsi dengan baik di bawah 8. --max-loras harus setidaknya max_staleness + 2: trainer mempertahankan max_staleness + 1 versi adapter yang terdaftar agar rollout yang dimulai di bawah kebijakan yang lebih lama dapat diselesaikan di bawah kebijakan tersebut, bukan beralih kebijakan di tengah generasi, dan setiap sinkronisasi memuat versi berikutnya sebelum memuat ulang yang tertua. Trainer memeriksa kedua nilai ini saat dimulai dan memberi tahu Anda dengan nilai apa server harus di-restart.

vLLM memuat adapter dari disk, bukan melalui jaringan, sehingga trainer menulis setiap versi di bawah <output_dir>/.vllm_lora/ dan memberikan path tersebut ke server. Jika server berjalan di mesin yang berbeda, direktori tersebut harus berada pada filesystem yang dapat diakses oleh keduanya.

<output_dir>/.vllm_lora/ adalah cache untuk serving, bukan tempat adapter yang Anda latih disimpan. Direktori ini hanya menyimpan versi-versi terbaru, dan masing-masing akan dihapus segera setelah keluar dari jendela keusangan (staleness window). Checkpoint-lah yang melestarikan adapter yang Anda latih, jadi jangan nonaktifkan fitur ini: dengan save_strategy="no", eksekusi berdurasi beberapa jam akan selesai tanpa meninggalkan apa pun. output_dir juga harus dapat dibaca oleh server dan harus bertahan lebih lama dari job, sehingga scratch node-lokal tidak akan berfungsi, meskipun hal itu akan menjaga sinkronisasi agar tidak membebani filesystem bersama yang lambat.

Sinkronisasi adapter tidak selalu lebih cepat. vLLM menghasilkan token sekitar 1,33–1,39x lebih lambat per token saat menyajikan adapter LoRA dibandingkan saat menyajikan bobot yang sudah digabungkan, dan proses generasi biasanya yang mendominasi waktu eksekusi. Sebagai gantinya, Anda mendapatkan jeda yang jauh lebih singkat di setiap sinkronisasi dan sekitar 1% dari pita lebar sinkronisasi. Jika throughput generasi adalah hambatan utama Anda, sinkronisasi gabungan sering kali tetap menjadi pilihan yang lebih baik.

Metrik yang dicatat

Sebuah rollout melewati beberapa tahap sebelum menjadi gradien. Mari kita lihat terlebih dahulu berbagai tahap tersebut, dari prompt hingga komposisi batch akhir kita:

Dataset row = PROMPT (message list + optional `teacher_id`)
 └─ ROLLOUT            (1 prompt -> 1 student completion -> 1 teacher scoring call)
     ├─ generate       (student's vLLM /v1/completions, sampled)
     └─ score          (the routed teacher's /v1/completions with `prompt_logprobs`, teacher-forced)
         └─ SAMPLE     (prompt + completion + the teacher's sparse per-position distribution)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

         └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
             └─ ROW            (Planner assigns it to one of `dp` rows, Σ Lᵢ²-balanced)
                 └─ MICRO-BATCH (`dp` rows, one per rank)
                     └─ PACKED ROW (1 concat sequence, `position_ids` reset per sample)
                         └─ FORWARD (`compute_loss`, bs=1, inter-rank padding stripped)
             └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

Membacanya dari atas ke bawah:

• a rollout adalah satu prompt yang dihasilkan sekali dan diberi skor sekali. Tidak ada grup: distilasi tidak memiliki baseline keuntungan untuk dihitung antar generasi, sehingga prompt tidak diulang dan satu rollout menghasilkan tepat satu sampel pelatihan.
• a sample adalah completion tersebut ditambah kandidat top-teacher_top_k dari teacher untuk setiap posisi completion. Inilah yang melewati batas proses.
• melewati batas, trainer merencanakan sampel ke dalam baris (satu per rank DP), mengemas setiap baris menjadi satu rangkaian yang digabungkan, dan mengakumulasi grad_accum mikro-batch menjadi satu langkah optimizer.

Jadi, kami membagi metrik kami berdasarkan tahap (atau entitas) yang menjadi fokusnya:

Semua yang diukur oleh loss itu sendiri dicatat tanpa awalan (jsd, entropy, teacher_entropy), dengan rincian teacher_*/<teacher_id> di bawah MOPD — lihat divergence di bawah.

Mari kita juga mendefinisikan beberapa istilah agar metrik token tidak ambigu:

• token generated adalah yang dihasilkan oleh student (completions/*).
• Token forwarded adalah setiap token yang diproses oleh jalur forward: prompt dan token yang dihasilkan.
• token trained adalah subset tempat loss dihitung, yaitu di mana completion_mask == 1. Trained ≠ generated: posisi completion yang tidak memiliki kandidat yang diberi skor oleh teacher akan di-mask dari divergence, tetapi tetap diteruskan.

Bagaimana nomor batch disusun

Sebuah baris adalah apa yang satu rank DP teruskan dalam satu micro-batch: beberapa sampel yang digabungkan menjadi satu urutan, dengan position_ids dimulai ulang di setiap batas sampel. Perencana menentukan sampel mana yang masuk ke baris mana (menyeimbangkan Σ Lᵢ² agar tidak ada rank yang tertinggal), dan data collator melakukan penggabungan tersebut.

Pada semua metrik kami, “step” selalu berarti satu langkah optimizer PENUH, bukan micro-batch. gradient_accumulation_steps micro-batch membentuk satu step, dan setiap metrik bernama _per_step adalah per langkah optimizer, sesuai dengan global_step, logging_steps, dan sisa kosakata Trainer. Jika suatu metrik adalah per micro-batch, hal itu dinyatakan secara eksplisit (batch/microbatches_per_step) atau merupakan kuantitas per baris (batch/row_*, batch/samples_per_row).

Oleh karena itu, satu langkah memiliki jumlah row-slots yang tetap:

row-slots per step = gradient_accumulation_steps x world_size

yang membuat metrik batch tersusun menjadi tangga yang dapat Anda periksa satu sama lain:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

Jadi batch/samples_per_step ≈ row-slots x batch/samples_per_row, dan hal yang sama berlaku untuk token. Kedua sisi hanya cocok hingga variasi antar micro-batch di dalam satu step — metrik per-step adalah jumlah, sedangkan metrik per-row adalah rata-rata — sehingga diharapkan selisihnya sekitar seperseperseratus persen, bukan kecocokan yang persis.

batch/row_fill_frac adalah metrik yang perlu diperhatikan ketika sampelnya panjang: sampel 10k token mengisi anggaran 32k dengan buruk (tiga muat, empat tidak pernah muat, sehingga packer sering hanya mendapatkan dua dan barisnya terisi sekitar 77%), sedangkan sampel 1k token mengisinya hampir sempurna. Itu adalah kuantisasi, bukan bug, dan token_budget adalah tuasnya — dengan mempertimbangkan bahwa attention adalah O(L²) per sekuens, sehingga baris yang lebih penuh berisi sekuens panjang tidak menambah memori secara linear.

Antrean rollout

Worker mendorong sampel yang telah diberi skor ke dalam antrean, dan trainer menarik darinya. Empat metrik menggambarkan antrean tunggal tersebut, dan mereka menjawab pertanyaan yang berbeda:

Dua metrik terakhir (perf/rollout_wait_s dan rollout/backpressure_s) adalah cermin satu sama lain dan tidak pernah keduanya besar! Membacanya bersama ukuran antrean memberi tahu sisi mana yang menjadi bottleneck:

• antrean hampir kosong, perf/rollout_wait_s tinggi → generation-bound. Trainer sedang kelaparan; periksa rollout/generated_tok_s, rollout/concurrency dan rollout/score_s.
• antrean hampir penuh, rollout/backpressure_s tinggi → trainer-bound. Proses generasi dibatasi dan keluarannya menua di dalam antrean, sehingga perhatikan sample/staleness_mean yang naik.
• keduanya mendekati nol → seimbang!

Completions

completion adalah apa yang dihasilkan oleh server vLLM siswa untuk satu prompt.

Rollouts

Rollout adalah satu prompt yang diproses hingga tuntas: dihasilkan oleh student, kemudian di-teacher-force melalui teacher-nya. Karena panggilan teacher berada di jalur kritis setiap rollout (berbeda dengan GRPO, di mana scoring adalah loop terpisah atas grup yang telah selesai), teacher yang lambat akan langsung terlihat sebagai rollout/duration_s.

Contoh yang tiba dari antrean

Setiap sampel adalah satu RolloutSample: sebuah prompt, penyelesaian dari student, dan penilaian spars dari teacher terhadapnya. Worker mendorong sampel-sampel tersebut ke dalam rollout_buffer; semua yang di bawah ini diukur di sisi trainer, saat menarik sampel-sampel tersebut untuk pelatihan.

Batch

Satu micro-batch adalah world_size baris (ingat bahwa packing meratakan menjadi 1 sequence), sehingga satu per DP rank, sehingga rank i memproses baris i. gradient_accumulation_steps dari micro-batch tersebut membentuk satu langkah optimizer. Satu langkah kemudian mencakup total gradient_accumulation_steps × world_size baris. Setiap metrik _per_step di bawah ini adalah jumlahan atas seluruh langkah dan di semua rank; metrik batch/row_* adalah rata-rata atas baris-baris tersebut.

Divergensi

Apa yang diukur oleh objektif itu sendiri, dirata-ratakan di atas token yang dilatih dalam jendela. Inilah angka-angka yang perlu diperhatikan untuk pembelajaran, bukan untuk throughput.

Tidak ada entropy per-teacher: entropi siswa adalah properti dari kebijakan miliknya sendiri, bukan dari teacher mana yang memberi skor pada sampel, sehingga metrik gabungan sudah mencakupnya.

Kinerja

Throughput mengukur berapa banyak token yang diproses trainer per detik. Model FLOPs Utilization (MFU) memperkirakan persentase kapasitas komputasi teoretis perangkat pelatihan yang digunakan untuk jalur forward dan backward model. Kapasitas yang digunakan untuk estimasi ini bergantung pada perangkat keras dan presisi pelatihan.

Throughput dan MFU dilaporkan berdasarkan dua basis waktu:

• _fwd_bwd mengukur kinerja selama proses maju dan mundur.
• _wall_clock mengukur kinerja pada langkah optimizer lengkap, termasuk pembaruan optimizer, sinkronisasi bobot, dan waktu yang dihabiskan menunggu rollout.

Perbedaan ini membantu membedakan kinerja komputasi dari biaya menunggu dan koordinasi. Kedua metrik MFU hanya menggambarkan perangkat pelatihan: mereka tidak mencakup komputasi atau kapasitas perangkat keras dari server rollout atau teacher eksternal. Wall-clock MFU mencerminkan waktu yang dihabiskan trainer untuk menunggu layanan tersebut, bukan utilisasi layanan itu sendiri.

MFU hanya dilaporkan ketika kapasitas komputasi puncak diketahui untuk perangkat pelatihan lokal pada presisi yang dipilih. Jika tidak, MFU diabaikan; pelatihan tetap berlanjut dan metrik throughput serta waktu tetap tersedia.

Filsafat desain

Trainer ini sengaja dibuat minimal dan tidak dimaksudkan untuk berkembang menjadi solusi umum. Jika Anda membutuhkan fitur yang tidak didukung, kami merekomendasikan untuk mengkloning repositori dan menyesuaikan trainer sesuai kebutuhan Anda secara langsung. Fitur baru hanya akan dipertimbangkan ketika ada permintaan komunitas yang signifikan.

AsyncDistillationConfig

AsyncDistillationTrainer

>>> from trl.experimental.async_distillation import AsyncDistillationTrainer
>>> from datasets import load_dataset

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

>>> trainer = AsyncDistillationTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()