In the previous chapter, we used the KV cache to speed up the generation of the GPT model. In this chapter, we explore another angle for improving generation efficiency: parallelism.
In real production environments, a GPT model often has to handle multiple generation requests at the same time. For example, a chatbot may have 100 users asking questions online at once. If they are processed one by one, users have to wait in line; if they can be processed in parallel as a batch, all requests can be computed at once, which greatly increases throughput.
Batch inference¶
The core idea of batch inference: pack multiple requests into one batch and complete the computation of all requests in a single forward pass.
How it works¶
Padding: requests have prompts of different lengths, so shorter sequences are filled with a special token (for example 0) up to the longest prompt length in the batch;
Masking: the padded positions have no real meaning, so a mask is needed to mark which positions are valid tokens and which are padding.
import time
import numpy as np
from src.core import *
from src.gpt import *np.random.seed(42)Model¶
GPT batch embedding layer¶
The GPT batch embedding layer (GPTEmbeddingWithPadding) needs to support custom position encoding: the positions of each sequence are encoded independently and are not affected by padding.
class GPTEmbeddingWithPadding(GPTEmbedding):
def forward(self, x: Tensor, positions):
token = self.embedding(x)
# Each sequence is offset independently
position = self.positional_embedding(Tensor(positions))
return self.dropout(token + position)GPT batch attention layer¶
The GPT batch attention layer (GPTAttentionWithPadding) needs to handle the padding mask:
after computing the attention scores, set the scores at padded positions to a very large negative number (-109);
after the Softmax, the weights at padded positions are 0 and do not affect valid tokens.
class GPTAttentionWithPadding(GPTAttention):
def forward(self, x: Tensor, valid):
norm = self.normalize(x)
head_embedding_size = self.embedding_size // self.heads
multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
shape = (-1, x.shape[1], self.embedding_size)
query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
scores = self.mask(query @ key * scale)
# Padding mask
scores = Tensor(np.where(valid[:, None, None, :], scores.data, self.mask.value))
weights = self.softmax(scores)
return x + self.dropout(self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape)))GPT batch Transformer block¶
The GPT batch Transformer block (GPTTransformerWithPadding) passes the mask parameter along, making sure padded positions do not take part in the computation.
class GPTTransformerWithPadding(Composite):
def __init__(self, embedding_size, heads):
self.attention = GPTAttentionWithPadding(embedding_size, heads)
self.feed_forward = GPTFeedForward(embedding_size)
super().__init__([self.attention,
self.feed_forward])
# Inference function, passing the padding mask
def forward(self, x: Tensor, valid):
return self.feed_forward(self.attention(x, valid))GPT batch model¶
The GPT batch model (GPTWithPadding) receives the mask and the position encodings and passes them on layer by layer.
class GPTWithPadding(Composite):
def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
self.embedding = GPTEmbeddingWithPadding(vocab_size, context_size, embedding_size)
self.transformers = [GPTTransformerWithPadding(embedding_size, heads) for _ in range(blocks)]
self.output = GPTOutput(embedding_size, vocab_size)
super().__init__([self.embedding] + self.transformers + [self.output])
# Inference function, passing the padding mask and the position numbers
def forward(self, x: Tensor, valid, positions):
x = self.embedding(x, positions)
for layer in self.transformers:
x = layer(x, valid)
return self.output(x)GPT batch trainer¶
The generate function of the GPT batch trainer GPTBatchTrainer:
Batch input: process multiple prompts at the same time;
Dynamic padding: before each generation step, pad the current sequences to the longest length in the batch;
Independent sampling: sample the next token independently for each sequence.
class GPTBatchTrainer(GPTTrainer):
def generate(self, dataset, prompts, steps=256):
# Inference mode
self.model.eval()
# Encode the prompts in a batch
batch_tokens = [dataset.encode(p) for p in prompts]
n = len(batch_tokens)
with Tensor.no_grad():
for _ in range(steps):
windows = [tokens[-dataset.context_size:] for tokens in batch_tokens]
width = max(len(w) for w in windows)
features = np.zeros((n, width), dtype=np.int64)
valid = np.zeros((n, width), dtype=bool)
positions = np.zeros((n, width), dtype=np.int64)
for i, window in enumerate(windows):
# Left padding: tokens are right-aligned, so the last position is always the newest token
pad = width - len(window)
features[i, pad:] = window
# Padding mask
valid[i, pad:] = True
# Position numbers
positions[i, pad:] = np.arange(len(window))
# Batch inference
logits = self.model(Tensor(features), valid, positions)
# Sample independently for each sequence
for i, tokens in enumerate(batch_tokens):
last_logits = logits.data[i, -1]
exp = np.exp(last_logits - np.max(last_logits))
probs = exp / np.sum(exp)
token = np.random.choice(len(probs), p=probs)
tokens.append(token)
return [dataset.decode(t) for t in batch_tokens]Hyperparameters¶
Model parameters¶
The model parameters must stay consistent with the pre-training stage.
BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"
dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPTWithPadding(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters)
trainer = GPTBatchTrainer(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)Batch prompts¶
prompts = ["ROMEO:", "JULIET:", "HAMLET:", "MACBETH:"]Serial inference¶
np.random.seed(42)
start = time.time()
for p in prompts:
print(trainer.generate(dataset, prompts=[p])[0])
print(f"{time.time() - start:.8f} seconds")ROMEO:
When I pray.
CAMILLO:
In my liest, I have done drunk.
CAMILLO:
Ay, since the speak, my doubt, fainners!
I am too, as to a kind in honourity,
Can a very disentent till take from here other
out, strike than 'Which children choistones
That would how have m
JULIET:
Sir, I will contracted up by the Gaunt,
And provost: ha? the desires of her babel,
Lest that love there of whose back high and stain: he he,
Where he I cannot die. The misfaction,
What, how, have unide a private way
met my noble--Surpose, let, her but, le
HAMLET:
We have hoped fathis professed of sorrow.
CAPULET:
Repress too victrate affair honour dead!
First Citizen:
I pray thou art field at the misicians win,
Be gentlewomen, and perform and concealunt stood
ass to as thine princes; for as from there to
Shool.
MACBETH:
Well, believe me; could never will good awake.
HASTINGS:
I'll cannot passlight the thanking of recovered
By Marcius nearer.
CLIFFORD:
What parter will us I more, if Hold answer;
Not we make of deserves dangerous are innocent
Whose times do your could tr
18.98765492 seconds
Batch inference¶
np.random.seed(42)
start = time.time()
results = trainer.generate(dataset, prompts=prompts)
for r in results:
print(r)
print(f"{time.time() - start:.8f} seconds")ROMEO:
Confess in to head; but I'll sea out
But all his otcher's eyes; it with their
flow; for I warrant then?
BRUTUS:
'IV Marcius, will I;'tis all-day:
Come, I part for me a sword,
But of it with hemies with my reting and prevail
If greet her laid, but well.
JULIET:
But I am lies bitter more and like a lost.
CLAUDIO:
O, my true, offence discretion.
JULIET:
Romeo, I'll canst upon. O, uncle, they do shall not
I cannot find upon; and seeing it them to me,
In to old unaviertain'd with clock'd.
ESCALUS:
I did thee, I g
HAMLET:
A blessing will steed tears, commands;
One would be so wonding of this acried:
So I bleem want full forget and that From
Offeriest! my cousin's hand be to my father,
And rebel, call me in this repeal'st;
For shake me Paduce consul,
Being being and with us
MACBETH:
Woo's committed for the life,
Where child belly to this knees man.
QUEEN MARGARET:
Let me King Henry, sir.
GREY:
With complimant your son is.
DUCHESS OF YORK:
And would not still the voices of supproof,
Minive I can a that death.
QUEEN ELIZABETH:
Your
19.44414520 seconds
In this chapter, we implemented batch inference: handle multiple generation requests at the same time, align the sequences with padding and masking, and complete all the computation in one forward pass.
But in the inference above we can see that batch inference did not show an efficiency gain (about 16.53 seconds serial versus about 16.50 seconds batched). There are two main reasons:
The total computation did not decrease: we did not use the KV cache, so every step recomputes the whole window. Batching only merges several requests into larger matrix multiplications, and the total computation is exactly the same as serial;
There is no idle compute to exploit: NumPy’s matrix multiplication is itself multi-threaded, and a single request can already saturate the CPU. On devices with a GPU, the benefit of batch inference would be much more significant.
With this, we have completed our discussion of model inference and tried to address the two core problems of controlling generation and improving efficiency:
Sampling strategies: control the quality and diversity of generation;
Stopping criteria: control when generation terminates;
KV cache: speeds up autoregressive generation;
Batch inference: raises the throughput of the model.
Exercises¶
Complete a comprehensive exercise: integrate sampling strategies, stopping criteria, the KV cache, and batch inference together.