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

Sorting support tickets by token length before batching, rather than looping item by item, is the third and final optimization in a KDnuggets series on small language model (SLM) narrow automation, with benchmarks run on Qwen2.5-0.5B-Instruct in float16 via Hugging Face Transformers on an M2 MacBook Air with 24GB RAM. The article argues that batch size 1 is memory-bandwidth bound, since the hardware streams every weight out of memory to serve a single sequence, and that padding each batch to its own local maximum instead of the dataset's global maximum avoids wasted compute on padding tokens.

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. Previous articles in this series discussed constraining output space https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization as well as reusing the prompt prefix with a key-value cache https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization , 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 https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct in float16 through Hugging Face Transformers https://huggingface.co/docs/transformers/index , 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 https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization . 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: python 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. python 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 < 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" def run batched order, batch size : predictions = None len prompts processed tokens = real tokens = 0 start = time.time for i in range 0, len order , batch size : idx = order i:i + batch size batch = tokenizer prompts j for j in idx , add special tokens=False, padding=True, return tensors="pt", .to model.device processed tokens += batch "input ids" .numel real tokens += int batch "attention mask" .sum with torch.no grad : logits = model batch, LAST LOGIT ONLY .logits :, -1, : best = logits :, label first ids .argmax dim=-1 for slot, choice in zip idx, best.tolist , strict=True : predictions slot = LABELS choice return predictions, time.time - start, processed tokens, real tokens def classify one prompt : """Reference path: a single unpadded sequence. Used only to verify.""" 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, : return LABELS int logits label first ids .argmax order = sorted range len prompts , key=lambda i: token lengths i predictions, duration batched, processed, real = run batched order, BATCH SIZE print f"Length-bucketed batching: {duration batched:.2f} seconds {len prompts / duration batched:.1f} items/sec " print f"Padding overhead: {100 1 - real / processed :.1f}% of processed tokens were padding" correctness: padded rows must score the same as unpadded ones. check a spread of lengths rather than all 600, since the point is to catch a padding-side or position bug, and such a bug shows up on the very first padded row. probe = order ::60 mismatches = i for i in probe if classify one prompts i = predictions i print f"Batched vs unbatched agreement on {len probe } probes: {len probe - len mismatches }/{len probe }" assert not mismatches, f"Batched path disagrees at indices {mismatches}" estimate the per-item cost on the same probe set, then scale start = time.time for i in probe: classify one prompts i Output: Loading weights: 100%|████████████████████████████████████████████████████████████████████| 290/290 00:01<00:00, 252.11it/s Prompt lengths: min 48, median 94, max 449 tokens Length-bucketed batching: 79.60 seconds 7.5 items/sec Padding overhead: 7.6% of processed tokens were padding Batched vs unbatched agreement on 10 probes: 10/10 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. - Setting padding side = "left" is required here, not a choice. With right padding, logits :, -1, : would land on a pad token for every row shorter than the batch maximum, deceptively producing garbage predictions. Left padding guarantees index -1 is the true final token of every sequence. - Left padding does shift each row's absolute token positions, because a plain model batch call numbers positions from zero across the padded width rather than deriving them from the attention mask. For a rotary-embedding model like Qwen2.5 this is harmless, since attention depends only on the relative distance between tokens and every real token in a row shifts by the same amount. For a model with learned absolute position embeddings it would not be harmless, and you would need to pass position ids built from the mask. Either way, the agreement check against the per-item loop is what tells you which situation you are in. - Ask for the last position's logits only. By default a causal LM returns a logit vector for every input position, and the vocabulary here is around 150k entries: at batch 32 by 400 tokens in float32 that is a multi-gigabyte tensor allocated and discarded on every single batch. logits to keep=1 named num logits to keep in older versions of Transformers suppresses it. This is negligible at batch size 1 with a short prompt, which is why it never came up in the earlier articles, and it dominates everything else once you start batching long inputs. - Sorting before chunking holds padding overhead to a few percent. The same data through fixed-size batches in arbitrary order pushes it far higher, which is the difference the script measures directly rather than asserting: every padding token is a token the hardware processed for no reason. - Sorting reorders the data, so keep the original indices around and write results back to their proper slots, as the order list does above. Losing the alignment between inputs and predictions is an easy and very expensive bug, and unlike a crash it produces plausible-looking output. - Pick BATCH SIZE by measuring, not by intuition, which is what the sweep at the end of the script is for. Throughput climbs steeply and then plateaus once you saturate compute; past that point you are only increasing the odds of an out-of-memory error on your longest bucket. The best value depends on your hardware and on your length distribution, so it is worth re-running the sweep when either changes. One caveat: prefix caching and batching need care to combine. The cache we built in the previous article has a batch dimension of 1, so reusing it across a batch means expanding every key and value tensor along that dimension to match, and cropping it back correctly afterwards. It is worth doing when your prefix is long, but do it deliberately and verify the predictions against the unbatched path rather than assuming the two optimizations compose for free. Wrapping Up This has been our third and final attempt at optimizing SLMs for our narrow automation series, and our target technique this time was length-bucketed batching. This technique replaces a one-item-at-a-time loop with sorted batches that pad to their own local maximum, which gets the hardware out of the memory-bandwidth-bound regime without paying for the padding that naive batching would introduce. By implementing it, we process the same 600 tickets in a fraction of the wall-clock time, with the same predictions coming out the other end. None of these techniques makes the model "smarter." Each one was verified by checking that its output was identical to the slower path it replaced, and that check is the whole reason to trust the speedup numbers at all. An optimization that changes your predictions is not an optimization, it is a regression with a stopwatch attached. To recap the series: - Constrained scoring https://www.kdnuggets.com/constraining-output-space-small-language-model-narrow-automation-optimization replaces free-form generation and string parsing with a single forward pass restricted to the valid label set, making malformed output structurally impossible while handing you a confidence score for routing edge cases to humans - Prefix key-value caching https://www.kdnuggets.com/reusing-the-prompt-prefix-with-a-key-value-cache-for-slm-optimization computes the static instruction block once instead of once per item, and pays off in proportion to how much of your prompt never changes - Length-bucketed batching https://www.kdnuggets.com/batching-by-length-instead-of-looping-item-by-item-for-slm-optimization gets the hardware out of the memory-bandwidth-bound constraint and keeps padding overhead near zero by grouping similarly sized inputs together \ Matthew Mayo\ https://www.kdnuggets.com/wp-content/uploads/./profile-pic.jpg https://www.linkedin.com/in/mattmayo13/ @mattmayo13 https://twitter.com/mattmayo13 holds a master's degree in computer science and a graduate diploma in data mining. As managing editor of KDnuggets https://www.kdnuggets.com/ & Statology https://www.statology.org/ , and contributing editor at Machine Learning Mastery https://machinelearningmastery.com/ , Matthew aims to make complex data science concepts accessible. His professional interests include natural language processing, language models, machine learning algorithms, and exploring emerging AI. He is driven by a mission to democratize knowledge in the data science community. Matthew has been coding since he was 6 years old.