[stanford-cs336] Lecture-1

Reference b站 https://www.bilibili.com/video/BV1pAjUzYEaK/?spm id from=333.337.search card.all.click&vd source=938c1d285952e2b4c3852a83df2851bb ; Youtube https://www.youtube.com/wat

Reference

b站;Youtube;课件;discord学习小组

FLOPs

定义

每秒浮点运算次数, Floating Point Operations per Second

浮点数的加法、乘法、除法。FLOPs 表示一个模型在一次前向计算中需要执行多少次浮点运算

MFLOPs是million FLOPs, GFLOPs是giga FLOPs

矩阵FLOPs

添加图片注释,不超过 140 字(可选)

添加图片注释,不超过 140 字(可选)

线性层FLOPs

添加图片注释,不超过 140 字(可选)

粗略估算

添加图片注释,不超过 140 字(可选)

代码

估算FLOPs的一段代码

fastai course2: 14_augment.ipynb

def _flops(x, h, w): # x是tensor,height, width
    if x.dim()<3: return x.numel() # numel返回的是元素数量,比如shape(3,4)返回的是12
    if x.dim()==4: return x.numel()*h*w

@fc.patch
def summary(self:Learner):
    res = '|Module|Input|Output|Num params|MFLOPS|\n|--|--|--|--|--|\n'
    totp,totf = 0,0
    def _f(hook, mod, inp, outp):
        nonlocal res,totp,totf
        nparms = sum(o.numel() for o in mod.parameters())
        totp += nparms
        *_,h,w = outp.shape
        flops = sum(_flops(o, h, w) for o in mod.parameters())/1e6
        totf += flops
        res += f'|{type(mod).__name__}|{tuple(inp[0].shape)}|{tuple(outp.shape)}|{nparms}|{flops:.1f}|\n'
    with Hooks(self.model, _f) as hooks: self.fit(1, lr=1, cbs=SingleBatchCB())
    print(f"Tot params: {totp}; MFLOPS: {totf:.1f}")
    if fc.IN_NOTEBOOK:
        from IPython.display import Markdown
        return Markdown(res)
    else: print(res)

添加图片注释,不超过 140 字(可选)

FLOPs占比

添加图片注释,不超过 140 字(可选)

在小模型里,multi head attention (MHA)的FLOPs 和 MLP(FFN) 不相上下。

但是在大模型里,大部分都是MLP,这个时候如果一味的optimize MHA是不work的。

Scaling规模扩大

可以指参数,数据集,或者计算量FLOPs

添加图片注释,不超过 140 字(可选)

较难任务只有在FLOPs突破一定规模的时候才能开始有效

添加图片注释,不超过 140 字(可选)

Modular arithmetic,模运算,17 mod5 =2 , 23mod7=2,最常见的就是时钟mod12, 通常对大模型比较难

IPA transliterate,把英文单词转换成音标,需要打模型对发音的理解。

word unscramble, 把英文单词里的字母打乱,让模型重新组合成单词

Persian QA, 波斯语言的问答,波斯语的语库很小,如果能正确问答,说明模型具有泛化能力。

Truthful QA,真实问答,问生鱼片能治病吗?错误是生鱼片是个好东西,正确是不能。

Grounded mapping, 概念映射,问一种生活在非洲的食草动物,脖子很长,答长颈鹿。

Multitask NLU, 多任务语言理解,打开了窗户意味着房间通风了。

word in context,词在上下文中的意思。he banked the plane, 这里bank不是银行而是指转弯。

一些概念

语言模型领域里有一些直觉的东西,很难解释,比如SwiGLU,都是通过实验验证得来的。

直觉的东西很难transfer。

技术性的东西可以transfer,比如transformer如何工作的,比如model parallel里GPU如何被分派,还有mindset上,scaling law很重要之类的。

efficiency

另外,模型的efficiency很重要,2012到2019年间,imagenet efficiency提升了44倍。

而且efficiency对超大模型很重要,因为训练一次的成本很高。

模型的公开程度

closed , 比如4o

open weights,比如deepseek, 不公开模型code,但是给了weights,huggingface上的一些模型,然后有paper作为支持。

open source,比如OLMo, 公开模型code,也给weights,方便重复。

Others

Common Crawl dump 是一个包含了数十亿网页原始内容的公开网络抓取数据集,按月发布,用于训练大型语言模型,包括 GPT、BERT、PaLM

五个模块

可以从五个方面最大化模型efficiency

添加图片注释,不超过 140 字(可选)

tokenizer,是可以把文字转换成integer整数的东西。比如apple转换成1,banana转换成2.

Byte pair encoding

BPE是一种分词方法,它将常见的字符对合并成新的“子词单位”,以减少词汇碎片。

下面是一个例子

添加图片注释,不超过 140 字(可选)

课程上的例子,这个merge的过程可以循环几次。

def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:  # @inspect string, @inspect num_merges
    Start with the list of bytes of string.
    indices = list(map(int, string.encode("utf-8")))  # @inspect indices
    merges: dict[tuple[int, int], int] = {}  # index1, index2 => merged index
    vocab: dict[int, bytes] = {x: bytes([x]) for x in range(256)}  # index -> bytes
    for i in range(num_merges):
        Count the number of occurrences of each pair of tokens
        counts = defaultdict(int)
        for index1, index2 in zip(indices, indices[1:]):  # For each adjacent pair
            counts[(index1, index2)] += 1  # @inspect counts
        Find the most common pair.
        pair = max(counts, key=counts.get)  # @inspect pair
        index1, index2 = pair
        Merge that pair.
        new_index = 256 + i  # @inspect new_index
        merges[pair] = new_index  # @inspect merges
        vocab[new_index] = vocab[index1] + vocab[index2]  # @inspect vocab
        indices = merge(indices, pair, new_index)  # @inspect indices
    return BPETokenizerParams(vocab=vocab, merges=merges)

Transformer architecture

添加图片注释,不超过 140 字(可选)

左边的是一个整体的GPT架构图,右边是灰色transformer block的具体架构图,也就是transformer具体的样子。

添加图片注释,不超过 140 字(可选)

Single-Head

这是一段self-attention的代码(single head),input是image, 来自于fastai 2: 27_attention.ipynb

class SelfAttention(nn.Module):
    def __init__(self, ni):
        super().__init__()
        self.scale = math.sqrt(ni)
        self.norm = nn.BatchNorm2d(ni)
        self.qkv = nn.Linear(ni, ni*3)
        self.proj = nn.Linear(ni, ni)
    
    def forward(self, inp):
        n,c,h,w = inp.shape # n:batch size, c:channel, h:height, w: width
        x = self.norm(inp).view(n, c, -1).transpose(1, 2) # view merge h,w to hw, transpose makes hw to the center
        q,k,v = torch.chunk(self.qkv(x), 3, dim=-1)
        s = (q@k.transpose(1,2))/self.scale
        x = s.softmax(dim=-1)@v
        x = self.proj(x).transpose(1,2).reshape(n,c,h,w)
        return x+inp

这里图片是input,shape是n,c,h,w -- n:batch size, c:channel, h:height, w: width

view(n,c-1)把图片shape变成n,c, hw, transpose(1,2)变成了n,hw(seq_len), c,因为transformer永远是在中间这个seq_len上work。

SelfAttention(ni=128)

ni等于c的大小

Multi-Head

多头注意力机制(Multi-Head Attention)是 Transformer 的核心创新之一。它的核心思想是: 并行学习多个注意力模式,比如一个头关注局部信息、一个头关注全局句法、一个头关注实体。

这是一个multi-head version

class SelfAttentionMultiHead(nn.Module):
    def __init__(self, ni, nheads):
        super().__init__()
        self.nheads = nheads
        self.scale = math.sqrt(ni/nheads)
        self.norm = nn.BatchNorm2d(ni)
        self.qkv = nn.Linear(ni, ni*3)
        self.proj = nn.Linear(ni, ni)
    
    def forward(self, inp):
        n,c,h,w = inp.shape
        x = self.norm(inp).view(n, c, -1).transpose(1, 2)
        x = self.qkv(x)
        x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
        q,k,v = torch.chunk(x, 3, dim=-1)
        s = (q@k.transpose(1,2))/self.scale
        x = s.softmax(dim=-1)@v
        x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)
        x = self.proj(x).transpose(1,2).reshape(n,c,h,w)
        return x+inp
     

这里实现multi-head的关键在于这行

x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)

把channel给拆成了(h d),每个维度是 d = c // nheads, (channel 在CNN里已经扩了很多,这里只是图片的例子。)

拆后 shape 是 (n * h, s, d)

和单头做一次 softmax 乘积 不同,多头的话可以做nheads个头并行 softmax,最后 concat 回去。表达能力高,可捕捉多个关注模式。

Quadratic blowup(平方级膨胀)

Quadratic blowup 是指 self-attention 的计算/内存复杂度为 O(n^2),当序列长度 n增大时,成本平方级增长。

每一步的计算复杂度:

添加图片注释,不超过 140 字(可选)

关键瓶颈是 Q 与 K 转置相乘, 生成一个 n x n 矩阵

因此,增长量会随着序列大小Quadratic(平方级别n^2)增长, (而不是expotential指数增长哦,2^n,远远快于平方)

添加图片注释,不超过 140 字(可选)

为了防止平方级增长,会有sliding window attention, 还有linear attention

Lower dimension

可以把前面linear的后面的ni*3改成更小的dimension,这里比如加一个小的kqv_dim,代替掉之前的ni x3

class SelfAttentionLowDim(nn.Module):
    def __init__(self, c, qkv_dim=32):
        super().__init__()
        self.scale = math.sqrt(qkv_dim)
        self.norm = nn.BatchNorm2d(c)
        self.qkv = nn.Linear(c, qkv_dim * 3)
        self.proj = nn.Linear(qkv_dim, c)  # 把低维输出映射回通道数

好处是可以降低计算复杂度与内存需求。

Lower dimension- group query attention

所有 heads 分为若干组,每组共享 Key/Value 权重。

每个 Query head 拥有独立的 weights Q, 但多个 Query head 共享同一组weights K & V

假设总共有 **16 个 attention heads,**我们设置 **每 4 个 Query head 共享一组 Key/Value。**这就是 GQA ratio = 4:1 的一种配置。

下面是一个chatgpt生成的coding例子

class GroupQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, kv_groups):
        super().__init__()
        self.n_heads = n_heads
        self.kv_groups = kv_groups
        self.d_head = d_model // n_heads

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model // kv_groups)
        self.v_proj = nn.Linear(d_model, d_model // kv_groups)

        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.d_head)
        k = self.k_proj(x).view(B, T, self.kv_groups, self.d_head * (self.n_heads // self.kv_groups))
        v = self.v_proj(x).view(B, T, self.kv_groups, self.d_head * (self.n_heads // self.kv_groups))
        # Broadcast k,v across heads within each group...

Low dimension - multi head latent attention

用一个固定数量的latent vectors(例如 256 个)来替代 n 个 queries

添加图片注释,不超过 140 字(可选)

Query 是 m个 learnable latent vectors

输出是 m个注意力结果,可以进一步处理或 decode 回原始序列

一段chatgpt生成的coding例子:

import torch
import torch.nn as nn

class MultiHeadLatentAttention(nn.Module):
    def __init__(self, input_dim, latent_dim, num_latents, num_heads):
        super().__init__()
        self.latents = nn.Parameter(torch.randn(num_latents, latent_dim))  # (m, d)
        self.cross_attn = nn.MultiheadAttention(embed_dim=latent_dim, num_heads=num_heads, batch_first=True)

        # 如果输入维度与 latent_dim 不一致,先投影
        self.input_proj = nn.Linear(input_dim, latent_dim)

    def forward(self, x):
        """
        x: (batch_size, seq_len, input_dim)
        """
        b, n, _ = x.shape
        x = self.input_proj(x)  # shape: (b, n, latent_dim)

        # Expand latent queries to batch: (b, m, d)
        latents = self.latents.unsqueeze(0).expand(b, -1, -1)

        # Cross-attention: Query = latent, Key/Value = input
        out, _ = self.cross_attn(query=latents, key=x, value=x)
        return out  # shape: (b, num_latents, latent_dim)

# ✨ 测试一下
x = torch.randn(4, 100, 128)  # batch=4, seq_len=100, input_dim=128
model = MultiHeadLatentAttention(input_dim=128, latent_dim=64, num_latents=16, num_heads=4)

output = model(x)
print(output.shape)  # ➜ (4, 16, 64)

GPU

结构

了解GPU的结构可以帮助写kernel代码

A100的样子,左边是单个 SM(Streaming Multiprocessors),右边是一个128 个 SM 的芯片结构。

添加图片注释,不超过 140 字(可选)

一些关于每个模块的解释

添加图片注释,不超过 140 字(可选)

Warp

Warp 是 GPU 的最小执行单元(32 个线程),由 SM 内的 Warp Scheduler 在每个周期挑选 ready warp 来调度执行,充分隐藏 memory latency 并提升并行效率。

每个时钟周期:

Warp Scheduler 会选一个 ready warp 来执行下一条指令

所谓 ready warp:所有线程不等待 memory、不阻塞、不发散

如果 warp A 在等 memory(如 global memory load),scheduler 会暂时不调它

会切换到另一个 ready warp B

这让 GPU 能在没有 cache 命中时仍保持高吞吐

如果遇到类似这样的case

if x[i] > 0:
    y[i] = x[i]
else:
    y[i] = -x[i]

GPU 会顺序执行两个分支路径 → 降低并行度(性能下降)

所以 Triton / CUDA 优化中经常使用

mask = cond
val = tl.where(mask, a, b)

Triton

下面是triton的常用命令和对应的GPU结构

添加图片注释,不超过 140 字(可选)

KV Cache

KV cache 是缓存 transformer 中已经计算过的 key 和 value,用来避免重复计算,极大加速自回归生成。

在生成任务(如 GPT)中,模型是 一个 token 一个 token 地生成,例如:

Input: "The cat"
→ 模型生成 "sat"

Input: "The cat sat"
→ 模型生成 "on"

Input: "The cat sat on"
→ 模型生成 "the

每次都重新计算整个序列的 attention(包括前面的所有 token),重复浪费大量算力。

添加图片注释,不超过 140 字(可选)

添加图片注释,不超过 140 字(可选)

每步只计算 1 个 token 的 Q/K/V

K/V 是增量累积、永久缓存,只追加不回算

GPT example

# 每层 transformer 都会存 KV cache
kv_cache = {
    "layer_0": {
        "k": tensor of shape (batch, num_heads, seq_len, head_dim),
        "v": tensor of shape (batch, num_heads, seq_len, head_dim)
    },
    ...
}

在每层中,当前新 token 的 Q 会去 attend 所有历史 K(来自 cache), 得到上下文感知的表示

时间复杂度

  • 没有 KV cache:每一步生成都要重新处理整段序列 → O(n²) 时间复杂度
  • 有 KV cache:每一步只需处理新 token → O(n) 时间复杂度(更精准地说是 O(1) per step)

Inference

分为两个阶段,左边是prefill phase 预填阶段,右边是decoding phase解码阶段

添加图片注释,不超过 140 字(可选)

预填阶段会先处理prompt,并行处理所有的token,计算每个token的KQV,然后保存到KV Cache

解码阶段会一个一个的生成,生成新的token,用上一个token作为新的输入。

添加图片注释,不超过 140 字(可选)

解码过程中有一些trick:

  • 比如用更小的模型(pruning、quantization、distillation)加速生成。
  • Speculative decoding,先用一个草稿模型(draft model)一次性预测多个 token, 再用大模型验证这几个 token 是否合理(可以并行打分)