[stanford-cs336] Assignment-1 Chapter-3

Reference

https://github.com/stanford-cs336/assignment1-basics/tree/maingithub.com/stanford-cs336/assignment1-basics/tree/main

Setup

下载并安装环境

git clone https://github.com/stanford-cs336/assignment1-basics.git
cd assignment1-basics
uv self update # optional
uv sync # 如果没安装uv，就去看我的uv文章
source .venv/bin/activate

先测试一下，应该是全部fail

uv run pytest

要求

因为一切都是手攒，所以

• 不可以用:
• torch.nn &
• torch.nn.functional &
• torch.optim 里面的组件
• 可以用:
• torch.nn.Parameter，
• Container classes in torch.nn (e.g.,Module, ModuleList, Sequential)
• torch.optim.Optimizer base class

所有的代码写到cs336_basics/* 下面，然后在adapters.py里调用自己的.py，需要通过所有的test。

最后的作业分为两部分，一个是code，一个是写的报告。

第三章 模型结构

3.4.1-2 Linear

要求

Xavier initialization

• 用torch.nn.init.trunc_normal_去initialize weights
• 用Module.load_state_dict来加载weights

答案

在cs336_basics里建一个layer.py

import torch
from torch import nn
from torch.nn import init
from einops import rearrange, einsum

class Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int, device=None, dtype=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        factory_kwargs = {'device': device, 'dtype': dtype}

        # Weight shape is (out_features, in_features)
        self.weight = nn.Parameter(torch.empty((out_features,in_features),**factory_kwargs))
        # Initialize weights using truncated normal
        std = (2 / (in_features + out_features)) ** 0.5
        init.trunc_normal_(self.weight, mean=0.0, std=std, a=-3*std, b=3*std)

    def forward(self, x):
        return einsum(x, self.weight, "... d_in, d_out d_in -> ... d_out") # equivalent to x@self.weight.T

课程里建议用einsum来表示@ matrix multiplication，可以更清晰的看到每个维度是如何变换的。

之后在adaptor.py里

import， 这里用*，因为后面我们还会在layer.py里写别的。

from cs336_basics.layer import *

然后代码部分

def run_linear(
    d_in: int,
    d_out: int,
    weights: Float[Tensor, " d_out d_in"],
    in_features: Float[Tensor, " ... d_in"],
) -> Float[Tensor, " ... d_out"]:
    device, dtype = in_features.device, in_features.dtype
    model = Linear(d_in, d_out, device=device, dtype=dtype)
    model.load_state_dict({'weight': weights})
    return model(in_features)

测试

uv run pytest -k test_linear

通过且没有warning

Xavier initialization

作业里给的initalize的例子是Xavier initialization。X@W.T, X shape (d_in), W shape(d_out, d_in), sum axis=d_in，如果X和W都是standard normal distribution N(0,1), forward pass 中matrix multiply 会导致一个N(0，d_in)的matrix，Var是d_in, σ是d_in sqrt。如何设置W的distribution，可以让最后的Var是1呢？

添加图片注释，不超过 140 字（可选）

所以forward中，W的distribution设置成N(0,1/d_in)，可以让输出的distribution为N(0,1)

反向也要考虑：

添加图片注释，不超过 140 字（可选）

所以Xavier 的normal initialization是N(mean=0, Var = 2/(d_in+d_out) )

关于normal distribution

N指normal distribution, 括弧里第一个位置是mean μ，第二个是variance 方差，标准差σ是方差的sqrt。

• N(mean μ, variance Var)
• Var = σ^2 ; σ = sqrt Var

Normal distribution里

• 68% 的值在 ±1σ里
• 95% 的值在 ±2σ里
• 99.7% 的值在 ±3σ里

±3σ 涵盖了99.7%的值。Standard normal distribution是指mean 为0，Var为1的分布，也就是N(0,1)

Kaiming initialization

何凯明的initialize版本考虑到了ReLU的存在，ReLU会把一半的负值归零。

那么最终X@W.T的数值就是之前的一半，d_in/2 , 想要compensate，需要把W的Var设置成2/d_in即可。

下面是两种init的比较：

添加图片注释，不超过 140 字（可选）

jaxtyping的格式

from jaxtyping import Float, Int
Float[Tensor, "batch dim"]
Int[Tensor, "length"]
Float[npt.NDArray, "n_features"]

前面是一个数据type(Int或者Float）然后跟一个container type(比如Tensor，或者npt.NDArray)，然后quotes""里面是dimension的定义。

einsum

rules：

• ->，右侧出现的dim必须出现在左侧
• ->, 左侧出现的dim，如果没出现在右侧，会被自动求和
• -> , 如果右侧出现的dim只出现在左侧的某一个输入里，那么会把另一个输入自动broadcast。

注意torch.einsum和einops.einops 的格式有些不同，einops里的是把两个item放在前面，equation放后面。torch里是把equation放在前面。

作业pdf的einsum的例子是einops的einsum的。除了上面代码里写的matrix multiplication，还有一个很有意思的例子：

images = torch.randn(64, 128, 128, 3) # (batch, height, width, channel)
dim_by = torch.linspace(start=0.0, end=1.0, steps=10)

## in one go:
dimmed_images = einsum(
images, dim_by,
"batch height width channel, dim_value -> batch dim_value height width channel"
)

一行einsum的代码，就可以broadcast展开dim_by(10,) 和images来 match右边的shape

images展开为(batch, 1, height, width, channel)

dim_by展开为(1, dim_value, 1, 1, 1)

然后右边没有任何dim缺少，所以只是单纯的做element wise multiply，也就是逐元素相乘。

最后变成(batch, dim_value, height, width, channel)

einops官网教程

einops官网pytorch教程

einx

作业pdf还建议已经熟练掌握einsum的童鞋去熟悉einx，更进阶。

安装pip install einx。格式是equation放在前面，item放后面，可以加graph=True可以看到具体如何实现。

einx.dot教程

比如，einx.dot可以在一行中同时做rearrange+ einsum

# Simple grouped linear layer
x = np.ones((20, 16))
w = np.ones((8, 4))
print(einx.dot("b (g c1), c1 c2 -> b (g c2)", x, w, g=2, graph=True))
import numpy as np
def op0(i0, i1):
    x0 = np.reshape(i0, (20, 2, 8))
    x1 = np.einsum("abc,cd->abd", x0, i1)
    x2 = np.reshape(x1, (20, 8))
    return x2

作业pdf的例子：

要对images (batch, height, width, channel),做一个 形状为(height × width, height × width)的linear transformation。

channels_last = torch.randn(64, 32, 32, 3) # (batch, height, width, channel)
B = torch.randn(32*32, 32*32)

如果不用einx，一个策略就是用rearrange先把channels_last里的32和32 merge到一起，然后把32*32和最后一位的3互换位置，然后matrix multiply B (out, in) 【可以@B.T，或者einsum做】。之后再把后面两位的位置换回来，然后用view恢复到原来的形状。

如果用einx.dot，一行代码就可以搞定：

height = width = 32
channels_last_transformed = einx.dot(
"batch row_in col_in channel, (row_out col_out) (row_in col_in)"
"-> batch row_out col_out channel",
channels_last, B,
col_in=width, col_out=width
)

3.4.3 Embedding

要求

添加图片注释，不超过 140 字（可选）

答案

class Embedding(nn.Module):
    def __init__(self, 
                 num_embeddings, # vocabulary size
                 embedding_dim, # embedding dimension
                 device=None, dtype=None):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        factory_kwargs = {'device': device, 'dtype': dtype}

        # Weight shape is (num_embeddings, embedding_dim)
        self.weight = nn.Parameter(torch.empty((num_embeddings, embedding_dim), **factory_kwargs))

        # Initialize weights using truncated normal
        std = 1
        init.trunc_normal_(self.weight, mean=0.0, std=std, a=-3*std, b=3*std)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        return self.weight[token_ids]  # Select rows corresponding to token_ids

在adapters.py里

def run_embedding(
    vocab_size: int,
    d_model: int,
    weights: Float[Tensor, " vocab_size d_model"],
    token_ids: Int[Tensor, " ..."],
) -> Float[Tensor, " ... d_model"]:
    device = token_ids.device
    model = Embedding(vocab_size, d_model, device=device) # if dtype is None, it will be float() in torch default
    model.load_state_dict({'weight': weights})
    return model(token_ids)

测试

uv run pytest -k test_embedding

PASS

3.5.1 RMSNorm

RMSNorm公式里面和gi的关系是* element wise multiplication, 不是@

要求

• RMSNorm 里的weights (gi) 用1来初始化
• eps设置成1e-5
• 为了防止x square overflow，用下面这段代码：

in_dtype = x.dtype
x = x.to(torch.float32)
# Your code here performing RMSNorm
...
result = ...
# Return the result in the original dtype
return result.to(in_dtype)

答案

class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None):
        super().__init__()
        self.d_model = d_model
        self.eps = eps
        factory_kwargs = {'device': device, 'dtype': dtype}

        # Initialize weights to 1
        self.weight = nn.Parameter(torch.ones(d_model, **factory_kwargs))
    
    def forward(self, x: Tensor) -> Tensor:

        # Prevent overflow in mean/sqrt calculations
        in_dtype = x.dtype
        x = x.to(torch.float32)

        # Perform RMSNorm calculation 
        
        # official implementation:
        # rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) 
        # normalized_x = x * rms

        # our implementation:
        RMS = (x.pow(2).mean(dim=-1, keepdim=True)+ self.eps).sqrt() 
        normalized_x = x / RMS
        
        results = normalized_x * self.weight # W will automatically broadcast to ..., d_model

        # Return the result in the original dtype
        return results.to(in_dtype)

adapters.py里

def run_rmsnorm(
    d_model: int,
    eps: float,
    weights: Float[Tensor, " d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    device, dtype = in_features.device, in_features.dtype
    model = RMSNorm(d_model, eps, device=device, dtype=dtype)
    model.load_state_dict({'weight': weights})
    return model(in_features)

测试

uv run pytest -k test_rmsnorm

PASSED

pytest里覆盖的情况并不是很广，如果不转换dtype或者一些细小的东西不加进来，也会显示pass，所以尽量贴近作业的要求。

关于keep_dim 和broadcast

pytorch里的keep_dim的默认是False，如果做什么sum或者mean在某个维度，会把那个dim自动削减掉。

在RMSNorm的代码中：

RMS = (x**2).mean(dim=-1, keepdim=True).clamp(min=self.eps).sqrt()

这里如果不加keepdim=True会报错

RMS = (x**2).mean(dim=-1)        # shape becomes [batch, seq_len]

x / RMS 的话，pytorch会试图broadcast[batch, seq_len] 到 [batch, seq_len, d_model]

但是broadcast只能应用于尾端对齐的时候，上面这个情况的尾端显然没有对齐，所以无法对齐

加keepdim的情况：

x.shape       = [batch, seq_len, d_model]      # e.g., [2, 4, 8]
RMS.shape     = [batch, seq_len, 1]            # broadcastable ✅

pytorch会把RMS最后那个最后那个维度复制，然后变成[batch, seq_len, d_model]

如果RMS一开始是[2, 4, 1]，那么broadcast之后就会变成[2, 4, 8] 如果x最后一维是8的话。

RMS vs. L2

注意RMS和L2有点像，但不一样

RMS在根号里多了1/d

添加图片注释，不超过 140 字（可选）

L2：

添加图片注释，不超过 140 字（可选）

它们的关系：

添加图片注释，不超过 140 字（可选）

所以RMS 可以理解为是一种scaled version of L2 norm

L2

L2 用||x||2来表示，也可以理解为对自身的dot product的sqrt

添加图片注释，不超过 140 字（可选）

L2也代表着vector length, 也可以理解为距离原点的Euclidean distance 欧几里得距离。

normalize的时候，用自身除以||x||2即可

添加图片注释，不超过 140 字（可选）

L2的导

添加图片注释，不超过 140 字（可选）

根据

添加图片注释，不超过 140 字（可选）

求导过程：

添加图片注释，不超过 140 字（可选）

从上面不难发现，对向量 x 做 L2 normalize和对 L2求导的数学公式都一样，不过意义不同。前者用于标准化、余弦相似度等，后者用于反向传播中的梯度计算。

和L1不同的是，L2丝滑可导，所以deep learning里很多时候都用L2。

L2 square

MSE，会用squared L2，也就是没有根号的L2

||x||右边下面的2代表的是L2，上面的2代表square

MSE的公式：

添加图片注释，不超过 140 字（可选）

MSE的导非常光滑，就是2x

添加图片注释，不超过 140 字（可选）

代码例子

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
loss = torch.sum(x**2)  # equivalent to ||x||_2^2
loss.backward()
print(x.grad)  # tensor([2.0, 4.0, 6.0])

如果初始x是[1,2,3]，那么L2 square之后，backward的导数就是在每个值上乘以2,变成了[2,4,6]

3.5.2 SwiGLU

SwiGLU公式

添加图片注释，不超过 140 字（可选）

• d_model：transformer的dimension
• d_ff：feed-forward network （FFN）的dimension
• 通常d_ff是d_model的4倍或者2.6倍的 （第三课，超参）
• 本质就是把x从d_model上放大到d_ff，然后再缩小回到d_model

其中GLU，是element-wise乘积

添加图片注释，不超过 140 字（可选）

其中SiLU

添加图片注释，不超过 140 字（可选）

要求

• 写SwiGLU FFN的实现，包括SiLU func和GLU func
• 确保dff/dmodel 大概是8/3， 同时是64的倍数; 然而作业里的dff是128然后dmodel是64，不是8/3倍数，是2倍

答案

def sigmoid(x: Tensor): return 1 / (1 + torch.exp(-x)) # sigmoid activation function

def silu(x: Tensor): return x * torch.sigmoid(x) # SiLU activation function

def glu(a:Tensor, b:Tensor): return a * b # element-wise multiplication

def swiglu_fn(a: Tensor, b: Tensor): return glu(silu(a), b) # SwiGLU activation function

class SwiGLU(nn.Module):
    def __init__(self, d_model: int, d_ff: int, device=None, dtype=None):
        super().__init__()
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.linear1 = Linear(d_model, d_ff, **factory_kwargs)  # W1
        self.linear2 = Linear(d_ff, d_model, **factory_kwargs)  # W2
        self.linear3 = Linear(d_model, d_ff, **factory_kwargs)  # W3

    def forward(self, x: Float[Tensor, "... d_model"]) -> Float[Tensor, "... d_model"]:
        w1x = self.linear1(x)
        w3x = self.linear3(x)
        h = swiglu_fn(w1x, w3x)
        return self.linear2(h)

又写了一个根据算d_ff （8/3倍的d_model）的func：

def get_compatible_dff(d_model: int) -> int:
    """
    Returns the nearest multiple of 64 to 8/3 * d_model.
    """
    raw = (8 * d_model) / 3
    rounded = int((raw + 32) // 64) * 64  # round to nearest multiple of 64
    return rounded

adapters.py里

SiLU

def run_silu(in_features: Float[Tensor, " ..."]) -> Float[Tensor, " ..."]:
    return silu(in_features)

SwiGLU

def run_swiglu(
    d_model: int,
    d_ff: int,
    w1_weight: Float[Tensor, " d_ff d_model"],
    w2_weight: Float[Tensor, " d_model d_ff"],
    w3_weight: Float[Tensor, " d_ff d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:

    # Optional: If you want to calculate d_ff from d_model based on 8/3 ratio
    d_ff_83 = get_compatible_dff(d_model)
    print(f"Expected d_ff to be {d_ff_83} based on d_model {d_model}; here we got d_ff to be {d_ff}.")

    device, dtype = in_features.device, in_features.dtype
    model = SwiGLU(d_model, d_ff, device=device, dtype=dtype)
    model.load_state_dict({
        "linear1.weight": w1_weight,
        "linear2.weight": w2_weight,
        "linear3.weight": w3_weight,
    })
    return model(in_features)

测试

uv run pytest -k test_silu_matches_pytorch # test silu
uv run pytest -k test_swiglu # test swiglu

PASS

element-wise multiply

⊙,⊗,* 都可以表示element-wise multiply，(其中，⊗也可以指outer product，容易混淆)

比如

import torch
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[10, 20], [30, 40]])
print(a * b)
# Output: tensor([[10, 40],
#                 [90, 160]])

需要两个tensor的shape相同

dot product 点积

对于两个长度相同的一维vector，先element-wise multiply，再sum

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
print(torch.dot(a, b))  # Output: 1*4 + 2*5 + 3*6 = 32

如果是两个matrix AB，dot(A,B)相当于矩阵乘法。

matrix multiply

@是matrix multiply，或者dot product

@,matrix multiplication, 适用于二维的，是左边绿色的行，和右边绿色的列进行dot product，得到最右侧C1的指，C1的列数等于b的列数，C1的行数等于a的行数。

添加图片注释，不超过 140 字（可选）

添加图片注释，不超过 140 字（可选）

outer product

外积是将两个 1D 向量进行“所有元素之间的乘积”形成一个矩阵

添加图片注释，不超过 140 字（可选）

接下来的RoPE里的position ⊗ freq 相当于两个1D的vector进行外积

添加图片注释，不超过 140 字（可选）

outer product也相当于在a和b多加一个维度，然后矩阵相乘；

或者a@b.T (自动broadcast，然后矩阵相乘）

a[:, None] @ b[None, :]
a @ b.T 

3.5.3 Rotary Position Embeddings (RoPE)

要求

• 对于sin和cos值的tensor，用self.register_buffer(persistent=False)，而不要用nn.Parameter，因为它们是固定的值，设置persistent=False意味着不需要在保存模型的时候保存这些值在state_dict里，每次现算就可以，反正都是一样的。

答案

class RotaryPositionalEmbedding(nn.Module):
    """
    Rotary Position Embedding (RoPE) layer.

    Args
    ----
    theta : float
        Base used to generate inverse frequencies (e.g. 10_000).
    d_k : int
        Dimension of the key / query vectors (must be even).
    max_seq_len : int
        Maximum sequence length expected at inference / training time.
    device : torch.device | None
        Where to place the pre-computed sine / cosine tables.
    """
    def __init__(self,
                 theta: float,
                 d_k: int,
                 max_seq_len: int,
                 device=None):
        super().__init__()
        if d_k % 2 != 0:
            raise ValueError("d_k must be even for RoPE.")
        self.d_k = d_k
        # ---- pre-compute inverse frequencies ----
        # freq[k] = 1 / theta ** (2k / d_k)          (k = 0,1,…,d_k/2-1)
        freq = 1.0 / (theta ** (torch.arange(0,d_k,2, device=device).float() / d_k))

        # shape: (max_seq_len, d_k // 2)
        positions = torch.arange(max_seq_len, device=device).float()
        freqs = torch.outer(positions, freq)

        # cache cos/sin; no gradients needed → persistent=False
        self.register_buffer('cos_cached', torch.cos(freqs),persistent=False) # persistent=False does not save to state_dict
        self.register_buffer('sin_cached', torch.sin(freqs), persistent=False)
    
    def forward(
        self,
        x: Float[Tensor, "... seq_len d_k"],
        token_positions: Int[Tensor, "... seq_len"]
        ) -> Float[Tensor, "... seq_len d_k"]:
        """
        Apply RoPE to `x`.  Works with any batch shape prefix.
        """
        # Check if the last dimension matches d_k
        if x.size(-1) != self.d_k:
            raise ValueError(f"Last dim of x ({x.size(-1)}) ≠ d_k ({self.d_k}).")
        
        # Gather the cached tables for the required positions
        cos_pos = self.cos_cached[token_positions]
        sin_pos = self.sin_cached[token_positions]

        # Split even / odd channels
        x_even = x[..., ::2]
        x_odd = x[..., 1::2]

        # Apply the 2-D rotation to each pair
        out_even = x_even * cos_pos - x_odd * sin_pos
        out_odd = x_even * sin_pos + x_odd * cos_pos

        # Re-interleave
        out = torch.empty_like(x)
        out[..., ::2] = out_even
        out[..., 1::2] = out_odd
        return out

adapters.py里

def run_rope(
    d_k: int,
    theta: float,
    max_seq_len: int,
    in_query_or_key: Float[Tensor, " ... sequence_length d_k"],
    token_positions: Int[Tensor, " ... sequence_length"],
) -> Float[Tensor, " ... sequence_length d_k"]:

    # Create and initialize the RoPE module
    device = in_query_or_key.device
    rope = RotaryPositionalEmbedding(theta=theta, d_k=d_k, max_seq_len=max_seq_len, device=device)

    # Apply RoPE to the input tensor
    return rope(in_query_or_key, token_positions)

测试

uv run pytest -k test_rope

PASS

RoPE的具体讲解

3.5.4a Softmax

self-attention里面有一个很重要的函数就是softmax

添加图片注释，不超过 140 字（可选）

要求

• 先写一个softmax函数
• softmax里，要用到减去最大值的方法，来避免overflow

softmax

• exponentiate 每个数
• 然后normalize by sum of them

公式：

添加图片注释，不超过 140 字（可选）

然而现实中如果实现，会有overflow 和underflow的风险，比如e^8就会超过电脑能计算的digits，以至于overflow; e^过小就会导致等于0，也就是underflow。

所以一个普遍的方法就是让exp(xi)减去一个x vector的最大值，这样可以确保exponentiate的值不会太大

添加图片注释，不超过 140 字（可选）

为什么左右两边相等？因为如果c是max的话，e^-c 可以在分子和分母上下抵消:

添加图片注释，不超过 140 字（可选）

减去max，exp（）里面的数永远小于等于0，那么exp后的数就在0到1之间，避免了overflow

代码

def softmax_stable(x: Tensor, dim: int = -1) -> Tensor:
    """Numerically stable softmax."""
    x_max = x.max(dim=dim, keepdim=True).values
    x_exp = torch.exp(x - x_max)
    return x_exp / x_exp.sum(dim=dim, keepdim=True)

adapters.py里

def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]:
    return softmax_stable(in_features, dim=dim)

测试

uv run pytest -k test_softmax
uv run pytest -k test_softmax_matches_pytorch

PASS

3.5.4b Scaled Dot-Product Attention

添加图片注释，不超过 140 字（可选）

要求

• Q, K, V的shape：

添加图片注释，不超过 140 字（可选）

• 在attention matrix加一个True/False的mask，shape为 (n,m)
• 需要把False的值转换成-inf

代码

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k: int):
        super().__init__()
        self.scale = 1.0 / math.sqrt(d_k)

    def forward(
        self,
        query: Float[Tensor, "... seq_len_q d_k"],
        key: Float[Tensor, "... seq_len_k d_k"],
        value: Float[Tensor, "... seq_len_k d_v"],
        mask: Bool[Tensor, "seq_len_q seq_len_k"] = None
    ) -> Float[Tensor, "... seq_len_q d_v"]:
        # Compute scaled dot product attention scores using einsum
        attn_scores = einsum(query, key, "... q d, ... k d -> ... q k") * self.scale

        if mask is not None:
            attn_scores = attn_scores.masked_fill(~mask, float("-inf"))

        attn_probs = softmax_stable(attn_scores, dim=-1)

        # Compute attention output using einsum again
        output = einsum(attn_probs, value, "... q k, ... k d -> ... q d")

        return output

adapters.py里

def run_scaled_dot_product_attention(
    Q: Float[Tensor, " ... queries d_k"],
    K: Float[Tensor, " ... keys d_k"],
    V: Float[Tensor, " ... values d_v"],
    mask: Float[Tensor, " ... queries keys"] | None = None,
) -> Float[Tensor, " ... queries d_v"]:
    model = ScaledDotProductAttention(d_k=Q.size(-1))
    return model(Q, K, V, mask=mask)

测试

uv run pytest -k test_scaled_dot_product_attention
uv run pytest -k test_4d_scaled_dot_product_attention

Q@K.T 分母为什么放d_k sqrt?

这个和我们前面讲linear weights initialization很像。Q@K.T后，会出现mean为0，Var为d_k （σ为d_k sqrt）的distribution。为了compensate Var，让它回归到Var为1，我们需要让result值除以σ （参考Z score，(x-μ)/σ）,就可以得到N(0,1)的值了。

等式右边a^2是d_k在外面，等式左边的a就是d_k sqrt；为了compensate，用1/d_k sqrt

关于为啥乘一个a时，方差变为原来的a^2倍：

添加图片注释，不超过 140 字（可选）

Q 和KV的seq_len不同

作业的Q的seq_len和K和V的不同：Q的seq_len是n， K和V是m

如果Q和KV的source不同，可能是cross-attention；如果Q和KV的source相同，就是targeted attention

添加图片注释，不超过 140 字（可选）

cross-attention

Q的source和KV是不一样的，很多时候它们的seq length不相同。

cross-attention里是encoder和decoder，适用于翻译的情景

比如从法语翻译成英文:

• Input (source，encoder): "Je suis étudiant"
• Output (target, decoder): "I am a student"

decoder Q如果已经生成I am， 那么n就是2

encoder K V是三个法语单词，那么m就是3

mask, convert 0 to -inf

为什么要用mask？

• Padding masks: 不让模型pay attention to padded tokens
• Causal masks: 在autoregressive 模型里, 把后面的信息罩住，保证模型只能看到前面的tokens。

为什么不用0而用-inf？

添加图片注释，不超过 140 字（可选）

exp(0)=1，模型还是会看它。

exp(-inf)=0，模型会ignore它。

为什么V可以有d_v，不同于Q，K的d_k?

Q和K.T因为要matrix multiplication，dimension一定是相同的。

但是对于Q来说，它的dimension不需要和它们相同。attention weights计算完成后，weighted sum over the values。

3.5.5 Causal Multi-Head Self-Attention

要求

• 加入head, head dim从d_model提取出来，和batch融到一起
• d_k和d_v =d_model/h
• RoPE加入到QK里，但不要加到V里；make RoPE optional in class
• 加入causal mask，把右上三角盖住

代码

class CausalMultiHeadSelfAttention(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        max_seq_len: int,
        rope_theta: float = 10000.0,
        use_rope: bool = True,
        device=None,
        dtype=None
    ):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.d_v = self.d_k  # match d_k for simplicity
        self.use_rope = use_rope

        factory_kwargs = {'device': device, 'dtype': dtype}
        self.q_proj, self.k_proj, self.v_proj, self.o_proj = [Linear(d_model, d_model, **factory_kwargs)
                                                              for _ in range(4)]
        self.attn = ScaledDotProductAttention(self.d_k)

        # Create a causal mask for the attention mechanism
        # Shape: (1, 1, max_seq_len, max_seq_len)
        mask = torch.tril(torch.ones(max_seq_len, max_seq_len, dtype=torch.bool, device=device))
        self.register_buffer("causal_mask", mask.unsqueeze(0).unsqueeze(0), persistent=False)

        if use_rope:
            self.rope = RotaryPositionalEmbedding(
                theta=rope_theta, d_k=self.d_k, max_seq_len=max_seq_len, device=device)

    def forward(
        self,
        x: Float[Tensor, "batch seq_len d_model"],
        token_positions: Int[Tensor, "batch seq_len"]| None = None,
    ) -> Float[Tensor, "batch seq_len d_model"]:
        B, S, _ = x.shape

        # Project to multi-head Q, K, V
        q,k,v = [rearrange(proj(x), "b s (h d) -> b h s d", h=self.num_heads) 
                 for proj in [self.q_proj, self.k_proj, self.v_proj]]

        # Apply RoPE to Q and K if enabled
        if self.use_rope: q,k = self.rope(q, token_positions),self.rope(k, token_positions)

        # Compute attention
        out = self.attn(q, k, v, mask=self.causal_mask[..., :S, :S])

        # Merge heads and project
        out = rearrange(out, "b h s d -> b s (h d)")
        return self.o_proj(out)

adapters.py里

def run_multihead_self_attention(
    d_model: int,
    num_heads: int,
    q_proj_weight: Float[Tensor, " d_k d_in"],
    k_proj_weight: Float[Tensor, " d_k d_in"],
    v_proj_weight: Float[Tensor, " d_v d_in"],
    o_proj_weight: Float[Tensor, " d_model d_v"],
    in_features: Float[Tensor, " ... sequence_length d_in"],
) -> Float[Tensor, " ... sequence_length d_out"]:
    device, dtype = in_features.device, in_features.dtype
    max_seq_len = in_features.shape[-2]  # Assuming in_features is of shape (..., sequence_length, d_in)
    model = CausalMultiHeadSelfAttention(d_model=d_model,
                                          num_heads=num_heads, 
                                          use_rope=False,
                                          max_seq_len=max_seq_len,
                                          device=device,
                                          dtype=dtype)
    model.load_state_dict({
        "q_proj.weight": q_proj_weight,
        "k_proj.weight": k_proj_weight,
        "v_proj.weight": v_proj_weight,
        "o_proj.weight": o_proj_weight,
    })
    return model(in_features)

然后还有下一个

def run_multihead_self_attention_with_rope(
    d_model: int,
    num_heads: int,
    max_seq_len: int,
    theta: float,
    q_proj_weight: Float[Tensor, " d_k d_in"],
    k_proj_weight: Float[Tensor, " d_k d_in"],
    v_proj_weight: Float[Tensor, " d_v d_in"],
    o_proj_weight: Float[Tensor, " d_model d_v"],
    in_features: Float[Tensor, " ... sequence_length d_in"],
    token_positions: Int[Tensor, " ... sequence_length"] | None = None,
) -> Float[Tensor, " ... sequence_length d_out"]:
    device, dtype = in_features.device, in_features.dtype
    model = CausalMultiHeadSelfAttention(
        d_model=d_model, 
        num_heads=num_heads, 
        rope_theta=theta,
        use_rope=(token_positions is not None),
        max_seq_len=max_seq_len,
        device=device,
        dtype=dtype
    )
    model.load_state_dict({
        "q_proj.weight": q_proj_weight,
        "k_proj.weight": k_proj_weight,
        "v_proj.weight": v_proj_weight,
        "o_proj.weight": o_proj_weight,
    })

    return model(in_features,token_positions)

测试

uv run pytest -k test_multihead_self_attention

PASS two

tril vs. triu

torch.tril(torch.ones(4, 4)) ， l是lower，keep lower， causal mask

tensor([[1, 0, 0, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 1]])

torch.triu(input)， u是upper，keep upper， future mask

tensor([[1, 1, 1, 1],
        [0, 1, 1, 1],
        [0, 0, 1, 1],
        [0, 0, 0, 1]])

3.6a Transformer Block

要求

• 有2个sublayers
• 每层，都有RMSNorm，MHA/FF，还有residual connection
• layer1: y= x + Causal Multi-Head Self-Attention(RMSNorm(x))
• layer2: z = y+SwiGLU(RMSNorm(y))

添加图片注释，不超过 140 字（可选）

代码

class TransformerBlock(nn.Module):
    """
    Pre-norm Transformer block with two sub-layers:

       x ──► RMSNorm ──► MHA ──► + ──►
         │                     ▲
         └─────────────────────┘     (sublayer-1)

       y ──► RMSNorm ──► FF  ──► + ──► out
         │                     ▲
         └─────────────────────┘     (sublayer-2)
    """
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        d_ff: int,
        max_seq_len: int,
        rope_theta: float = 10_000.0,
        use_rope: bool = True,
        device=None,
        dtype=None,
    ) -> None:
        super().__init__()
        kwargs = {"device": device, "dtype": dtype}

        # ── sub-layer 1: (RMSNorm → causal MHA) ──────────────────────────────
        self.norm1 = RMSNorm(d_model, **kwargs)
        self.attn  = CausalMultiHeadSelfAttention(
            d_model=d_model,
            num_heads=num_heads,
            max_seq_len=max_seq_len,
            rope_theta=rope_theta,
            use_rope=use_rope,
            **kwargs,
        )

        # ── sub-layer 2: (RMSNorm → feed-forward) ────────────────────────────
        self.norm2 = RMSNorm(d_model, **kwargs)
        self.ff    = SwiGLU(d_model=d_model, d_ff=d_ff, **kwargs)

    # -----------------------------------------------------------------------
    def forward(
        self,
        x: torch.Tensor,               # (batch, seq_len, d_model)
        token_positions: torch.Tensor | None = None,  # (batch, seq_len)
    ) -> torch.Tensor:
        b, s, _ = x.shape

        # ---- sub-layer-1: RMSNorm → MHA → residual -------------------------
        attn_out = self.attn(self.norm1(x), token_positions=token_positions)
        x = x + attn_out                       # residual connection

        # ---- sub-layer-2: RMSNorm → FF → residual --------------------------
        ff_out   = self.ff(self.norm2(x))
        x        = x + ff_out                  # residual connection
        return x

adapters.py里

def run_transformer_block(
    d_model: int,
    num_heads: int,
    d_ff: int,
    max_seq_len: int,
    theta: float,
    weights: dict[str, Tensor],
    in_features: Float[Tensor, " batch sequence_length d_model"],
) -> Float[Tensor, " batch sequence_length d_model"]:
    device, dtype = in_features.device, in_features.dtype
    model = TransformerBlock(
        d_model=d_model,
        num_heads=num_heads,
        d_ff=d_ff,
        use_rope=True,
        max_seq_len=max_seq_len,
        rope_theta=theta,
        device=device,
        dtype=dtype
    )

    # 2. Load the reference weights -----------------------------------------
    # Load attention projection weights
    model.attn.q_proj.weight.data.copy_(weights["attn.q_proj.weight"])
    model.attn.k_proj.weight.data.copy_(weights["attn.k_proj.weight"])
    model.attn.v_proj.weight.data.copy_(weights["attn.v_proj.weight"])
    model.attn.o_proj.weight.data.copy_(weights["attn.output_proj.weight"])

    # Load RMSNorm weights
    model.norm1.weight.data.copy_(weights["ln1.weight"])
    model.norm2.weight.data.copy_(weights["ln2.weight"])

    # Load FFN weights (transpose needed for 1 and 3)
    model.ff.linear1.weight.data.copy_(weights["ffn.w1.weight"])
    model.ff.linear2.weight.data.copy_(weights["ffn.w2.weight"]) 
    model.ff.linear3.weight.data.copy_(weights["ffn.w3.weight"])

    # _load_reference_weights(model.state_dict(), weights)

    # 3. Form token-position indices for RoPE -------------------------------
    B, S, _ = in_features.shape
    positions = torch.arange(S, device=device).expand(B, S)  # (B, S)

    # 4. Run the forward pass ----------------------------------------------
    out = model(in_features, token_positions=positions)      # (B, S, d_model)
    return out

测试

uv run pytest -k test_transformer_block

3.6b Transformer Language Model (LM)

要求

• args: vocab_size， context_length， num_layers
• follow graph below

添加图片注释，不超过 140 字（可选）

代码

def _copy_param(target: torch.Tensor, source: torch.Tensor) -> None:
    """
    Copy `source` into `target` in-place, transposing `source` if that
    is what makes the shapes line up.
    """
    if source.shape == target.shape:
        target.data.copy_(source)
    elif source.T.shape == target.shape:
        target.data.copy_(source.T)
    else:
        raise ValueError(f"Shape mismatch: cannot load parameter of shape {source.shape} "
                         f"into tensor of shape {target.shape}")
    

class TransformerLM(nn.Module):
    def __init__(self,
                 vocab_size: int,
                 context_length: int,
                 num_layers: int,
                 d_model: int,
                 num_heads: int,
                 d_ff: int,
                 rope_theta: float,
                 device=None,
                 dtype=None):
        super().__init__()
        kw = dict(device=device, dtype=dtype)

        # token embedding  (no separate pos-emb: RoPE lives inside blocks)
        self.tok_emb = Embedding(vocab_size, d_model, **kw)

        # L Transformer blocks
        self.blocks = nn.ModuleList([
            TransformerBlock(
                d_model=d_model,
                num_heads=num_heads,
                d_ff=d_ff,
                max_seq_len=context_length,
                rope_theta=rope_theta,
                use_rope=True,
                **kw,
            )
            for _ in range(num_layers)
        ])

        # final norm
        self.ln_final = RMSNorm(d_model, **kw)
        self.lm_head = Linear(d_model, vocab_size, **kw)

        self.context_length = context_length

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        b, s = token_ids.shape
        if s > self.context_length:
            raise ValueError(f"seq_len {s} exceeds context_length {self.context_length}")

        # token embeddings
        x = self.tok_emb(token_ids)                         # (b, s, d)

        # token positions for RoPE
        pos = torch.arange(s, device=token_ids.device).unsqueeze(0).expand(b, s)

        # transformer stack
        for blk in self.blocks:
            x = blk(x, token_positions=pos)                # (b, s, d)

        # final norm → tied linear projection (logits)
        x = self.ln_final(x)                                 # (b, s, d)

        logits = self.lm_head(x)  # (b, s, vocab_size)
        return logits

adapters.py

from cs336_basics.layer import _copy_param

def run_transformer_lm(
    vocab_size: int,
    context_length: int,
    d_model: int,
    num_layers: int,
    num_heads: int,
    d_ff: int,
    rope_theta: float,
    weights: dict[str, Tensor],
    in_indices: Int[Tensor, " batch_size sequence_length"],
) -> Float[Tensor, " batch_size sequence_length vocab_size"]:

    device = in_indices.device
    dtype  = next(iter(weights.values())).dtype             # assume all same dtype

    # 1) construct the model skeleton
    model = TransformerLM(
        vocab_size     = vocab_size,
        context_length = context_length,
        num_layers     = num_layers,
        d_model        = d_model,
        num_heads      = num_heads,
        d_ff           = d_ff,
        rope_theta     = rope_theta,
        device         = device,
        dtype          = dtype,
    ).eval()                                                 # inference mode

    # 2) load the weights ----------------------------------------------------
    with torch.no_grad():
        # (a) token embedding  (also implicitly ties lm_head)
        _copy_param(model.tok_emb.weight,
                    weights["token_embeddings.weight"])

        # (b) per-layer parameters
        for layer_idx in range(num_layers):
            pfx   = f"layers.{layer_idx}."
            block = model.blocks[layer_idx]

            # ── attention projections
            _copy_param(block.attn.q_proj.weight,  weights[pfx + "attn.q_proj.weight"])
            _copy_param(block.attn.k_proj.weight,  weights[pfx + "attn.k_proj.weight"])
            _copy_param(block.attn.v_proj.weight,  weights[pfx + "attn.v_proj.weight"])
            _copy_param(block.attn.o_proj.weight,  weights[pfx + "attn.output_proj.weight"])

            # ── RMSNorm weights
            _copy_param(block.norm1.weight, weights[pfx + "ln1.weight"])
            _copy_param(block.norm2.weight, weights[pfx + "ln2.weight"])

            # ── feed-forward (SwiGLU) weights
            _copy_param(block.ff.linear1.weight, weights[pfx + "ffn.w1.weight"])
            _copy_param(block.ff.linear2.weight, weights[pfx + "ffn.w2.weight"])
            _copy_param(block.ff.linear3.weight, weights[pfx + "ffn.w3.weight"])

        # (c) final layer-norm
        _copy_param(model.ln_final.weight, weights["ln_final.weight"])

        # (d) (optional) make sure tied output embedding matches lm_head if provided
        _copy_param(model.lm_head.weight, weights["lm_head.weight"])

    # 3) run the forward pass and return logits
    with torch.no_grad():
        return model(in_indices)            # (batch, seq_len, vocab_size)

测试

uv run pytest -k test_transformer_lm

PASS

添加图片注释，不超过 140 字（可选）

To be continue..