train-llm-from-scratch 完整笔记

1
def process_files(input_dir, output_file, tokenizer_name, max_data=None):
2
    enc = tiktoken.get_encoding(tokenizer_name)   # 默认 'r50k_base'（GPT-3 用的分词器）
3

4
    with h5py.File(output_file, 'w') as out_f:
5
        # 创建一个一维、可动态扩容的 dataset，存所有 token
6
        dataset = out_f.create_dataset('tokens', (0,), maxshape=(None,), dtype='i')
7
        start_index = 0
8

9
        for filename in sorted(os.listdir(input_dir)):
10
            if filename.endswith(".jsonl.zst"):
11
                with zstd.open(in_file, 'rt', encoding='utf-8') as in_f:  # 流式解压 + 文本模式读
12
                    for line in tqdm(in_f, desc=f"Processing {filename}", total=max_data):
13
                        data = json.loads(line)
14
                        text = data.get('text')
15
                        if text:
16
                            # 每条文本末尾追加 <|endoftext|>，再编码
17
                            encoded = enc.encode(text + "<|endoftext|>", allowed_special={'<|endoftext|>'})
18
                            encoded_len = len(encoded)
19
                            end_index = start_index + encoded_len
20
                            dataset.resize(dataset.shape[0] + encoded_len, axis=0)  # 扩容
21
                            dataset[start_index:end_index] = encoded                # 追加写入
22
                            start_index = end_index
23
                        processed_lines += 1
24
                        if max_data is not None and processed_lines >= max_data:
25
                            break  # 每个文件最多处理 max_data 行

1
def get_batch_iterator(data_path, batch_size, context_length, device="cpu"):
2
    with h5py.File(data_path, 'r') as hdf5_file:
3
        dataset = hdf5_file['tokens']
4
        dataset_size = dataset.shape[0]
5
        # 能切出多少个不重叠的样本（-1 是因为目标要往后错一位）
6
        n_examples = (dataset_size - 1) // context_length
7
        example_idxs = np.arange(n_examples)
8
        np.random.shuffle(example_idxs)   # 打乱样本顺序
9

10
        epochs = 0
11
        counter = 0
12
        while True:                       # 无限循环，训练循环自己控制何时停
13
            if counter + batch_size > n_examples:
14
                np.random.shuffle(example_idxs)   # 一个 epoch 用完，重新打乱
15
                counter = 0
16
                print(f"Finished epoch {epochs}")
17
                epochs += 1
18

19
            # 把"第几个样本"换算成 token 流里的起始下标
20
            random_indices = example_idxs[counter:counter+batch_size] * context_length
21
            # 每个样本取 context_length+1 个 token（多取一个用于错位）
22
            random_samples = torch.tensor(np.array(
23
                [dataset[idx:idx+context_length+1] for idx in random_indices]
24
            ))
25
            xb = random_samples[:, :context_length].to(device)    # 输入：前 context_length 个
26
            yb = random_samples[:, 1:context_length+1].to(device) # 目标：错后一位
27
            counter += batch_size
28
            yield xb, yb

1
class MLP(nn.Module):
2
    def __init__(self, n_embed):
3
        super().__init__()
4
        self.hidden = nn.Linear(n_embed, 4 * n_embed)  # 升维到 4 倍
5
        self.relu = nn.ReLU()
6
        self.proj = nn.Linear(4 * n_embed, n_embed)    # 投影回原维度
7

8
    def forward(self, x):
9
        x = self.forward_embedding(x)   # hidden + relu
10
        x = self.project_embedding(x)   # proj
11
        return x
12

13
    def forward_embedding(self, x):
14
        return self.relu(self.hidden(x))
15

16
    def project_embedding(self, x):
17
        return self.proj(x)

1
class Head(nn.Module):
2
    def __init__(self, head_size, n_embed, context_length):
3
        super().__init__()
4
        self.key   = nn.Linear(n_embed, head_size, bias=False)   # K 投影
5
        self.query = nn.Linear(n_embed, head_size, bias=False)   # Q 投影
6
        self.value = nn.Linear(n_embed, head_size, bias=False)   # V 投影
7
        # 下三角矩阵，注册成 buffer（不是参数，不训练，但随模型存取/搬设备）
8
        self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
9

10
    def forward(self, x):
11
        B, T, C = x.shape
12
        head_size = self.key.out_features
13
        k = self.key(x)      # (B, T, head_size)
14
        q = self.query(x)    # (B, T, head_size)
15
        scale_factor = 1 / math.sqrt(head_size)
16
        # 注意力分数：q 和 k 做点积
17
        attn_weights = q @ k.transpose(-2, -1) * scale_factor       # (B, T, T)
18
        # 因果掩码：把"未来"位置置成 -inf
19
        attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
20
        attn_weights = F.softmax(attn_weights, dim=-1)              # 每行归一化成概率
21
        v = self.value(x)    # (B, T, head_size)
22
        out = attn_weights @ v                                      # (B, T, head_size)
23
        return out

1
class Transformer(nn.Module):
2
    def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
3
        super().__init__()
4
        self.context_length = context_length
5
        self.N_BLOCKS = N_BLOCKS
6
        self.token_embed    = nn.Embedding(vocab_size, n_embed)        # token → 向量
7
        self.position_embed = nn.Embedding(context_length, n_embed)    # 位置 → 向量
8
        self.attn_blocks    = nn.ModuleList(
9
            [Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)]
10
        )
11
        self.layer_norm = nn.LayerNorm(n_embed)              # 最后一层归一化
12
        self.lm_head    = nn.Linear(n_embed, vocab_size)     # 投影到词表，得到 logits
13
        self.register_buffer('pos_idxs', torch.arange(context_length))
14

15
    def _pre_attn_pass(self, idx):
16
        B, T = idx.shape
17
        tok_embedding = self.token_embed(idx)                 # (B, T, C)
18
        pos_embedding = self.position_embed(self.pos_idxs[:T]) # (T, C)，广播相加
19
        return tok_embedding + pos_embedding
20

21
    def forward(self, idx, targets=None):
22
        x = self._pre_attn_pass(idx)         # 词嵌入 + 位置嵌入
23
        for block in self.attn_blocks:       # 逐层 Transformer Block
24
            x = block(x)
25
        x = self.layer_norm(x)               # 最终归一化
26
        logits = self.lm_head(x)             # (B, T, vocab_size)
27
        loss = None
28
        if targets is not None:
29
            B, T, C = logits.shape
30
            flat_logits = logits.view(B * T, C)
31
            targets = targets.view(B * T).long()
32
            loss = F.cross_entropy(flat_logits, targets)   # 交叉熵
33
        return logits, loss

1
class DemoChatDataset(Dataset):
2
    def __getitem__(self, idx):
3
        conversations = self.samples[idx]
4
        # 用 chat template 把整段对话拼接 + 分词
5
        input_ids = self.tokenizer.apply_chat_template(
6
            conversations, tokenize=True, add_generation_prompt=False,
7
            return_tensors="pt", max_length=self.max_length,
8
            truncation=True, padding="max_length"
9
        ).squeeze(0)
10

11
        loss_mask = torch.zeros_like(input_ids, dtype=torch.long)   # 默认全 0（都不算 loss）
12
        bos_assistant_ids = self.tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
13
        eos_ids = self.tokenizer.encode("<|im_end|>", add_special_tokens=False)
14

15
        # 扫描 token 序列，把每段 "assistant\n ... <|im_end|>" 之间的 token 标成 1
16
        i = 0
17
        input_ids_list = input_ids.tolist()
18
        while i < len(input_ids_list):
19
            # 命中一段 assistant 回复的起始标记
20
            if input_ids_list[i : i+len(bos_assistant_ids)] == bos_assistant_ids:
21
                start_of_response = i + len(bos_assistant_ids)
22
                # 往后找该回复的结束标记 <|im_end|>
23
                end_marker = -1
24
                j = start_of_response
25
                while j < len(input_ids_list):
26
                    if input_ids_list[j : j+len(eos_ids)] == eos_ids:
27
                        end_marker = j
28
                        break
29
                    j += 1
30
                if end_marker != -1:                                # 找到结束标记
31
                    loss_mask[start_of_response : end_marker + len(eos_ids)] = 1
32
                    i = end_marker + len(eos_ids)                   # 跳到这段之后继续扫
33
                    continue
34
                else:                                               # 没找到（被截断），标到末尾
35
                    loss_mask[start_of_response:] = 1
36
                    break
37
            i += 1
38

39
        loss_mask[input_ids == self.tokenizer.pad_token_id] = 0   # padding 不算
40
        X = input_ids[:-1]; Y = input_ids[1:]
41
        mask_for_loss_calculation = loss_mask[1:]
42
        return X, Y, mask_for_loss_calculation

1
class RotaryEmbedding(nn.Module):
2
    def __init__(self, dim, max_seq_len, theta=10000.0, device=None):
3
        super().__init__()
4
        # θ_i = 1 / theta^(2i/d)，i = 0,1,...,d/2-1
5
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))   # (d/2,)
6
        t = torch.arange(max_seq_len).float()                             # (max_seq_len,)
7
        freqs = torch.outer(t, freqs)                                     # (max_seq_len, d/2)，元素 = m·θ_i
8
        # 用模长 1、角度 m·θ_i 构造复数 cos(m·θ_i) + i·sin(m·θ_i)
9
        self.register_buffer("freqs_cis", torch.polar(torch.ones_like(freqs), freqs), persistent=False)
10

11
    def forward(self, xq, xk, seq_len):
12
        # xq/xk: (bsz, num_heads, seq_len, head_dim)
13
        # 把相邻两维拼成复数：(..., head_dim) → (..., head_dim/2) 复数
14
        xq_c = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
15
        xk_c = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
16
        freqs_cis_pos = self.freqs_cis[:seq_len].unsqueeze(0).unsqueeze(0)  # (1,1,seq_len,d/2)
17
        # 复数乘法 = 旋转
18
        xq_out = torch.view_as_real(xq_c * freqs_cis_pos).flatten(3)
19
        xk_out = torch.view_as_real(xk_c * freqs_cis_pos).flatten(3)
20
        return xq_out.type_as(xq), xk_out.type_as(xk)

1
# 辅助方法：把 KV 头复制 n_rep 次，(bsz, num_kv_heads, seq, d) → (bsz, num_kv_heads*n_rep, seq, d)
2
def _repeat_kv(self, x, n_rep):
3
    bs, num_kv_heads, slen, head_dim = x.shape
4
    if n_rep == 1:
5
        return x
6
    return (x[:, :, None, :, :]
7
            .expand(bs, num_kv_heads, n_rep, slen, head_dim)
8
            .reshape(bs, num_kv_heads * n_rep, slen, head_dim))
9

10
def forward(self, hidden_states, attention_mask=None, position_ids=None,
11
            past_key_value=None, use_cache=False):
12
    bsz, q_len, _ = hidden_states.shape
13
    # 1) 投影 + 拆头：(bsz, heads, q_len, head_dim)
14
    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_q_heads,  self.head_dim).transpose(1, 2)
15
    key_states   = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
16
    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
17

18
    # 总的 K/V 序列长度（含 KV cache 里的历史 token）
19
    kv_seq_len = q_len
20
    if past_key_value is not None:
21
        kv_seq_len += past_key_value[0].shape[2]
22

23
    # 2) 给 Q、K 施加 RoPE（按当前新 token 的长度）
24
    query_states, key_states = self.rotary_emb(query_states, key_states, seq_len=q_len)
25

26
    # 3) KV cache：把历史的 K/V 拼到前面（推理加速）
27
    if past_key_value is not None:
28
        key_states   = torch.cat([past_key_value[0], key_states], dim=2)
29
        value_states = torch.cat([past_key_value[1], value_states], dim=2)
30
    current_key_value = (key_states, value_states) if use_cache else None
31

32
    # 4) GQA：把 KV 头复制 num_kv_groups 次，匹配 Q 头数
33
    key_states   = self._repeat_kv(key_states,   self.num_kv_groups)
34
    value_states = self._repeat_kv(value_states, self.num_kv_groups)
35

36
    # 5) 算注意力：优先 Flash，否则手动
37
    if self.flash_available and attention_mask is None:
38
        attn_output = F.scaled_dot_product_attention(
39
            query_states, key_states, value_states,
40
            dropout_p=self.config.dropout if self.training else 0.0,
41
            is_causal=(q_len == kv_seq_len))        # 无 KV cache 时即标准因果注意力
42
    else:
43
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
44
        # 因果掩码；causal_shift 兼容 KV cache（q_len < kv_seq_len）的情形
45
        causal_shift = kv_seq_len - q_len
46
        mask = torch.triu(torch.full((q_len, kv_seq_len), float("-inf"),
47
                                     device=query_states.device), diagonal=1 + causal_shift)
48
        attn_weights = attn_weights + mask[None, None, :, :]
49
        if attention_mask is not None:              # 额外的 padding 掩码（加性）
50
            attn_weights = attn_weights + attention_mask
51
        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(query_states)
52
        attn_output = torch.matmul(attn_weights, value_states)
53

54
    # 6) 合并头 + 输出投影
55
    attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
56
    return self.o_proj(attn_output), current_key_value