批量推理 - 从零开始大模型

上一章，我们利用键值缓存加速了 GPT 模型的生成效率。本章，我们探讨一下从另一个角度提升模型生成效率：并行。

实际的生产环境中，GPT 模型经常需要同时处理多个生成请求。比如一个聊天机器人，同时有 100 个用户在线提问。如果逐个处理，用户需要排队等待；如果可以并行批量处理，则可以一次计算所有请求，大大提高吞吐量。

批量推理¶

批量推理（Batch Inference）的核心思想：把多个请求打包成一个批次，一次前向传播完成所有请求的计算。

实现原理¶

填充：不同请求的提示长度不同，因此要用特殊词元（比如：0）把短的序列补齐到批次内最大的提示长度；
掩码：填充的位置没有实际意义，因此需要用掩码标记哪些位置是有效词元，哪些是填充。

import time

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

模型¶

GPT 批量嵌入层¶

GPT 批量嵌入层（GPTEmbeddingWithPadding）需要支持自定义的位置编码：每个序列的位置独立编码，不受填充影响。

class GPTEmbeddingWithPadding(GPTEmbedding):

    def forward(self, x: Tensor, positions):
        token = self.embedding(x)
        # 每个序列独立偏移
        position = self.positional_embedding(Tensor(positions))
        return self.dropout(token + position)

GPT 批量注意力层¶

GPT 批量注意力层（GPTAttentionWithPadding）需要处理填充掩码：

计算注意力得分后，把填充位置的得分设置为一个极大的负数（-10⁹）；
Softmax 之后，填充位置的权重为 0，不影响有效词元。

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)

        # 填充掩码
        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 批量 Transformer 模块¶

GPT 批量 Transformer 模块（GPTTransformerWithPadding）传递掩码参数，确保填充位置不参与计算。

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

    # 推理函数，传递填充掩码
    def forward(self, x: Tensor, valid):
        return self.feed_forward(self.attention(x, valid))

GPT 批量模型¶

GPT 批量模型（GPTWithPadding）接收掩码和位置编码，逐层传递。

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

    # 推理函数，传递填充掩码和位置编号
    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 批量训练器¶

GPT 批量训练器 GPTBatchTrainer 的 generate 函数：

批量输入：同时处理多个提示；
动态填充：每步生成前，把当前序列填充到批次内最大长度；
独立采样：每个序列独立采样生成下一个词元。

class GPTBatchTrainer(GPTTrainer):

    def generate(self, dataset, prompts, steps=256):
        # 推理模式
        self.model.eval()
        # 批量编码提示词
        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):
                    # 左填充：词元靠右对齐，最后一个位置始终是最新的词元
                    pad = width - len(window)
                    features[i, pad:] = window
                    # 填充掩码
                    valid[i, pad:] = True
                    # 位置编号
                    positions[i, pad:] = np.arange(len(window))

                # 批量推理
                logits = self.model(Tensor(features), valid, positions)

                # 每个序列独立采样
                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]

超参数¶

模型参数¶

模型参数需要保持和预训练阶段一致。

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

建模¶

加载模型¶

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

批量提示词¶

prompts = ["ROMEO:", "JULIET:", "HAMLET:", "MACBETH:"]

串行推理¶

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
19.20661974 seconds

批量推理¶

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.30349398 seconds

本章，我们实现了批量推理：同时处理多个生成请求，通过填充和掩码对齐序列，一次前向传播完成所有计算。

但是在上面的推理中我们可以看到，批量推理并没有表现出效率提高（串行约 16.53 秒，批量约 16.50 秒）。这主要有两个原因：

总计算量没有减少：我们没有使用 KV 缓存，每一步都要对整个窗口重新计算。批量只是把多个请求合并成更大的矩阵乘法，总的计算量与串行完全相同；
没有闲置的算力可以利用：NumPy 的矩阵乘法本身就是多线程的，单个请求已经能够占满 CPU。在支持 GPU 的设备上，批量推理的收益会更加显著。

至此，我们完成了关于模型推理的讨论，尝试解决控制生成和提高效率这两个核心问题：

采样策略：控制生成的质量和多样性；
停止条件：控制生成的终止时机；
键值缓存：加速自回归生成；
批量推理：提高模型吞吐量。

课后练习¶

完成一个综合练习：把采样策略、停止条件、键值缓存和批量推理集成到一起。