
Batching by Length Instead of Looping Item by Item for SLM Optimization
We finish off our short series on SLM optimization with the third entry, focused on batching by length instead of looping item by item.
以下正文同步自 KDnuggets,版权归原站所有,已转换为易读排版。
Previous articles in this series discussed constraining output space as well as reusing the prompt prefix with a key-value cache, both framed as approaches to small language model (SLM) narrow automation optimization. Let's finish this series off with the third entry, focused on batching by length instead of looping item by item.
As in our previous articles, all benchmarks below use Qwen2.5-0.5B-Instruct in float16 through Hugging Face Transformers, running on an M2 Macbook Air with 24GB RAM and a 16-core Neural Engine.
Don't forget to set up a Python environment and install your requirements:
pip install torch transformers accelerate
We will continue to use the support ticket framing from our first article.
Processing one ticket per forward pass is the single largest source of waste in the whole pipeline. At batch size 1, a small model is memory-bandwidth bound rather than compute-bound: the hardware streams every weight out of memory in order to serve one sequence, then does it again for the next one, and the arithmetic units sit mostly idle in between. This is true on a GPU and it is true on the CPU we have been using throughout this series, which is where a 0.5B model most often actually runs.
Batching amortizes that weight read across many sequences. But the obvious implementation introduces its own waste, since sequences in a batch must be padded to a common length. Real-world text has a long tail: if the longest item in your dataset is a few hundred tokens and the median is well under a hundred, padding every batch to the global maximum means most of what you compute is padding.
The answer is to sort by token length before forming batches, so each batch contains similarly sized items and pads to its own local maximum.
Here is the per-item baseline on a realistic length distribution, with the constrained scoring from the first article in this series carried forward so that each item costs exactly one forward pass:
import os import time import inspect import torch import numpy as np from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
torch.set_num_threads(os.cpu_count() or 1)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" # keeps the last real token at index -1
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32) model.eval()
LABELS = ["billing", "technical", "account"]
constrained scoring, carried over from the first article in this series
label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS] assert len(set(label_first_ids)) == len(LABELS), ( "Labels share a first token; score full label sequences instead." ) label_first_ids = torch.tensor(label_first_ids, device=model.device)
a causal LM returns a logit vector for every position by default; at batch 32 by
400 tokens that is a multi-gigabyte tensor we would immediately throw away, so
ask for the last position only where the installed version supports it
_forward_params = inspect.signature(model.forward).parameters if "logits_to_keep" in _forward_params: LAST_LOGIT_ONLY = {"logits_to_keep": 1} elif "num_logits_to_keep" in _forward_params: LAST_LOGIT_ONLY = {"num_logits_to_keep": 1} else: LAST_LOGIT_ONLY = {}
def build_prompt(ticket): messages = [ { "role": "system", "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.", }, {"role": "user", "content": f"Ticket: {ticket}\nCategory:"}, ] return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
simulate a long-tailed ticket length distribution: most items are short, a few
are very long. Each ticket keeps a real, classifiable sentence at the front and
is extended with filler, so length varies without the label signal disappearing.
BASE_TICKETS = [ "My card was charged twice for the same invoice.", "The mobile app crashes whenever I open the settings page.", "I need to change the email address on my profile.", ] FILLER = ( "I have been waiting for a response for several days now and would really " "appreciate an update on this whenever someone gets a chance to look at it." ).split()
rng = np.random.default_rng(0) target_words = np.clip(rng.lognormal(np.log(60), 0.9, size=600), 12, 400).astype(int)
def make_ticket(base, n_words): words = base.split() while len(words)
Output:
Prompt lengths: min 48, median 94, max 449 tokens Padding every item to the global maximum would process 3.7x the necessary tokens 100/600 tickets (0.24s each) 200/600 tickets (0.23s each) 300/600 tickets (0.23s each) 400/600 tickets (0.23s each) 500/600 tickets (0.24s each) 600/600 tickets (0.24s each) One at a time: 144.35 seconds (4.2 items/sec)
Batching by Length
Now the batched version. It runs the same data twice: once in arbitrary order, to isolate what batching alone is worth, and once sorted by length, to show what the sorting adds on top. Both runs track how much of the processed token budget went to padding.
import os import time import inspect import torch import numpy as np from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct" BATCH_SIZE = 32
torch.set_num_threads(os.cpu_count() or 1)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
keep the last real token at index -1
tokenizer.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32) model.eval()
LABELS = ["billing", "technical", "account"] label_first_ids = [tokenizer.encode(label, add_special_tokens=False)[0] for label in LABELS] assert len(set(label_first_ids)) == len(LABELS), ( "Labels share a first token; score full label sequences instead." ) label_first_ids = torch.tensor(label_first_ids, device=model.device)
_forward_params = inspect.signature(model.forward).parameters if "logits_to_keep" in _forward_params: LAST_LOGIT_ONLY = {"logits_to_keep": 1} elif "num_logits_to_keep" in _forward_params: LAST_LOGIT_ONLY = {"num_logits_to_keep": 1} else: LAST_LOGIT_ONLY = {}
def build_prompt(ticket): messages = [ { "role": "system", "content": "You classify support tickets. Answer with exactly one of: billing, technical, account.", }, {"role": "user", "content": f"Ticket: {ticket}\nCategory:"}, ] return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
BASE_TICKETS = [ "My card was charged twice for the same invoice.", "The mobile app crashes whenever I open the settings page.", "I need to change the email address on my profile.", ] FILLER = ( "I have been waiting for a response for several days now and would really " "appreciate an update on this whenever someone gets a chance to look at it." ).split()
rng = np.random.default_rng(0) target_words = np.clip(rng.lognormal(np.log(60), 0.9, size=600), 12, 400).astype(int)
def make_ticket(base, n_words): words = base.split() while len(words)
Output:
Loading weights: 100%|████████████████████████████████████████████████████████████████████| 290/290 [00:01
A large throughput increase on identical hardware and an identical model, purely from how the work was scheduled. Note that the two batched runs do the same arithmetic per real token and differ only in how much padding they carry, so the gap between them is a direct measurement of what the sort buys you.
正文由 探How 从来源站点 RSS 同步,内容未经改写;遇到排版缺失或需要图片、视频时请以原文为准。