Batching by Length Instead of Looping Item by Item for SLM Optimization

來源原標題

Batching by Length Instead of Looping Item by Item for SLM Optimization

作者

Matthew Mayo

原始連結

https://www.kdnuggets.com/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization

摘要原文

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. By Matthew Mayo , KDnuggets Managing Editor on September 25, 2026 in Language Models 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 . Why Batch by Length Instead of Looping Item by Item 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. Looping Item by Item 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) < n_words: words += FILLER return " ".join(words[:n_words]) tickets_var = [make_ticket(BASE_TICKETS[i % 3], int(n)) for i, n in enumerate(target_words)] prompts = [build_prompt(t) for t in tickets_var] token_lengths = [len(tokenizer(p, add_special_tokens=False)["input_ids"]) for p in prompts] print( f"Prompt lengths: min {min(token_lengths)}, " f"median {int(np.median(token_lengths))}, " f"max {max(token_lengths)} tokens" ) print(f"Padding every item to the global maximum would process " f"{max(token_lengths) * len(prompts) / sum(token_lengths):.1f}x the necessary tokens") # time inference baseline_predictions = [] start = time.time() for n, prompt in enumerate(prompts, start=1): # this loop runs for minutes on CPU, so report progress rather than sitting silent if n % 100 == 0: print(f" {n}/{len(prompts)} tickets ({(time.time() - start) / n:.2f}s each)", flush=True) inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(model.device) with torch.no_grad(): logits = model(**inputs, **LAST_LOGIT_ONLY).logits[0, -1, :] baseline_predictions.append(LABELS[int(logits[label_first_ids].argmax())]) duration_loop = time.time() - start # output task metrics print(f"One at a time: {duration_loop:.2f} seconds ({len(prompts) / duration_loop:.1f} items/sec)") 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.