自然语言处理（九）：大语言模型架构深度解析

1
2
3
4
5

from transformers import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM

enc  = AutoModel.from_pretrained("bert-base-uncased")            # encoder-only
dec  = AutoModelForCausalLM.from_pretrained("gpt2")               # decoder-only
ed   = AutoModelForSeq2SeqLM.from_pretrained("t5-small")          # encoder-decoder

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

import torch
import torch.nn as nn
import torch.nn.functional as F

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.g

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)   # gate
        self.w3 = nn.Linear(d_model, d_ff, bias=False)   # value
        self.w2 = nn.Linear(d_ff, d_model, bias=False)   # down-proj
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class LlamaBlock(nn.Module):
    """Pre-norm + RMSNorm + SwiGLU；attention 由外部注入。"""
    def __init__(self, d_model, attn, d_ff):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn  = attn
        self.norm2 = RMSNorm(d_model)
        self.ffn   = SwiGLU(d_model, d_ff)
    def forward(self, x, **kw):
        x = x + self.attn(self.norm1(x), **kw)
        x = x + self.ffn(self.norm2(x))
        return x

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class KVCache:
    """每层 KV cache，layout: [B, H_kv, T, D_h]。"""
    def __init__(self, B, H_kv, D_h, max_T, device, dtype=torch.float16):
        self.k = torch.empty(B, H_kv, max_T, D_h, device=device, dtype=dtype)
        self.v = torch.empty(B, H_kv, max_T, D_h, device=device, dtype=dtype)
        self.t = 0                                     # 已写入长度

    def append(self, k_new, v_new):
        T_new = k_new.size(2)
        self.k[:, :, self.t : self.t + T_new] = k_new
        self.v[:, :, self.t : self.t + T_new] = v_new
        self.t += T_new
        return self.k[:, :, : self.t], self.v[:, :, : self.t]

def attention_step(q_new, cache: KVCache, k_new, v_new, scale):
    """一次 decode 步：q_new 形状 [B, H_q, 1, D_h]。"""
    K, V = cache.append(k_new, v_new)                  # 完整前缀
    scores = (q_new @ K.transpose(-2, -1)) * scale     # [B, H, 1, T]
    attn   = scores.softmax(-1)
    return attn @ V                                    # [B, H, 1, D_h]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def precompute_rope(d_head, max_T, base=10000.0, device="cpu"):
    half = d_head // 2
    freqs = 1.0 / (base ** (torch.arange(0, half, device=device).float() / half))
    t = torch.arange(max_T, device=device).float()
    angles = torch.outer(t, freqs)                     # [T, d_head/2]
    return torch.cos(angles), torch.sin(angles)

def apply_rope(x, cos, sin):
    """x: [B, H, T, D_h]。旋转 (i, i+half) 维度对。"""
    half = x.size(-1) // 2
    x1, x2 = x[..., :half], x[..., half:]
    cos = cos[: x.size(-2)].to(x.dtype)
    sin = sin[: x.size(-2)].to(x.dtype)
    rot1 = x1 * cos - x2 * sin
    rot2 = x1 * sin + x2 * cos
    return torch.cat([rot1, rot2], dim=-1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_q_heads, n_kv_heads):
        super().__init__()
        assert n_q_heads % n_kv_heads == 0
        self.h_q  = n_q_heads
        self.h_kv = n_kv_heads
        self.rep  = n_q_heads // n_kv_heads
        self.d_h  = d_model // n_q_heads
        self.wq = nn.Linear(d_model, n_q_heads  * self.d_h, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.d_h, bias=False)
        self.wv = nn.Linear(d_model, n_kv_heads * self.d_h, bias=False)
        self.wo = nn.Linear(n_q_heads * self.d_h, d_model, bias=False)

    def forward(self, x, cos, sin, cache: KVCache | None = None):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.h_q,  self.d_h).transpose(1, 2)
        k = self.wk(x).view(B, T, self.h_kv, self.d_h).transpose(1, 2)
        v = self.wv(x).view(B, T, self.h_kv, self.d_h).transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        if cache is not None:
            k, v = cache.append(k, v)
        # 把每个 KV 头广播到对应的 query 组
        k = k.repeat_interleave(self.rep, dim=1)
        v = v.repeat_interleave(self.rep, dim=1)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=cache is None)
        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.wo(out)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# 现代 transformers/torch 里你不需要自己写 kernel，开关打开即可：
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    attn_implementation="flash_attention_2",   # 也可选 "sdpa" / "eager"
    torch_dtype=torch.float16,
    device_map="auto",
)
# 或者用 PyTorch 自带的 scaled_dot_product_attention，
# 它会自动派发到 flash / mem-efficient / math 三种内核。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

class TopKRouter(nn.Module):
    def __init__(self, d_model, n_experts, k=2):
        super().__init__()
        self.gate = nn.Linear(d_model, n_experts, bias=False)
        self.k = k

    def forward(self, x):
        logits = self.gate(x)                                   # [*, N]
        topv, topi = logits.topk(self.k, dim=-1)
        weights = topv.softmax(-1)                              # 重新归一
        return weights, topi                                    # [*, k], [*, k]

class SparseMoE(nn.Module):
    def __init__(self, d_model, d_ff, n_experts, k=2):
        super().__init__()
        self.experts = nn.ModuleList(
            [SwiGLU(d_model, d_ff) for _ in range(n_experts)]
        )
        self.router = TopKRouter(d_model, n_experts, k)

    def forward(self, x):
        B, T, D = x.shape
        flat = x.view(-1, D)
        w, idx = self.router(flat)                              # [BT, k]
        out = torch.zeros_like(flat)
        for slot in range(self.router.k):
            for e in range(len(self.experts)):
                mask = idx[:, slot] == e
                if mask.any():
                    out[mask] += w[mask, slot, None] * self.experts[e](flat[mask])
        return out.view(B, T, D)

1
2
3
4
5

def quantize_int8_per_channel(W):
    """W: [out, in]。返回 int8 权重 + per-row fp16 scale。"""
    scale = W.abs().amax(dim=1, keepdim=True) / 127.0          # [out, 1]
    Wq = (W / scale).round().clamp(-127, 127).to(torch.int8)
    return Wq, scale.to(torch.float16)

1
2
3
4
5
6

# 加载预量化的 AWQ checkpoint
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-2-7B-AWQ", device_map="auto", torch_dtype="auto",
)
tok = AutoTokenizer.from_pretrained("TheBloke/Llama-2-7B-AWQ")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

from vllm import LLM, SamplingParams

llm = LLM(
    model="TheBloke/Llama-2-7B-AWQ",
    quantization="awq",
    dtype="float16",
    gpu_memory_utilization=0.90,
    max_model_len=8192,
)
params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=128)
for o in llm.generate(["用两句话解释 MoE。"], params):
    print(o.outputs[0].text)