1

git clone https://github.com/stanford-cs336/assignment1-basics.git

1
2
3
4

#先在WSL上安装uv
curl -LsSf https://astral.sh/uv/install.sh | sh 
source $HOME/.local/bin/env
uv run pytest tests

1
2
3
4
5
6
7
8

PAT = r"""
'(?:[sdmt]|ll|ve|re) #英文缩写如 's,'d等
| ?\p{L}+ #单词
| ?\p{N}+ #阿拉伯数字
| ?[^\s\p{L}\p{N}]+ #标点符号
|\s+(?!\S) #行尾空白
|\s+ # 其他空白
"""

1
2
3

# 需要安装 `regex` 包
import regex as re
re.findall(PAT, "some text that i'll pre-tokenize")

1

['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

1
2
3
4

vocab = {i: bytes([i]) for i in range(256)}
    for token in special_tokens:
        if token not in vocab.values():
            vocab[len(vocab)] = token.encode("utf-8")

1
2
3
4
5
6
7
8
9

with open(input_path, "rb") as f:
        num_processes = os.cpu_count() or 4
        boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
    ranges = list(zip(boundaries[:-1], boundaries[1:]))
    with mp.Pool(processes=num_processes) as pool:
        results = pool.map(
            _pretokenize_range,
            [(input_path,s,e) for s,e in ranges],
        )

1
2
3
4
5
6
7
8
9

def split_keep_special(text: str, special_tokens: list[str]) -> list[str]:
    #这里需要补充的一点在于，这里之所以对special token进行保留，是因为后续在tokenizer中我们可以直接复用这个函数
    if not special_tokens:
        return [text]
    special_tokens = sorted(special_tokens, key=lambda x: -len(x))
    #需要避免special token中有|等正则符号
    pattern = "(" + "|".join(re.escape(token) for token in special_tokens) + ")"
    parts = re.split(pattern, text)
    return [p for p in parts if p != ""]

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

def _pretokenize_range(args: tuple[str | os.PathLike, int, int]) -> dict[bytes, int]:
    #首先将chunk按照special token进行split,避免出现跨doc的合并问题
    input_path, start, end = args
    token_counts = Counter()
    PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
    with open(input_path, "rb") as f:
        f.seek(start)
        to_read = end - start
        data = f.read(to_read)
        parts = split_keep_special(data.decode("utf-8", errors="ignore"), ["<|endoftext|>"])
    for part in parts:
        if part == "<|endoftext|>":
            #token_counts[part.encode("utf-8")] += 1
            continue
        for m in re.finditer(PAT, part):
            token_counts[m.group(0).encode("utf-8")] += 1
    return token_counts

 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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

def run_train_bpe(
    input_path: str | os.PathLike,
    vocab_size: int,
    special_tokens: list[str],
    **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """Given the path to an input corpus, run train a BPE tokenizer and
    output its vocabulary and merges.

    Args:
        input_path (str | os.PathLike): Path to BPE tokenizer training data.
        vocab_size (int): Total number of items in the tokenizer's vocabulary (including special tokens).
        special_tokens (list[str]): A list of string special tokens to be added to the tokenizer vocabulary.
            These strings will never be split into multiple tokens, and will always be
            kept as a single token. If these special tokens occur in the `input_path`,
            they are treated as any other string.

    Returns:
        tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
            vocab:
                The trained tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
                to bytes (token bytes)
            merges:
                BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
                representing that <token1> was merged with <token2>.
                Merges are ordered by order of creation.
    """
    #词表的初始化，初始时词表应该只有从字节串 token 到整数 ID 的一一映射,以及规定的special tokens
    vocab = {i: bytes([i]) for i in range(256)}
    merges = []
    for token in special_tokens:
        b = token.encode("utf-8")
        if b not in vocab.values():
            vocab[len(vocab)] = b
    
    with open(input_path, "rb") as f:
        num_processes = os.cpu_count() or 4
        boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
    ranges = list(zip(boundaries[:-1], boundaries[1:]))
    total_token_counts = Counter()
    with mp.Pool(processes=num_processes) as pool:
        results = pool.map(
            _pretokenize_range,
            [(input_path,s,e) for s,e in ranges],
        )
    for token_counts in results:
        total_token_counts.update(token_counts)
    word_freqs = total_token_counts
    word_symbols = {
        token: tuple(vocab[b] for b in token)
        for token in word_freqs.keys()
    }
    while len(vocab) < vocab_size:
        pair_counts = Counter()
        for token,freq in word_freqs.items():
            symbols = word_symbols[token]
            if len(symbols) < 2:
                continue
            for a,b in zip(symbols, symbols[1:]):
                pair_counts[(a,b)] += freq
        if not pair_counts:
            break
        
        #合并出现频次最高的pair,如果有多个pair出现频次相同，则选择字典序最大的那个
        best_pair = max(pair_counts.items(), key=lambda kv: (kv[1], kv[0]))[0]
        #更新vocab和word_symbols
        merges.append(best_pair)
        merged_token = best_pair[0] + best_pair[1]
        vocab[len(vocab)] = merged_token

        for token in list(word_symbols.keys()):
            symbols = word_symbols[token]
            if len(symbols) < 2:
                continue
            new_symbols = []
            i = 0
            while i < len(symbols):
                if i < len(symbols) - 1 and (symbols[i], symbols[i+1]) == best_pair:
                    new_symbols.append(merged_token)
                    i += 2
                else:
                    new_symbols.append(symbols[i])
                    i += 1
            word_symbols[token] = tuple(new_symbols)
    return vocab, merges

  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
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

def run_train_bpe(
    input_path: str | os.PathLike,
    vocab_size: int,
    special_tokens: list[str],
    **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """Given the path to an input corpus, run train a BPE tokenizer and
    output its vocabulary and merges.

    Args:
        input_path (str | os.PathLike): Path to BPE tokenizer training data.
        vocab_size (int): Total number of items in the tokenizer's vocabulary (including special tokens).
        special_tokens (list[str]): A list of string special tokens to be added to the tokenizer vocabulary.
            These strings will never be split into multiple tokens, and will always be
            kept as a single token. If these special tokens occur in the `input_path`,
            they are treated as any other string.

    Returns:
        tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
            vocab:
                The trained tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
                to bytes (token bytes)
            merges:
                BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
                representing that <token1> was merged with <token2>.
                Merges are ordered by order of creation.
    """
    #词表的初始化，初始时词表应该只有从字节串 token 到整数 ID 的一一映射,以及规定的special tokens
    vocab = {i: bytes([i]) for i in range(256)}
    merges = []
    for token in special_tokens:
        b = token.encode("utf-8")
        if b not in vocab.values():
            vocab[len(vocab)] = b
    
    with open(input_path, "rb") as f:
        num_processes = os.cpu_count() or 4
        boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
    ranges = list(zip(boundaries[:-1], boundaries[1:]))
    total_token_counts = Counter()
    with mp.Pool(processes=num_processes) as pool:
        results = pool.map(
            _pretokenize_range,
            [(input_path,s,e) for s,e in ranges],
        )
    for token_counts in results:
        total_token_counts.update(token_counts)
    word_freqs = total_token_counts
    word_symbols = {
        token: tuple(vocab[b] for b in token)
        for token in word_freqs.keys()
    }
    pair2token = defaultdict(set)
    pair_counts = Counter()
    for token,freq in word_freqs.items():
        symbols = word_symbols[token]
        if len(symbols) < 2:
            continue
        for a,b in zip(symbols, symbols[1:]):
            pair_counts[(a,b)] += freq
            pair2token[(a,b)].add(token)
    while len(vocab) < vocab_size:
        if not pair_counts:
            break
        
        #合并出现频次最高的pair,如果有多个pair出现频次相同，则选择字典序最大的那个
        best_pair = max(pair_counts.items(), key=lambda kv: (kv[1], kv[0]))[0]
        #更新vocab和word_symbols
        merges.append(best_pair)
        merged_token = best_pair[0] + best_pair[1]
        vocab[len(vocab)] = merged_token

        tokens_to_update = pair2token[best_pair]

        for token in list(tokens_to_update):
            freq = word_freqs[token]
            symbols = word_symbols[token]
            if len(symbols) < 2:
                continue
            new_symbols = []
            #这里和原来不同了，我们需要减去旧的pair的计数
            i = 0
            while i < len(symbols) - 1:
                p = (symbols[i], symbols[i+1])
                pair_counts[p] -= freq
                pair2token[p].discard(token)
                i += 1
            #然后进行合并
            i = 0
            while i < len(symbols):
                if i < len(symbols) - 1 and (symbols[i], symbols[i+1]) == best_pair:
                    new_symbols.append(merged_token)
                    i += 2
                else:
                    new_symbols.append(symbols[i])
                    i += 1
            word_symbols[token] = tuple(new_symbols)
            for i in range(len(new_symbols) - 1):
                p = (new_symbols[i], new_symbols[i+1])
                pair_counts[p] += freq
                pair2token[p].add(token)
    return vocab, merges

 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
35
36
37
38
39
40
41
42
43
44
45
46
47

import multiprocessing as mp
import regex as re
import os
import json
import pickle
from pathlib import Path
from typing import BinaryIO
from collections import Counter
from adapters import run_train_bpe
from common import gpt2_bytes_to_unicode

def main():
    filepath = "data/TinyStoriesV2-GPT4-train.txt"
    # filepath = "data/owt_train.txt"
    vocab_size = 10000
    special_tokens = ["<|endoftext|>"]
    vocab, merges = run_train_bpe(filepath, vocab_size, special_tokens)
    print(f"The longest tokens in the vocabulary are:{sorted(vocab.values(), key=len, reverse=True)[:10]}")
    output_dir = Path(__file__).resolve().parent / "outputs"
    output_dir.mkdir(parents=True, exist_ok=True)

    bytes_to_unicode = gpt2_bytes_to_unicode()

    def encode_token(token_bytes: bytes) -> str:
        return "".join(bytes_to_unicode[b] for b in token_bytes)

    vocab_items = sorted(vocab.items(), key=lambda kv: kv[0])
    vocab_json = {encode_token(token_bytes): token_id for token_id, token_bytes in vocab_items}

    vocab_path = output_dir / "trained_vocab.json"
    merges_path = output_dir / "trained_merges.txt"
    serialized_path = output_dir / "trained_vocab_merges.pkl"

    with open(vocab_path, "w", encoding="utf-8") as f:
        json.dump(vocab_json, f, ensure_ascii=False, indent=2)

    with open(merges_path, "w", encoding="utf-8") as f:
        for token_a, token_b in merges:
            f.write(encode_token(token_a))
            f.write(" ")
            f.write(encode_token(token_b))
            f.write("\n")

    with open(serialized_path, "wb") as f:
        pickle.dump({"vocab": vocab, "merges": merges}, f, protocol=pickle.HIGHEST_PROTOCOL)
if __name__ == "__main__":
    main()

 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
35
36
37
38
39
40

def encode(self,text):
        from .adapters import split_keep_special
        #对单个文本进行BPE编码，返回token id列表
        #首先对文本进行预分词(此时假设text为"Hello <PAD> world!")
        parts = split_keep_special(text, self.special_tokens)
        #此时parts = [""Hello ","<PAD>"," world!"]
        token_ids = []
        for part in parts:
            if part in self.special_tokens:
                token_ids.append(self.vocab_inv[part.encode("utf-8")])
                continue
            for m in re.finditer(self.PAT, part):
                token = m.group(0).encode("utf-8")
                #此时Token为b"Hello"
                symbols = [bytes([b]) for b in token]
                #此时symbols为[b"H",b"e",b"l",b"l",b"o"]

                #接下来进行BPE合并
                while True:
                    pairs = [(symbols[i], symbols[i+1]) for i in range(len(symbols)-1)]
                    if not pairs:
                        break
                    candidate_pairs = [p for p in pairs if p in self.merges]
                    if not candidate_pairs:
                        break
                    merged_token = min(candidate_pairs, key=lambda p: self.merges[p])
                    new_symbols = []
                    i = 0
                    while i < len(symbols):
                        if i < len(symbols) - 1 and (symbols[i], symbols[i+1]) == merged_token:
                            new_symbols.append(merged_token[0] + merged_token[1])
                            i += 2
                        else:
                            new_symbols.append(symbols[i])
                            i += 1
                    symbols = new_symbols
                #将合并后的symbols转换为token ids
                for sym in symbols:
                    token_ids.append(self.vocab_inv[sym])
        return token_ids

1
2
3
4
5
6

def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        #对输入的流式文本的指定范围进行BPE编码，返回token id生成器
        for text in iterable:
            token_ids = self.encode(text)
            for tid in token_ids:
                yield tid

1
2
3
4
5

def decode(self,token_ids):
        #将token id列表解码为文本,很简单，遍历一遍就行
        bytes_list = [self.vocab[tid] for tid in token_ids]
        text = b"".join(bytes_list).decode("utf-8",errors="replace")
        return text

1
2
3
4
5
6
7
8

import torch
from einops import rearrange, einsum
Y = D @ A.T
#很难看出输入输出的张量形状以及具体的含义
#若我们采用einsum，就很直观了
Y = einsum(D,A,"batch sequence d_in,d_out d_in -> batch sequence d_out")
#我们还有一个更加简便的例子:
Y = einsum(D,A,"... d_in,d_out d_in -> ... d_out")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

images = torch.randn(64,128,128,3)#(batch,height,width,channel)
dim_dy = torch.linspace(start = 0.0,end = 1.0,steps = 10)
## 通过reshape和逐元素相乘实现
dim_value = rearrange(dim_dy,"dim_value -> 1 dim_value 1 1")
image_rearr = rearrange(images,"b height width channel -> b 1 height width channel")
dimmed_images = images_rearr * dim_value
## 若我们通过enisum实现:
dimmed_images = enisum(
    images,dim_dy,
    "batch height width channel,dim_value -> batch dim_value height width channel"
)

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

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

#传统实现方法就是通过view + transpose
channels_last_flat = channels_last.view(
    -1,channels_last.size(1) * channels_last.size(2),channels_last.size(3)
)
channels_first_flat_transformed = channels_first_flat @ B.T
channels_last_flat_transformed = channels_first_flat_transformed.transpose(1, 2)
channels_last_transformed = channels_last_flat_transformed.view(*channels_last.shape)

#如果我们用enisum
height = width = 32
channels_last_transformed = einsum(
    channels_last,B,
    "batch h_in w_in channel,(h_out w_out)(h_in w_in) -> batch h_out w_out channel"
)

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

import torch
import math
from einops import rearrange, einsum
class Linear(torch.nn.Module):
    def __init__(self,in_features,out_features,device = None,dtype = None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.W = torch.nn.Parameter(torch.empty((out_features,in_features),device = device,dtype = dtype))
        sigma = math.sqrt(2 / (in_features + out_features))
        torch.nn.init.trunc_normal_(self.W,mean = 0.0,std = sigma,a = -3.0 * sigma,b = 3.0 * sigma)
    def forward(self,x):
        return einsum(x,self.W,"... d_in,d_out d_in->... d_out")

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

import torch
from einops import rearrange, einsum

class Embedding(torch.nn.Module):
    def __init__(self,num_embeddings,embedding_dim,device = None,dtype = None):
        """
        num_embeddings: int 表示词表大小
        embedding_dim: int 表示每个词向量的维度
        """
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.weight = torch.nn.Parameter(torch.empty((num_embeddings,embedding_dim),device = device,dtype = dtype))
        torch.nn.init.trunc_normal_(self.weight,mean = 0.0,std = 1.0,a = -3.0,b = 3.0)
    def forward(self,token_ids):
        """
        根据给定的token_ids返回对应的Embedding向量
        """
        return self.weight[token_ids]

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

import torch
from einops import rearrange, einsum
import math
class RMSNorm(torch.nn.Module):
    def __init__(self, d_model, eps = 1e-5,device = None,dtype = None):
        super().__init__()
        self.d_model = d_model
        self.eps = eps
        #依据实验指导书，初始化为全1
        self.weight = torch.nn.Parameter(torch.ones(d_model, device=device, dtype=dtype))
    def forward(self,x):
        #为了数值稳定，先将类型转化为float32
        in_dtype = x.dtype
        x = x.to(torch.float32)
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        x_normed = x / rms
        x_normed = x_normed * self.weight
        return x_normed.to(in_dtype)

 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

import torch
import torch.nn as nn
import math
from einops import rearrange, einsum

class SwiGLU(nn.Module):
    def __init__(self,d_model,device = None,dtype = None):
        super().__init__()
        self.d_model = d_model
        #根据指导手册要求，但是需要保证是64的整数倍
        d_ff = int(math.ceil((8/3) * d_model / 64) * 64)
        
        self.W1 = torch.nn.Parameter(torch.empty((d_ff,d_model),device=device,dtype=dtype))
        self.W2 = torch.nn.Parameter(torch.empty((d_model,d_ff),device=device,dtype=dtype))
        self.W3 = torch.nn.Parameter(torch.empty((d_ff,d_model),device=device,dtype=dtype))

        sigma1 = (2 / (d_model + d_ff)) ** 0.5
        sigma2 = (2 / (d_ff + d_model)) ** 0.5

        torch.nn.init.trunc_normal_(self.W1,mean = 0.0,std = sigma1,a = -3.0 * sigma1,b = 3.0 * sigma1)
        torch.nn.init.trunc_normal_(self.W2,mean = 0.0,std = sigma2,a = -3.0 * sigma2,b = 3.0 * sigma2)
        torch.nn.init.trunc_normal_(self.W3,mean = 0.0,std = sigma1,a = -3.0 * sigma1,b = 3.0 * sigma1)
    def forward(self,x):
        x_proj1 = einsum(x,self.W1,"... d_model,d_ff d_model->... d_ff")
        x_proj1 = x_proj1 * torch.sigmoid(x_proj1) #计算SiLU激活
        x_proj2 = einsum(x,self.W3,"... d_model,d_ff d_model->... d_ff")
        x_glu = x_proj1 * x_proj2
        out = einsum(x_glu,self.W2,"... d_ff,d_model d_ff->... d_model")
        return out

 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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

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

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self,theta,d_k,max_seq_len,device):
        """
        d_k: int, 维度大小，必须为偶数
        theta: float, RoPE中的\Theta值
        max_seq_len: int, 最大序列长度
        device: torch.device, 设备
        """
        super().__init__()
        assert d_k % 2 == 0, "d_k must be even"
        self.theta = theta
        self.d_k = d_k
        self.max_seq_len = max_seq_len
        self.device = device

        #一共有d / 2个频率
        half_dk = d_k // 2
        k = torch.arange(0,half_dk,device=device).float()
        inv_freq = 1.0 / (self.theta ** (2.0 * k / d_k))

        positions = torch.arange(0,max_seq_len,device = device).float()

        angles = einsum(positions,inv_freq,"max_seq_len,half_dk->max_seq_len half_dk")
        cos = torch.cos(angles)
        sin = torch.sin(angles)

        self.register_buffer("cos",cos,persistent = False)
        self.register_buffer("sin",sin,persistent = False)

    def forward(self,x,token_positions):
        """
        inputs:
            x: ...,seq_len,d_k
            token_positions:...,seq_len
        returns:
            x_rotated: ...,seq_len,d_k
        """
        cos = self.cos[token_positions]  # ...,seq_len,half_dk
        sin = self.sin[token_positions]  # ...,seq_len,half_dk

        x_even = x[...,0::2]
        x_odd = x[...,1::2]

        x_rot_even = x_even * cos - x_odd * sin
        x_rot_odd = x_even * sin + x_odd * cos

        out = torch.empty_like(x)
        out[...,0::2] = x_rot_even
        out[...,1::2] = x_rot_odd
        return out

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

def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]:
    """
    Given a tensor of inputs, return the output of softmaxing the given `dim`
    of the input.

    Args:
        in_features (Float[Tensor, "..."]): Input features to softmax. Shape is arbitrary.
        dim (int): Dimension of the `in_features` to apply softmax to.

    Returns:
        Float[Tensor, "..."]: Tensor of with the same shape as `in_features` with the output of
        softmax normalizing the specified `dim`.
    """
    max_num = torch.max(in_features,dim=dim,keepdim=True).values
    exp_tensor = torch.exp(in_features - max_num)
    sum_exp = torch.sum(exp_tensor,dim=dim,keepdim=True)
    return exp_tensor / sum_exp

 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

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: Bool[Tensor, " ... queries keys"] | None = None,
) -> Float[Tensor, " ... queries d_v"]:
    """
    Given key (K), query (Q), and value (V) tensors, return
    the output of your scaled dot product attention implementation.

    Args:
        Q (Float[Tensor, " ... queries d_k"]): Query tensor
        K (Float[Tensor, " ... keys d_k"]): Key tensor
        V (Float[Tensor, " ... values d_v"]): Values tensor
        mask (Bool[Tensor, " ... queries keys"] | None): Mask tensor
    Returns:
        Float[Tensor, " ... queries d_v"]: Output of SDPA
    """
    d_k = Q.shape[-1]
    scores = einsum(Q,K,"... q d_k,... k d_k -> ... q k") / (d_k ** 0.5)
    if mask is not None:
        #对于mask的，我们直接将mask为False的位置设置为负无穷
        scores = scores.masked_fill(~mask,float("-inf"))
    attn_weights = run_softmax(scores,dim=-1)
    return einsum(attn_weights,V,"... q k,... k d_v -> ... q d_v")

 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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

class CausalMultiHeadSelfAttention(nn.Module):
    def __init__(self,d_model,num_heads,seq_len,device,use_rope=True):
        """
        d_model: int,模型维度
        num_heads: int,注意力头数
        device: torch.device,设备
        use_rope: bool,是否使用RoPE
        """
        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.device = device
        self.theta = 10000.0

        self.W_q = Linear(self.d_k,self.d_model,self.device)
        self.W_k = Linear(self.d_k,self.d_model,self.device)
        self.W_v = Linear(self.d_k,self.d_model,self.device)
        self.W_o = Linear(self.d_model,self.d_k,self.device)
        if use_rope:
            self.use_rope = True
            self.RoPE = RotaryPositionalEmbedding(theta=self.theta,d_k = self.d_k, max_seq_len = seq_len, device = device)
        else:
            self.use_rope = False

    def forward(self,x,token_positions):
        """
        inputs:
        x: Float[Tensor, "batch_size seq_len d_model"]
        token_positions: Long[Tensor, "batch_size seq_len"]
        returns:
        out: Float[Tensor, "batch_size seq_len d_model"]
        """
        batch_size,seq_len,_ = x.shape

        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        
        Q = rearrange(Q,"b s (h d_k) -> b h s d_k",h=self.num_heads)
        K = rearrange(K,"b s (h d_k) -> b h s d_k",h=self.num_heads)
        V = rearrange(V,"b s (h d_k) -> b h s d_k",h=self.num_heads)
        if self.use_rope:
            Q = self.RoPE(Q,token_positions)
            K = self.RoPE(K,token_positions)
        
        mask = torch.tril(torch.ones((seq_len,seq_len),device=self.device)).bool()
        attn_output = run_scaled_dot_product_attention(Q,K,V,mask=mask)
        attn_output = rearrange(attn_output,"b h s d_k -> b s (h d_k)")
        out = self.W_o(attn_output)
        return out

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

import torch.nn as nn
from cs336_basics.CausalMultiHeadSelfAttention import CausalMultiHeadSelfAttention
from cs336_basics.Linear import Linear
from cs336_basics.RMSNorm import RMSNorm
from cs336_basics.SwiGLU import SwiGLU

class TransformerBlock(nn.Module):
    def __init__(self,d_model,num_heads,seq_len,ffn_hidden_dim,device,use_rope=True):
        super().__init__()
        self.attention = CausalMultiHeadSelfAttention(d_model,num_heads,seq_len,device,use_rope)
        self.ffn = SwiGLU(d_model,device)
        # pass device explicitly to avoid treating it as eps
        self.attn_norm = RMSNorm(d_model, device=device)
        self.ffn_norm = RMSNorm(d_model, device=device)
    def forward(self,x,token_positions):
        x_norm = self.attn_norm(x)
        attn_out = self.attention(x_norm,token_positions)
        x = x + attn_out
        x_norm = self.ffn_norm(x)
        ffn_out = self.ffn(x_norm)
        x = x + ffn_out
        return x

 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

from cs336_basics.Transformer_Block import TransformerBlock
from cs336_basics.Embedding import Embedding
from cs336_basics.RMSNorm import RMSNorm
from cs336_basics.Linear import Linear
import torch.nn as nn
import torch
from einops import rearrange, einsum

class Transformer:
    def __init__(self,vocab_size,context_length,num_layers,d_model,num_heads,device):
        self.TokenEmbedding = Embedding(vocab_size,d_model,device)
        self.TransformerBlocks = nn.ModuleList([
            TransformerBlock(d_model,num_heads,context_length,ffn_hidden_dim=None,device=device,use_rope=True)
            for _ in range(num_layers)
        ])
        self.FinalNorm = RMSNorm(d_model,device=device)
        self.OutputLayer = Linear(d_model,vocab_size,device)
    def forward(self,token_ids):
        """
        inputs:
        token_ids: Long[Tensor, "batch_size seq_len"]
        returns:
        logits: Float[Tensor, "batch_size seq_len vocab_size"]
        """
        batch_size,seq_len = token_ids.shape
        x = self.TokenEmbedding(token_ids)  # [batch_size, seq_len, d_model]
        token_positions = torch.arange(seq_len,device=x.device).unsqueeze(0).expand(batch_size,-1)  # [batch_size, seq_len]
        for block in self.TransformerBlocks:
            x = block(x,token_positions)
        x = self.FinalNorm(x)
        logits = self.OutputLayer(x)
        return logits

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

import torch
from einops import rearrange, einsum

def cross_entropy(logits, targets):
    """
    logits: Float[Tensor, "batch vocab_size"]
    targets: Int[Tensor, "batch"]
    """
    max_logits = logits.max(dim = -1,keepdim = True).values #shape: [batch_size,1]
    #如果直接对logits减去最大值，最后会因为log操作出现下溢的情况
    log_sum_exp = torch.log(torch.exp(logits - max_logits).sum(dim=-1,keepdim=True)) + max_logits  # [batch,1]
    # [batch,1] -> [batch]
    log_sum_exp = rearrange(log_sum_exp,'b 1 -> b')
    loss = log_sum_exp - logits[torch.arange(logits.shape[0]),targets]
    return loss.mean()

 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
35
36
37
38
39
40
41
42

import torch
from einops import rearrange, einsum
import math

class Adamw(torch.optim.Optimizer):
    def __init__(self,params,lr=1e-3,betas=(0.9,0.999),eps=1e-8,weight_decay=0.01):
        defaults = dict(lr=lr,betas=betas,eps=eps,weight_decay=weight_decay)
        super().__init__(params,defaults)
    def step(self,closure=None):
        loss = None
        for group in self.param_groups:
            lr = group['lr']
            betas = group['betas']
            eps = group['eps']
            weight_decay = group['weight_decay']
            
            # step 放在 group 级别，所有参数共享
            step = group.get('step', 0) + 1
            group['step'] = step
            
            # 计算偏差修正后的学习率
            lr_t = lr * math.sqrt(1 - betas[1] ** step) / (1 - betas[0] ** step)
            
            for p in group['params']:
                if p.grad is None:
                    continue
                state = self.state[p]  # 只包含一阶矩、二阶矩
                m = state.get('m', torch.zeros_like(p))
                v = state.get('v', torch.zeros_like(p))
                
                # 更新一阶矩、二阶矩
                m = betas[0] * m + (1 - betas[0]) * p.grad
                v = betas[1] * v + (1 - betas[1]) * p.grad * p.grad

                # 更新参数
                p.data = p.data - lr_t * (m / (torch.sqrt(v) + eps))
                # 参数衰减
                p.data = p.data * (1 - weight_decay * lr)
                
                state['m'] = m
                state['v'] = v
        return loss

1
2
3
4
5
6
7
8
9

import math

def lr_cosine_schedule(t,a_max,a_min,t_w,t_c):
    if t < t_w:
        return t / t_w * a_max
    elif t <= t_c:
        return a_min + (a_max - a_min) * (1 + math.cos(math.pi * (t - t_w) / (t_c - t_w))) / 2
    else:
        return a_min

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import math
import torch

def gradient_clip(params, max_norm, epsilon=1e-6):
    total_norm = torch.sqrt(sum([torch.norm(p.grad.data, p=2) ** 2 for p in params if p.grad is not None]))
    scale = max_norm / (total_norm + epsilon)
    if scale < 1:
        for p in params:
            if p.grad is not None:
                p.grad.data *= scale

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

from torch.utils.data import Dataset
import numpy as np
import torch

def get_batch(x, batch_size, context_length, device):
    """
    input:
        x:  Int[Tensor, "seq_len"] 或 numpy array
        batch_size: int
        context_length: int
        device: torch.device
    output:
        xb: Int[Tensor, "batch_size context_length"]
        yb: Int[Tensor, "batch_size context_length"]
    """
    if not isinstance(x, torch.Tensor):
        x = torch.tensor(x)
    
    idx = torch.randint(0, len(x) - context_length, size=(batch_size,))
    xb = torch.stack([x[i:i+context_length] for i in idx])
    yb = torch.stack([x[i+1:i+context_length+1] for i in idx])
    xb, yb = xb.to(device), yb.to(device)
    return xb, yb

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

def save_checkpoint(model,optimizer,epoch, out):
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
    }, out)

def load_checkpoint(src,model,optimizer):
    checkpoint = torch.load(src)
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint["epoch"]