Transformers documentation

Batching berkelanjutan

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Batching berkelanjutan

Continuous batching memaksimalkan pemanfaatan GPU dengan menjadwalkan ulang batch secara dinamis di setiap langkah generasi. Saat permintaan selesai, permintaan baru langsung bergabung tanpa menunggu seluruh batch selesai. GPU tetap penuh dan throughput tetap tinggi.

Untuk deployment produksi, gunakan transformers serve. Fitur ini dibangun di atas ContinuousBatchingManager dan mengekspos endpoint HTTP yang kompatibel dengan OpenAI.

generate_batch

Continuous batching didukung melalui generate_batch(). Berikan daftar prompt yang telah ditokenisasi dan Anda akan menerima hasil untuk semuanya saat selesai. generate_batch menangani penjadwalan secara internal dan memblokir hingga semua permintaan selesai.

Untuk kasus penggunaan penyajian dan streaming, gunakan ContinuousBatchingManager secara langsung untuk mengelola permintaan.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

prompts = [
    "What's up?",
    "Name a cat breed.",
    "Write a detailed history of quantum mechanics.",
]
inputs = [tokenizer.encode(p) for p in prompts]

generation_config = GenerationConfig(
    max_new_tokens=64,
    eos_token_id=tokenizer.eos_token_id,
)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

for request_id, output in outputs.items():
    text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
    print(f"[{request_id}] {text}")

ContinuousBatchingManager

ContinuousBatchingManager menjalankan thread latar belakang dan memungkinkan Anda mengirim permintaan serta mengambil hasil secara independen. Pada setiap langkah generasi, ia memeriksa permintaan yang telah selesai dan menjadwalkan permintaan baru untuk bergabung ke dalam batch. Hal ini berguna untuk streaming, penyajian real-time, atau pengiriman permintaan saat mereka tiba.

Gunakan continuous_batching_context_manager() untuk memulai dan menghentikan manager dengan aman. Contoh di bawah ini berisi input dengan panjang variabel. Begitu prompt terpendek selesai, ia keluar dari batch sementara prompt yang lebih panjang terus menghasilkan. Dengan static batching, Anda harus melakukan padding pada semuanya hingga panjang yang sama. Continuous batching membebaskan prompt yang sudah selesai sehingga Anda dapat segera mulai memproses prompt berikutnya.

with model.continuous_batching_context_manager(generation_config=generation_config) as manager:
    manager.add_request(
        input_ids=tokenizer.encode("Write a detailed history of quantum mechanics."),
        request_id="long",
        max_new_tokens=512,
    )
    manager.add_request(
        input_ids=tokenizer.encode("What's up?"),
        request_id="short_0",
        max_new_tokens=32,
    )
    manager.add_request(
        input_ids=tokenizer.encode("Name a cat breed."),
        request_id="short_1",
        max_new_tokens=32,
    )

    for result in manager:
        text = tokenizer.decode(result.generated_tokens, skip_special_tokens=True)
        print(f"[{result.request_id}] {text}")

Anda juga dapat memanggil init_continuous_batching() untuk mengelola siklus hidupnya sendiri.

manager = model.init_continuous_batching(generation_config=generation_config)
manager.start()

# submit and retrieve requests...

Menutup manager

Manajer menjalankan thread latar belakang dan memegang sumber daya terdistribusi. Penutupan terjadi dalam dua tahap sehingga Anda dapat memilih apa yang harus dilakukan dengan pekerjaan yang sedang berjalan.

Panggil stop() untuk menghentikan thread latar belakang. Secara default, manager berhenti menerima pengiriman baru dan menunggu hingga permintaan yang antre dan aktif selesai sebelum thread keluar.

manager.stop()

Lewati hard_stop=True untuk membatalkan pekerjaan yang tertunda secara langsung. Permintaan yang berada dalam antrean dan yang sedang aktif akan gagal dengan RuntimeError alih-alih diselesaikan.

manager.stop(hard_stop=True)

Setelah stop dipanggil, add_request() dan add_requests() akan menolak pengiriman baru dan mencatat peringatan. Anda masih dapat memanggil start lagi untuk menjalankan sesi generasi lain dengan manajer yang sama.

Panggil destroy() untuk melepaskan sumber daya terdistribusi. destroy menghentikan manager terlebih dahulu jika masih berjalan, dan manager tidak dapat dimulai ulang setelahnya. Gunakan ini ketika Anda selesai menggunakan continuous batching untuk sisa masa hidup proses.

manager.destroy()

continuous_batching_context_manager() menangani proses ini. Fungsi ini memanggil stop saat keluar dan destroy kecuali Anda meneruskan persistent_manager=True untuk meng-cache manager pada model untuk sesi berikutnya.

Menambahkan permintaan

add_request() mengirimkan satu permintaan. Berikan request_id atau biarkan manajer membuatnya secara otomatis.

manager.add_request(input_ids=input_ids, request_id="my_request")

add_requests() mengirim batch sekaligus. Fungsi ini mengurutkan input secara otomatis untuk memaksimalkan hit cache prefix ketika block sharing diaktifkan.

manager.add_requests(inputs=inputs)

Batalkan permintaan dengan cancel_request().

manager.cancel_request(request_id="my_request")

Parameter sampling per permintaan

Aktifkan per_request_processors untuk menerapkan temperature, top_k, dan top_p secara independen per permintaan dalam satu forward pass yang sama, sehingga memungkinkan parameter sampling yang berbeda untuk permintaan yang berbeda (misalnya, output kreatif dengan suhu tinggi versus output presisi dengan suhu rendah).

cb_config = ContinuousBatchingConfig(per_request_processors=True)

# each request gets its own sampling parameters
manager.add_request(input_ids=inputs_a, temperature=0.9, top_p=0.95)
manager.add_request(input_ids=inputs_b, temperature=0.1, top_k=10)

Setiap parameter dalam GenerationConfig harus berupa nilai non-default agar logits processor terkait dapat dibuat saat runtime. Sebagai contoh, atur temperature ke nilai selain None atau 1 untuk mendukung kontrol suhu per permintaan. Permintaan dengan suhu 1 tetap dapat dibuat setelahnya.

Mengambil hasil

Iterasi pada manager untuk menerima hasil saat tiba.

for result in manager:
    print(tokenizer.decode(result.generated_tokens, skip_special_tokens=True))

get_result() mengambil hasil berikutnya dari antrean keluaran. Sampaikan request_id untuk memfilter permintaan tertentu. Jika hasil berikutnya dalam antrean tidak cocok, hasil tersebut diproses ulang dan metode mengembalikan None.

# next available result
result = manager.get_result()

# filter for a specific request
result = manager.get_result(request_id="my_request")

Streaming

Setel streaming=True pada permintaan, lalu gunakan request_id_iter() untuk mengiterasi keluaran parsial saat token dihasilkan.

from transformers.generation.continuous_batching import RequestStatus

manager.add_request(input_ids=input_ids, request_id="streamed", streaming=True)

for chunk in manager.request_id_iter(request_id="streamed"):
    token = tokenizer.decode(chunk.generated_tokens[-1:], skip_special_tokens=True)
    print(token, end="", flush=True)
    if chunk.status == RequestStatus.FINISHED:
        break

ContinuousBatchingConfig

ContinuousBatchingConfig mengontrol KV cache, penjadwalan, CUDA graphs, penggunaan memori, dan lainnya. Sampaikan ke generate_batch atau init_continuous_batching() dengan argumen continuous_batching_config, terpisah dari GenerationConfig. Kedua konfigurasi ini mendeskripsikan hal yang berbeda. GenerationConfig berpusat pada model dan menyimpan parameter sampling serta penghentian, sedangkan ContinuousBatchingConfig berpusat pada perangkat keras dan menyimpan parameter memori serta penjadwalan.

Pengaturan continuous_batching_config pada GenerationConfig sudah tidak didukung lagi, mengeluarkan FutureWarning, dan akan dihapus pada v5.19.

Secara default, max_batch_tokens adalah 8192, dibatasi oleh memori GPU yang tersedia dan tidak pernah di bawah 256, sementara num_blocks mengisi sisa memori. Gunakan tabel di bawah untuk membantu Anda memilih fitur yang tepat.

FiturMemoriThroughputLatensi
max_memory_percent / page_size✓ mengontrol anggaran KV
max_batch_tokens↑ buffer input yang lebih besar✓ batch prefill yang lebih besar✓ TTFT saat dibatasi oleh prefill
scheduler✓ kebijakan penjadwalan✓ TTFT
Graf CUDA↑ penyimpanan graf✓ overhead dispatch yang lebih rendah✓
Batching asinkron↑ ~2× buffer I/O✓ tumpang tindih CPU/GPU
Kompilasi↑ hanya waktu pemanasan✓ jalur maju yang lebih cepat✓
Jalur cepat decode↑ tabel blok per permintaan✓ langkah decode-only yang lebih cepat✓
Offloading CPU↑ memori CPU yang di-*pin*✓ melewati beberapa re-prefill
Penyimpanan sementara awalan↓ blok KV bersama✓ melewati prefill yang redundan✓ TTFT
Perhatian berhalaman↓ tidak ada fragmentasi✓ keanggotaan batch dinamis
Jendela geser↓ KV terbatas per lapisan
Pemroses per permintaan✓ parameter sampling campuran per batch
max_requests_per_batch↓ membatasi buffer logits✓ membatasi ukuran batch
safety_margin✓ melindungi latensi decode
from transformers.generation import ContinuousBatchingConfig

cb_config = ContinuousBatchingConfig(
    max_memory_percent=0.8,
    page_size=256,           # KV cache page size in tokens
    scheduler_type="fifo",        # "fifo" or "prefill_first"
)

outputs = model.generate_batch(
    inputs=inputs,
    generation_config=generation_config,
    continuous_batching_config=cb_config,
)

Ukuran halaman cache KV

page_size mengatur berapa banyak token KV yang ditampung oleh satu halaman perhatian penuh untuk satu lapisan. Nilainya harus minimal 4, atau inisialisasi akan memicu ValueError. Nilai default 256 cocok untuk sebagian besar beban kerja.

Parameter block_size sudah tidak digunakan lagi dan diganti oleh page_size.

Ukuran halaman yang lebih besar memangkas pencatatan, tetapi membuang-buang ruang ketika sebuah sekuens berhenti di tengah halaman. Halaman yang lebih kecil membuang-buang ruang lebih sedikit, tetapi membutuhkan lebih banyak pencatatan. Anda biasanya tidak perlu menyetel page_size. Jika beban kerja Anda memiliki banyak permintaan pendek, coba nilai yang lebih kecil. Jika memiliki beberapa permintaan yang sangat panjang, coba nilai yang lebih besar. Lihat Manajemen memori untuk memahami bagaimana halaman, blok, dan potongan saling terhubung.

cb_config = ContinuousBatchingConfig(page_size=128)

Ukuran batch prefill

max_batch_tokens mengatur anggaran token untuk satu langkah forward, yaitu jumlah maksimum token query yang diproses model sekaligus. Anggaran yang lebih besar memuat lebih banyak token prompt ke dalam setiap prefill, yang meningkatkan throughput prefill dan menurunkan time-to-first-token pada workload yang berat prompt. Scheduler membagi prompt yang melebihi anggaran ke dalam beberapa langkah. Lihat chunked prefill untuk cara penanganan prompt yang melebihi ukuran.

Secara default, max_batch_tokens adalah 8192. Nilainya dibatasi oleh memori GPU yang tersedia agar tensor input per-batch tidak menggeser KV cache, dan nilainya tidak pernah turun di bawah 256. Pada GPU dengan memori kosong yang sedikit, batas tersebut akan menurunkannya di bawah 8192.

max_batch_tokens dan num_blocks mengambil dari anggaran memori yang sama, sehingga keduanya saling bertolak belakang. Anggaran token yang lebih besar mengalokasikan buffer input yang lebih besar dan menyisakan lebih sedikit blok KV cache untuk permintaan konkuren atau panjang. Naikkan max_batch_tokens untuk meningkatkan throughput prefill ketika Anda memiliki lebih banyak memori yang tersedia, dan turunkan untuk membebaskan memori bagi lebih banyak blok KV atau untuk menghindari error out-of-memory.

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

Batas batch dan penjadwalan

max_requests_per_batch membatasi jumlah permintaan yang dijalankan dalam satu forward pass. Model menghasilkan tensor logits yang ukurannya ditentukan oleh jumlah token batch dan ukuran kosakata, lalu dicast ke fp32 untuk sampling, yang menggandakan footprint. Setiap permintaan hanya membutuhkan satu prediksi token berikutnya, sehingga batas yang lebih kecil menjaga tensor ini tetap kecil dan menghindari error out-of-memory pada batch yang berat prefill dengan kosakata besar. Secara default, nilainya diatur ke jumlah prompt yang dikirim, dengan fallback 1024, lalu dibatasi agar muat dalam max_batch_tokens dan num_blocks.

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin menyisakan sebagian dari kapasitas cache yang bebas untuk permintaan aktif. Ketika kapasitas bebas turun di bawah safety_margin (sebagian dari ukuran total cache), penjadwal berhenti menerima prefill baru dan hanya melanjutkan permintaan decoding yang sedang berlangsung. Hal ini memprioritaskan penyelesaian pekerjaan aktif di atas memulai permintaan baru, yang melindungi latensi decoding dan menunda offloading. Nilainya harus berada di antara 0 dan 1, di mana 0 menonaktifkan margin.

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

Nilai default bergantung pada scheduler. Untuk scheduler FIFO, nilainya adalah 0.15 dan untuk prefill_first, nilainya adalah 0.0. Lihat Admission untuk memahami bagaimana margin berinteraksi dengan penjadwalan.

Probabilitas log

ContinuousBatchingConfig mengembalikan log probabilitas dari setiap token yang dihasilkan ketika return_logprobs=True. Hal ini berguna untuk RL di mana logprobs merupakan input bagi beberapa loop pelatihan.

cb_config = ContinuousBatchingConfig(return_logprobs=True)

# generate_batch()

for request_id, output in outputs.items():
    for token_id, log_prob in zip(output.generated_tokens, output.logprobs):
        token = tokenizer.decode([token_id])
        print(f"{token} | logprob: {log_prob}")

Graf CUDA

Graf CUDA menghilangkan overhead dispatch CPU dengan merekam graf eksekusi GPU sekali dan memutar ulang untuk batch dengan bentuk yang cocok. Aktifkan secara eksplisit dengan use_cuda_graph=True.

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

Saat aktif, manager mengisi panjang query dan KV hingga interval tetap sehingga bentuk berulang dan grafik dapat digunakan kembali. Nilai yang lebih kecil untuk q_padding_interval_size dan kv_padding_interval_size mengurangi komputasi yang terbuang untuk padding, tetapi ini berarti ada lebih banyak bentuk unik yang harus direkam dan disimpan oleh grafik, yang membutuhkan lebih banyak memori.

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    q_padding_interval_size=64,
    kv_padding_interval_size=16384,
)

Batching asinkron

Async batching menumpuk penjadwalan CPU untuk batch berikutnya dengan komputasi GPU untuk batch saat ini. Fitur ini memerlukan CUDA graphs dan secara kasar menggandakan penggunaan VRAM untuk tensor input.

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    use_async_batching=True,
)

Kompilasi

default_compile_level menerapkan torch.compile pada jalur forward model. Kompilasi menukar biaya warmup sekali jalan dengan jalur forward yang lebih cepat selama generasi. Level yang lebih tinggi menjalankan optimasi yang lebih agresif, yang meningkatkan throughput tetapi memperpanjang waktu warmup. Pertahankan level rendah untuk pengujian dan iterasi benchmark, dan tingkatkan untuk workload serving jangka panjang di mana biaya warmup terbayar kembali melalui banyak permintaan.

Tingkat berkisar dari 0 hingga 3. Tingkat 0 adalah nilai default dan sepenuhnya melewati kompilasi.

LevelmodedynamicKompromi
0——Tidak ada kompilasi (default), startup tercepat
1defaultTruePeningkatan kecepatan moderat, pemanasan singkat
2max-autotune-no-cudagraphsTruePercepatan lebih besar, pemanasan lebih lama
3max-autotune-no-cudagraphsFalseThroughput terbaik, warmup terlama
cb_config = ContinuousBatchingConfig(default_compile_level=1)

Level menyediakan CompileConfig default untuk jalur eksekusi varlen dan decode. Ini hanya berlaku untuk jalur yang tidak memiliki konfigurasi eksplisit, sehingga varlen_compile_config dan decode_compile_config memiliki prioritas ketika diset. Di bawah FlashAttention, jalur varlen melewati kompilasi karena max_length_k memicu kompilasi ulang yang sering, sehingga level hanya memengaruhi jalur decode dalam kasus tersebut.

Jalur cepat decode

Ketika sebuah batch hanya berisi permintaan decode (satu token query per sekuens), manajer dapat mendispatch ke kernel flash_attn_with_kvcache alih-alih kernel panjang variabel. Ini lebih cepat daripada jalur varlen karena kernel membaca dan menulis cache KV berhalaman secara in-place melalui tabel blok, bukan melalui pembaruan manual. Lihat Paged attention untuk detail tingkat kernel.

Fast path diukur berdasarkan max_blocks_per_request, yang menentukan dimensi tabel blok per permintaan. Secara default, nilai ini disimpulkan secara otomatis. Jika max_prompt_length dan max_generated_length diset pada manager, tabel blok diukur agar muat dengan panjang sekuens maksimum. Jika tidak, nilai default cadangan (32 blok per permintaan) digunakan.

Set max_blocks_per_request ke nilai tertentu untuk menentukan ukuran tabel blok secara eksplisit. Ini berguna ketika Anda mengetahui panjang maksimum urutan per permintaan dan ingin membatasi biaya memori tabel blok.

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

Set max_blocks_per_request=0 untuk menonaktifkan jalur cepat dan memaksa setiap batch melewati kernel varlen. Ini memulihkan perilaku sebelum default dan berguna ketika jalur cepat tidak tersedia untuk implementasi attention Anda (manajer juga menonaktifkannya secara otomatis ketika kernel yang mendasarinya tidak dapat digunakan).

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

Jalur cepat bergantung pada kernel flash_attn_with_kvcache, yang tersedia untuk dua kombinasi perangkat dan implementasi perhatian.

Perangkatattn_implementation
CUDAflash_attention_2 atau flash_attention_3
XPUflash_attention_2

Untuk kombinasi lain, atau ketika kernel tidak dapat diimpor, manager kembali ke jalur varlen. Peringatan hanya dicatat ketika Anda mengatur max_blocks_per_request secara eksplisit.

Full attention mendukung kernel block tables, tetapi sliding attention tidak. Jika model memiliki lapisan sliding attention, termasuk model hybrid yang juga menggunakan full attention, inisialisasi memaksa max_blocks_per_request menjadi 0, dan nilai non-nol yang Anda atur akan ditimpa. Manager kemudian mencatat bahwa jalur cepat decode tidak tersedia karena beberapa jenis perhatian model tidak mendukung kernel block tables.

Offloading CPU

CPU offloading menyalin blok KV cache yang dieviksi ke buffer CPU pinned yang dialokasikan sebelumnya ketika KV cache GPU penuh. Setelah ruang cache tersedia, manajer menyalin blok-blok tersebut kembali ke GPU dan melanjutkan permintaan tanpa menghitung ulang prompt dan token yang dihasilkan.

Set cpu_offload_space ke ruang tukar CPU dalam GiB. Nilai default, 0.0, menonaktifkan offloading CPU.

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

Secara default, cpu_offload_space_safety_threshold=0.8 membatasi ruang yang diminta menjadi 80% dari RAM sistem yang tersedia ketika psutil terpasang. Atur cpu_offload_space=None untuk menentukan ukuran pool swap berdasarkan ambang batas keamanan.

Timeout paralel tensor

Di bawah tensor parallelism, manager membuat grup komunikasi CPU untuk mengoordinasikan pengiriman, pembatalan, dan shutdown permintaan antar rank. cpu_group_timeout membatasi berapa lama sebuah collective pada grup ini dapat memblokir sebelum proses crash. Jika satu rank macet, timeout mencegah rank lainnya menunggu selamanya.

Atur timeout yang lebih panjang untuk beban kerja yang mengeluarkan kolektif secara jarang, atau lewati None untuk menonaktifkannya.

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

Penyimpanan sementara awalan

Ketika beberapa permintaan berbagi awalan yang sama, seperti prompt sistem, manajer memblokir ulang blok cache KV mereka alih-alih menghitung ulang. allow_block_sharing secara default bernilai True. Berbagi awalan hanya aktif ketika setiap alokator tipe perhatian mendukung berbagi, sehingga hanya model yang sepenuhnya berbasis perhatian penuh yang menggunakannya. Model sliding murni dan hibrida menonaktifkannya meskipun allow_block_sharing bernilai True.

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

Pilih backend perhatian

Continuous batching mendukung FlashAttention, SDPA, dan "paged|eager". Atur attn_implementation saat Anda memuat model.

Backendattn_implementationPersyaratan
FlashAttention 2“flash_attention_2”paket flash-attn atau kernels
FlashAttention 3“flash_attention_3”paket flash-attn-3 atau kernels
SDPA (PyTorch asli)“sdpa”None
Eager“paged|eager”None
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)

Continuous batching berjalan lebih cepat dengan FlashAttention dibandingkan dengan eager atau SDPA, terutama karena FlashAttention tidak memerlukan attention mask. Jika model mendukung FlashAttention, model yang dimuat dengan "eager" atau "sdpa" akan beralih ke FlashAttention 3, atau ke FlashAttention 2 jika FlashAttention 3 tidak terpasang, dan mencatat peringatan. Jika kedua versi tersebut tidak terpasang, backend tidak berubah.

Untuk mempertahankan eager atau SDPA, misalnya saat pengujian, atur auto_switch_to_flash=False di ContinuousBatchingConfig. Setiap kali backend tetap "eager", continuous batching menjalankan "paged|eager" sebagai gantinya.

Paralelisme tensor

Untuk model yang terlalu besar untuk muat di satu GPU, pecah bobotnya ke beberapa perangkat dengan tensor parallelism. Atur jumlah perangkat dengan DistributedConfig(tp_size=N). Continuous batching membaca ukuran tensor parallel dari model untuk menentukan ukuran paged KV cache per shard. Lihat Tensor parallelism untuk daftar arsitektur yang didukung dan cara kerja sharding.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",
    attn_implementation="flash_attention_2",
    distributed_config=distributed_config,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-32B")

inputs = [tokenizer.encode(p) for p in ["What's up?", "Name a cat breed."]]
generation_config = GenerationConfig(max_new_tokens=64, eos_token_id=tokenizer.eos_token_id)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

Jalankan skrip dengan torchrun, atur --nproc-per-node ke jumlah GPU yang ingin Anda shard.

torchrun --nproc-per-node 4 cb_tp.py

Ukuran paralel tensor harus membagi num_key_value_heads model (periksa konfigurasi model). Cache berhalaman akan memicu error saat startup jika tidak, sehingga pilih --nproc-per-node yang sesuai.

Jangan mengatur device_map bersamaan dengan distributed_config. Keduanya saling bertentangan karena device_map menempatkan modul utuh pada GPU tertentu, sedangkan tensor parallelism memecah parameter yang sama tersebut di seluruh GPU.

Perhatian jendela geser

Model dengan sliding window attention (Mistral, Gemma 2) bekerja dengan continuous batching. Lapisan sliding menggunakan SlidingAttentionCacheAllocator, yang hanya menyimpan jendela tetap dari token-token terbaru dengan mendaur ulang sejumlah kecil halaman. Jenis attention berasal dari layer_types atau sliding_window model, bukan dari ContinuousBatchingConfig.

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("google/gemma-2-2b")
config.sliding_window = 4096

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b",
    config=config,
    attn_implementation="sdpa",
    device_map="auto",
    dtype=torch.bfloat16,
)

Prefix caching dan decode fast path dinonaktifkan setiap kali sliding attention hadir, termasuk model hybrid yang juga memiliki lapisan full attention.

Langkah selanjutnya

Update on GitHub