Lô liên tục

Continuous batching tối đa hóa việc sử dụng GPU bằng cách lên lịch lại động cho batch ở mỗi bước sinh. Khi các yêu cầu hoàn tất, các yêu cầu mới sẽ tham gia ngay lập tức thay vì chờ đợi toàn bộ batch hoàn thành. GPU luôn được lấp đầy và thông lượng duy trì ở mức cao.

Đối với triển khai trong môi trường sản xuất, hãy sử dụng transformers serve. Công cụ này được xây dựng dựa trên ContinuousBatchingManager và cung cấp một điểm cuối HTTP tương thích với OpenAI.

generate_batch

Continuous batching được hỗ trợ thông qua generate_batch(). Truyền vào một danh sách các prompt đã được token hóa và nhận lại kết quả cho tất cả chúng khi hoàn tất. generate_batch xử lý lịch trình bên trong và chặn cho đến khi tất cả các yêu cầu hoàn thành.

Đối với các trường hợp sử dụng phục vụ và phát trực tiếp, hãy sử dụng ContinuousBatchingManager trực tiếp để quản lý các yêu cầu.

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 chạy một luồng nền và cho phép bạn gửi yêu cầu cũng như lấy kết quả một cách độc lập. Ở mỗi bước sinh, nó kiểm tra các yêu cầu đã hoàn thành và lên lịch cho các yêu cầu mới tham gia vào batch. Điều này hữu ích cho streaming, phục vụ thời gian thực, hoặc gửi yêu cầu khi chúng đến.

Sử dụng continuous_batching_context_manager() để khởi động và dừng trình quản lý một cách an toàn. Ví dụ bên dưới chứa các đầu vào có độ dài khác nhau. Ngay khi prompt ngắn nhất hoàn tất, nó sẽ rời khỏi batch trong khi các prompt dài hơn tiếp tục tạo ra. Với static batching, bạn sẽ phải đệm tất cả chúng để có cùng độ dài. Continuous batching giải phóng prompt đã hoàn tất để bạn có thể bắt đầu xử lý prompt tiếp theo ngay lập tức.

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}")

Bạn cũng có thể gọi init_continuous_batching() để tự quản lý vòng đời.

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

# submit and retrieve requests...

Tắt trình quản lý

Quản lý chạy một luồng nền và giữ các tài nguyên phân tán. Việc tắt diễn ra theo hai giai đoạn để bạn có thể chọn cách xử lý công việc đang chạy.

Gọi stop() để dừng luồng nền. Mặc định, trình quản lý sẽ ngừng nhận các yêu cầu mới và chờ cho đến khi các yêu cầu đang xếp hàng và đang hoạt động hoàn tất trước khi luồng thoát.

manager.stop()

Truyền hard_stop=True để hủy bỏ công việc đang chờ ngay lập tức. Các yêu cầu trong hàng đợi và đang hoạt động sẽ bị lỗi với RuntimeError thay vì hoàn tất.

manager.stop(hard_stop=True)

Khi stop được gọi, add_request() và add_requests() sẽ từ chối các yêu cầu mới và ghi lại cảnh báo. Bạn vẫn có thể gọi start một lần nữa để chạy phiên sinh mới với cùng một trình quản lý.

Gọi destroy() để giải phóng các tài nguyên phân tán. destroy sẽ dừng manager trước nếu nó vẫn đang chạy, và manager không thể được khởi động lại sau đó. Hãy sử dụng nó khi bạn đã hoàn tất việc sử dụng continuous batching trong suốt vòng đời của tiến trình.

manager.destroy()

continuous_batching_context_manager() xử lý quy trình này. Nó gọi stop khi thoát và destroy trừ khi bạn truyền persistent_manager=True để lưu bộ nhớ đệm trình quản lý trên mô hình cho phiên tiếp theo.

Thêm các yêu cầu

add_request() gửi một yêu cầu đơn lẻ. Hãy cung cấp request_id hoặc để trình quản lý tự động tạo một cái.

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

add_requests() gửi một lô yêu cầu cùng lúc. Hàm này tự động sắp xếp các đầu vào để tối đa hóa số lần truy cập bộ nhớ đệm tiền tố khi chia sẻ khối được bật.

manager.add_requests(inputs=inputs)

Hủy một yêu cầu bằng cancel_request().

manager.cancel_request(request_id="my_request")

Tham số lấy mẫu cho từng yêu cầu

Bật per_request_processors để áp dụng temperature, top_k và top_p độc lập cho từng yêu cầu trong cùng một bước forward, cho phép sử dụng các tham số lấy mẫu khác nhau cho các yêu cầu khác nhau (ví dụ: đầu ra sáng tạo, nhiệt độ cao so với đầu ra chính xác, nhiệt độ thấp).

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)

Mỗi tham số trong GenerationConfig phải là một giá trị không mặc định để tạo ra logits processor tương ứng khi chạy. Ví dụ, đặt temperature thành một giá trị khác None hoặc 1 để hỗ trợ kiểm soát nhiệt độ theo từng yêu cầu. Các yêu cầu có nhiệt độ là 1 vẫn có thể được tạo ra sau đó.

Lấy kết quả

Lặp qua trình quản lý để nhận kết quả khi chúng đến.

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

get_result() lấy kết quả tiếp theo từ hàng đợi đầu ra. Truyền request_id để lọc cho một yêu cầu cụ thể. Nếu kết quả tiếp theo trong hàng đợi không khớp, nó sẽ được đưa lại vào hàng đợi và phương thức trả về None.

# next available result
result = manager.get_result()

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

Streaming

Đặt streaming=True cho một yêu cầu, sau đó sử dụng request_id_iter() để lặp qua các đầu ra một phần khi các token được tạo ra.

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 kiểm soát bộ đệm KV, lịch trình, đồ thị CUDA, mức sử dụng bộ nhớ và nhiều khía cạnh khác. Truyền nó vào generate_batch hoặc init_continuous_batching() bằng đối số continuous_batching_config, tách biệt với GenerationConfig. Hai cấu hình này mô tả những thứ khác nhau. GenerationConfig lấy mô hình làm trung tâm và chứa các tham số lấy mẫu và dừng, trong khi ContinuousBatchingConfig lấy phần cứng làm trung tâm và chứa các tham số bộ nhớ và lịch trình.

Việc đặt continuous_batching_config trên một GenerationConfig đã bị ngừng hỗ trợ, sẽ phát ra FutureWarning và sẽ bị loại bỏ trong v5.19.

Mặc định, max_batch_tokens là 8192, bị giới hạn bởi bộ nhớ GPU khả dụng và không bao giờ thấp hơn 256, trong khi num_blocks lấp đầy phần bộ nhớ còn lại. Hãy sử dụng bảng bên dưới để giúp bạn chọn các tính năng phù hợp.

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,
)

Kích thước trang bộ đệm KV

page_size xác định số lượng token KV mà một trang attention đầy đủ chứa cho một lớp. Giá trị này phải ít nhất là 4, nếu không quá trình khởi tạo sẽ gây ra ValueError. Giá trị mặc định 256 hoạt động tốt cho hầu hết các tác vụ.

Tham số block_size đã bị ngừng hỗ trợ và được thay thế bằng page_size.

Kích thước trang lớn hơn giúp giảm chi phí quản lý nhưng lại lãng phí không gian khi một chuỗi kết thúc giữa trang. Trang nhỏ hơn lãng phí ít không gian hơn nhưng đòi hỏi nhiều công việc quản lý hơn. Thông thường bạn không cần phải tinh chỉnh page_size. Nếu khối lượng công việc của bạn có nhiều yêu cầu ngắn, hãy thử một giá trị nhỏ hơn. Nếu có một vài yêu cầu rất dài, hãy thử một giá trị lớn hơn. Xem Quản lý bộ nhớ để biết cách các trang, khối và đoạn kết hợp với nhau.

cb_config = ContinuousBatchingConfig(page_size=128)

Kích thước lô prefill

max_batch_tokens đặt ngân sách token cho một lần forward pass, tức là số lượng token query tối đa mà mô hình xử lý cùng lúc. Ngân sách lớn hơn sẽ nhồi nhiều token prompt hơn vào mỗi lần prefill, từ đó tăng thông lượng prefill và giảm thời gian đến token đầu tiên trên các tác vụ nặng về prompt. Bộ lập lịch sẽ chia các prompt vượt quá ngân sách thành nhiều bước. Xem chunked prefill để biết cách xử lý các prompt quá lớn.

Mặc định, max_batch_tokens là 8192. Giá trị này bị giới hạn bởi bộ nhớ GPU khả dụng để các tensor đầu vào theo lô không chiếm hết bộ nhớ KV cache, và không bao giờ thấp hơn 256. Trên GPU có ít bộ nhớ trống, giới hạn này sẽ hạ giá trị xuống dưới 8192.

max_batch_tokens và num_blocks cùng lấy từ một ngân sách bộ nhớ, do đó chúng đánh đổi lẫn nhau. Ngân sách token lớn hơn sẽ phân bổ các bộ đệm đầu vào lớn hơn và để lại ít khối KV cache hơn cho các yêu cầu đồng thời hoặc dài. Tăng max_batch_tokens để đẩy cao thông lượng prefill khi bạn có thêm bộ nhớ khả dụng, và giảm nó để giải phóng bộ nhớ cho nhiều khối KV hơn hoặc để tránh lỗi hết bộ nhớ.

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

Giới hạn về lô và lịch trình

max_requests_per_batch giới hạn số lượng yêu cầu được thực thi trong một lần forward pass duy nhất. Mô hình tạo ra một tensor logits có kích thước bằng số token trong batch nhân với kích thước từ vựng, và tensor này được chuyển sang fp32 để lấy mẫu, khiến bộ nhớ sử dụng tăng gấp đôi. Mỗi yêu cầu chỉ cần một dự đoán token tiếp theo, do đó việc đặt giới hạn nhỏ hơn sẽ giữ tensor này ở kích thước nhỏ hơn và tránh lỗi hết bộ nhớ trên các batch nặng prefill với từ vựng lớn. Mặc định, giá trị này được đặt bằng số prompt được gửi, với giá trị dự phòng là 1024, sau đó được giới hạn để phù hợp với max_batch_tokens và num_blocks.

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin dành ra một phần dung lượng cache trống cho các yêu cầu đang hoạt động. Khi dung lượng trống giảm xuống dưới safety_margin (một phần của tổng kích thước cache), trình điều độ sẽ ngừng chấp nhận các prefill mới và chỉ tiếp tục giải mã các yêu cầu đang trong quá trình xử lý. Điều này ưu tiên hoàn thành công việc đang hoạt động hơn là bắt đầu các yêu cầu mới, giúp bảo vệ độ trễ giải mã và trì hoãn offloading. Giá trị phải nằm trong khoảng từ 0 đến 1, trong đó 0 vô hiệu hóa biên an toàn.

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

Giá trị mặc định phụ thuộc vào trình lập lịch. Với trình lập lịch FIFO, giá trị là 0.15 và với prefill_first, giá trị là 0.0. Xem Admission để biết cách biên độ tương tác với việc lập lịch.

Xác suất log

ContinuousBatchingConfig trả về log probability của từng token được sinh ra khi return_logprobs=True. Điều này hữu ích cho RL, nơi logprobs là đầu vào của một số vòng lặp huấn luyện.

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}")

Biểu đồ CUDA

CUDA graphs loại bỏ chi phí điều phối CPU bằng cách ghi lại đồ thị thực thi GPU một lần và phát lại nó cho các lô có hình dạng khớp. Bật chúng một cách rõ ràng bằng use_cuda_graph=True.

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

Khi đang hoạt động, trình quản lý sẽ đệm (pad) độ dài query và KV lên các khoảng cố định để các shape lặp lại và đồ thị được tái sử dụng. Các giá trị nhỏ hơn của q_padding_interval_size và kv_padding_interval_size sẽ giảm thiểu tính toán lãng phí do đệm, nhưng điều này đồng nghĩa với việc có nhiều shape duy nhất hơn mà đồ thị phải ghi nhận và lưu trữ, dẫn đến tốn nhiều bộ nhớ hơn.

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

Đóng gói đồng bộ

Async batching cho phép xếp chồng việc lên lịch CPU cho batch tiếp theo với việc tính toán GPU cho batch hiện tại. Tính năng này yêu cầu CUDA graphs và làm tăng gấp đôi lượng VRAM sử dụng cho các tensor đầu vào.

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

Biên dịch

default_compile_level áp dụng torch.compile cho các bước forward của model. Biên dịch đánh đổi một chi phí khởi động ban đầu để đổi lấy các bước forward nhanh hơn trong quá trình sinh. Các mức cao hơn thực hiện tối ưu hóa mạnh mẽ hơn, giúp cải thiện thông lượng nhưng kéo dài thời gian khởi động. Giữ mức thấp cho các bài kiểm tra và lặp lại benchmark, và tăng mức lên cho các tác vụ phục vụ chạy dài, nơi chi phí khởi động được bù đắp qua nhiều yêu cầu.

Mức độ dao động từ 0 đến 3. Mức 0 là mặc định và bỏ qua hoàn toàn việc biên dịch.

cb_config = ContinuousBatchingConfig(default_compile_level=1)

Mức này cung cấp CompileConfig mặc định cho các đường dẫn thực thi varlen và decode. Nó chỉ áp dụng cho đường dẫn không có cấu hình rõ ràng, do đó varlen_compile_config và decode_compile_config có độ ưu tiên cao hơn khi được đặt. Dưới FlashAttention, đường dẫn varlen bỏ qua việc biên dịch vì max_length_k gây ra việc biên dịch lại thường xuyên, nên trong trường hợp đó, mức này chỉ ảnh hưởng đến đường dẫn decode.

Đường dẫn giải mã nhanh

Khi một batch chỉ chứa các yêu cầu decode (một token truy vấn cho mỗi chuỗi), trình quản lý có thể điều phối đến kernel flash_attn_with_kvcache thay vì kernel biến độ dài. Cách này nhanh hơn đường dẫn varlen vì kernel đọc và ghi bộ nhớ đệm KV phân trang tại chỗ thông qua bảng khối, thay vì phải thực hiện cập nhật thủ công. Xem Paged attention để biết chi tiết ở cấp độ kernel.

Fast path được xác định kích thước bởi max_blocks_per_request, tham số này quyết định kích thước bảng khối (block table) cho mỗi yêu cầu. Mặc định, giá trị này được suy luận tự động. Nếu max_prompt_length và max_generated_length được đặt trên manager, bảng khối sẽ được xác định kích thước để chứa độ dài chuỗi tối đa. Ngược lại, một giá trị mặc định dự phòng (32 khối mỗi yêu cầu) sẽ được sử dụng.

Đặt max_blocks_per_request thành một giá trị cụ thể để xác định kích thước bảng khối một cách rõ ràng. Điều này hữu ích khi bạn biết độ dài tối đa của chuỗi cho mỗi yêu cầu và muốn giới hạn chi phí bộ nhớ của bảng khối.

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

Đặt max_blocks_per_request=0 để tắt fast path và buộc mọi batch đi qua varlen kernel. Thao tác này khôi phục hành vi mặc định trước đó và hữu ích khi fast path không khả dụng cho implementation attention của bạn (manager cũng tự động tắt nó khi kernel nền tảng không thể được sử dụng).

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

Đường dẫn nhanh dựa trên kernel flash_attn_with_kvcache, có sẵn cho hai tổ hợp thiết bị và triển khai attention.

Với bất kỳ tổ hợp nào khác, hoặc khi không thể import kernel, manager sẽ quay lại đường dẫn varlen. Nó chỉ ghi log cảnh báo khi bạn đặt max_blocks_per_request một cách tường minh.

Full attention hỗ trợ bảng khối kernel, nhưng sliding attention thì không. Nếu mô hình có bất kỳ lớp sliding attention nào, bao gồm cả các mô hình hybrid cũng sử dụng full attention, quá trình khởi tạo sẽ ép max_blocks_per_request thành 0, và giá trị khác 0 mà bạn đặt sẽ bị ghi đè. Sau đó, trình quản lý sẽ ghi log rằng đường nhanh decode không khả dụng vì một số loại attention của mô hình không hỗ trợ bảng khối kernel.

Tách bộ nhớ ra CPU

CPU offloading sao chép các khối KV cache bị loại bỏ vào bộ đệm CPU đã được cố định và phân bổ trước khi bộ nhớ KV cache trên GPU đầy. Sau khi có không gian cache khả dụng, trình quản lý sẽ sao chép các khối này trở lại GPU và tiếp tục yêu cầu mà không cần tính toán lại prompt và các token đã tạo.

Đặt cpu_offload_space thành dung lượng bộ nhớ đệm CPU tính bằng GiB. Giá trị mặc định, 0.0, sẽ tắt tính năng chuyển bộ nhớ ra CPU.

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

Mặc định, cpu_offload_space_safety_threshold=0.8 giới hạn dung lượng được yêu cầu ở mức 80% RAM hệ thống khả dụng khi psutil được cài đặt. Đặt cpu_offload_space=None để xác định kích thước pool swap từ ngưỡng an toàn.

Timeout song song tensor

Dưới chế độ song song tensor, trình quản lý tạo một nhóm giao tiếp CPU để điều phối việc gửi yêu cầu, hủy yêu cầu và tắt hệ thống giữa các rank. cpu_group_timeout giới hạn thời gian tối đa mà một thao tác tập thể trên nhóm này có thể bị chặn trước khi tiến trình bị sập. Nếu một rank bị treo, thời gian chờ quá hạn sẽ ngăn các rank khác chờ đợi vô hạn.

Đặt timeout dài hơn cho các tác vụ phát sinh các thao tác tập thể không thường xuyên, hoặc truyền None để tắt nó.

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

Bộ nhớ đệm tiền tố

Khi nhiều yêu cầu chia sẻ một tiền tố chung, chẳng hạn như một prompt hệ thống, trình quản lý sẽ tái sử dụng các khối KV cache của chúng thay vì tính toán lại. allow_block_sharing mặc định là True. Chia sẻ tiền tố chỉ được kích hoạt khi mọi bộ cấp phát kiểu chú ý đều hỗ trợ chia sẻ, do đó chỉ các mô hình chỉ dùng chú ý đầy đủ mới sử dụng tính năng này. Các mô hình trượt thuần túy và mô hình lai sẽ tắt tính năng này ngay cả khi allow_block_sharing là True.

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

Chọn một backend chú ý

Continuous batching hỗ trợ FlashAttention, SDPA và "paged|eager". Đặt attn_implementation khi bạn tải model.

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

Continuous batching chạy nhanh hơn với FlashAttention so với eager hoặc SDPA, chủ yếu vì FlashAttention không cần attention mask. Nếu model hỗ trợ FlashAttention, một model được nạp với "eager" hoặc "sdpa" sẽ chuyển sang FlashAttention 3, hoặc sang FlashAttention 2 khi FlashAttention 3 chưa được cài đặt, và ghi lại một cảnh báo. Nếu cả hai phiên bản đều chưa được cài đặt, backend sẽ không thay đổi.

Để giữ nguyên eager hoặc SDPA, chẳng hạn khi kiểm thử, hãy đặt auto_switch_to_flash=False trong ContinuousBatchingConfig. Mỗi khi backend vẫn là "eager", continuous batching sẽ chạy "paged|eager" thay thế.

Song song tensor

Đối với các mô hình quá lớn để chạy trên một GPU duy nhất, hãy chia nhỏ trọng số giữa các thiết bị bằng tensor parallelism. Đặt số lượng thiết bị bằng DistributedConfig(tp_size=N). Continuous batching đọc kích thước tensor parallel từ mô hình để xác định kích thước paged KV cache cho từng shard. Xem Tensor parallelism để biết danh sách các kiến trúc được hỗ trợ và cách thức chia nhỏ hoạt động.

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)

Chạy script bằng torchrun, đặt --nproc-per-node bằng số lượng GPU mà bạn muốn chia nhỏ (shard) dữ liệu.

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

Kích thước song song tensor phải chia hết cho num_key_value_heads của mô hình (kiểm tra cấu hình mô hình). Nếu không, bộ đệm trang sẽ gây ra lỗi khi khởi động, vì vậy hãy chọn --nproc-per-node phù hợp.

Không đặt device_map cùng với distributed_config. Hai tham số này xung đột vì device_map đặt các module hoàn chỉnh lên các GPU cụ thể, trong khi tensor parallelism chia nhỏ chính các tham số đó trên tất cả các GPU.

Chú ý cửa sổ trượt

Các mô hình có chú ý cửa sổ trượt (Mistral, Gemma 2) hoạt động với continuous batching. Các lớp trượt sử dụng SlidingAttentionCacheAllocator, chỉ giữ một cửa sổ cố định các token gần nhất bằng cách tái sử dụng một tập hợp nhỏ các trang. Loại chú ý được lấy từ layer_types hoặc sliding_window của mô hình, không phải từ 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 và decode fast path sẽ bị vô hiệu hóa khi có sliding attention, bao gồm cả các mô hình hybrid có thêm các lớp full attention.

Các bước tiếp theo

• Bài đăng trên blog Continuous batching đề cập đến KV caching, chunked prefill và dynamic scheduling kèm theo các con số benchmark hiệu năng.
• Để tìm hiểu sâu hơn về cách hệ thống liên tục hoạt động, hãy xem tài liệu Continuous batching architecture.