Add: ch5.3 code

docs/chapter5/code/model.py

@@ -0,0 +1,430 @@
+import math
+import struct
+import inspect
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+@dataclass
+class ModelArgs:
+    # 自定义超参数
+    dim: int = 288  # 模型维度
+    n_layers: int = 6  # Transformer层数
+    n_heads: int = 6  # 注意力机制的头数
+    n_kv_heads: Optional[int] = 6  # 键/值头数，如果未指定，则默认为n_heads
+    vocab_size: int = 32000  # 词汇表大小
+    hidden_dim: Optional[int] = None  # 隐藏层维度，如果未指定，则使用其他规则确定
+    multiple_of: int = 32  # MLP隐藏层大小是这个数的倍数
+    norm_eps: float = 1e-5  # 归一化层的epsilon值
+    max_seq_len: int = 256  # 最大序列长度
+    dropout: float = 0.0  # 丢弃率
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        # eps是为了防止除以0的情况
+        self.eps = eps
+        # weight是一个可学习的参数，全部初始化为1
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        # 计算RMSNorm的核心部分
+        # x.pow(2).mean(-1, keepdim=True)计算了输入x的平方的均值
+        # torch.rsqrt是平方根的倒数，这样就得到了RMSNorm的分母部分，再加上eps防止分母为0
+        # 最后乘以x，得到RMSNorm的结果
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        # forward函数是模型的前向传播
+        # 首先将输入x转为float类型，然后进行RMSNorm，最后再转回原来的数据类型
+        # 最后乘以weight，这是RMSNorm的一个可学习的缩放因子
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+# 获得旋转嵌入的实部和虚部
+# 注意：此处的dim应为 dim//n_head，因为我们是对每个head进行旋转嵌入
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    # torch.arange(0, dim, 2)[: (dim // 2)].float()生成了一个从0开始，步长为2的序列，长度为dim的一半
+    # 然后每个元素除以dim，再取theta的倒数，得到频率
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    # 生成一个从0到end的序列，长度为end
+    t = torch.arange(end, device=freqs.device)
+    # 计算外积，得到一个二维矩阵，每一行是t的元素乘以freqs的元素
+    freqs = torch.outer(t, freqs).float()
+    # 计算频率的余弦值，得到实部
+    freqs_cos = torch.cos(freqs)
+    # 计算频率的正弦值，得到虚部
+    freqs_sin = torch.sin(freqs)
+    return freqs_cos, freqs_sin
+
+# 此函数的作用是将freqs_cis调整为与x的形状相同，以便能够与x进行广播操作
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    # 获取x的维度数
+    ndim = x.ndim
+    # 断言，确保1在x的维度范围内
+    assert 0 <= 1 < ndim
+    # 断言，确保freqs_cis的形状与x的第二维和最后一维相同
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    # 构造一个新的形状，除了第二维和最后一维，其他维度都为1，这样做是为了能够将freqs_cis与x进行广播操作
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    # 将freqs_cis调整为新的形状，并返回
+    return freqs_cis.view(shape)
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    # 将查询和键张量转换为浮点数，并重塑形状以分离实部和虚部
+    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
+    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
+
+    # 重新塑形频率张量以进行广播
+    freqs_cos = reshape_for_broadcast(freqs_cos, xq_r)
+    freqs_sin = reshape_for_broadcast(freqs_sin, xq_r)
+
+    # 应用旋转，分别计算旋转后的实部和虚部
+    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
+    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
+    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
+    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos
+
+    # 将最后两个维度合并，并还原为原始张量的形状
+    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
+    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    # 获取输入张量的形状：批量大小、序列长度、键/值对头的数量、每个头的维度大小
+    bs, slen, n_kv_heads, head_dim = x.shape
+    
+    # 如果重复次数为1，则不需要重复，直接返回原始张量
+    if n_rep == 1:
+        return x
+    
+    # 对张量进行扩展和重塑操作以重复键值对
+    return (
+        x[:, :, :, None, :]  # 在第四个维度（头的维度前）添加一个新的维度
+        .expand(bs, slen, n_kv_heads, n_rep, head_dim)  # 将新添加的维度扩展到n_rep大小，实现重复的效果
+        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)  # 重新塑形，合并键/值对头的数量和重复次数的维度
+    )
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        # 根据是否指定n_kv_heads，确定用于键（key）和值（value）的头的数量。
+        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
+        # 确保总头数可以被键值头数整除。
+        assert args.n_heads % self.n_kv_heads == 0
+
+        # 模型并行处理大小，默认为1。
+        model_parallel_size = 1
+        # 本地计算头数，等于总头数除以模型并行处理大小。
+        self.n_local_heads = args.n_heads // model_parallel_size
+        # 本地键值头数，等于键值头数除以模型并行处理大小。
+        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
+        # 重复次数，用于扩展键和值的尺寸。
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        # 每个头的维度，等于模型维度除以头的总数。
+        self.head_dim = args.dim // args.n_heads
+
+        # 定义权重矩阵。
+        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
+        # 输出权重矩阵。
+        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
+
+        # 定义dropout。
+        self.attn_dropout = nn.Dropout(args.dropout)
+        self.resid_dropout = nn.Dropout(args.dropout)
+        # 保存dropout概率。
+        self.dropout = args.dropout
+
+        # 检查是否使用Flash Attention（需要PyTorch >= 2.0）。
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            # 若不支持Flash Attention，则使用手动实现的注意力机制，并设置mask。
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            # 创建一个上三角矩阵，用于遮蔽未来信息。
+            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
+            mask = torch.triu(mask, diagonal=1)
+            # 注册为模型的缓冲区
+            self.register_buffer("mask", mask)
+
+    def forward(self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor):
+        # 获取批次大小和序列长度，[batch_size, seq_len, dim]
+        bsz, seqlen, _ = x.shape
+
+        # 计算查询（Q）、键（K）、值（V）。
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        # 调整形状以适应头的维度。
+        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
+
+        # 应用旋转位置嵌入（RoPE）。
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
+
+        # 对键和值进行扩展以适应重复次数。
+        xk = repeat_kv(xk, self.n_rep)
+        xv = repeat_kv(xv, self.n_rep)
+
+        # 将头作为批次维度处理。
+        xq = xq.transpose(1, 2)
+        xk = xk.transpose(1, 2)
+        xv = xv.transpose(1, 2)
+
+        # 根据是否支持Flash Attention，选择实现方式。
+        if self.flash:
+            # 使用Flash Attention。
+            output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
+        else:
+            # 使用手动实现的注意力机制。
+            scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
+            assert hasattr(self, 'mask')
+            scores = scores + self.mask[:, :, :seqlen, :seqlen]
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            scores = self.attn_dropout(scores)
+            output = torch.matmul(scores, xv)
+
+        # 恢复时间维度并合并头。
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+
+        # 最终投影回残差流。
+        output = self.wo(output)
+        output = self.resid_dropout(output)
+        return output
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
+        super().__init__()
+        # 如果没有指定隐藏层的维度，我们将其设置为输入维度的4倍
+        # 然后将其减少到2/3，最后确保它是multiple_of的倍数
+        if hidden_dim is None:
+            hidden_dim = 4 * dim
+            hidden_dim = int(2 * hidden_dim / 3)
+            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        # 定义第一层线性变换，从输入维度到隐藏维度
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        # 定义第二层线性变换，从隐藏维度到输入维度
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        # 定义第三层线性变换，从输入维度到隐藏维度
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        # 定义dropout层，用于防止过拟合
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # 前向传播函数
+        # 首先，输入x通过第一层线性变换和SILU激活函数
+        # 然后，结果乘以输入x通过第三层线性变换的结果
+        # 最后，通过第二层线性变换和dropout层
+        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+    
+
+class DecoderLayer(nn.Module):
+    def __init__(self, layer_id: int, args: ModelArgs):
+        super().__init__()
+        # 定义多头注意力的头数
+        self.n_heads = args.n_heads
+        # 定义输入维度
+        self.dim = args.dim
+        # 定义每个头的维度，等于输入维度除以头数
+        self.head_dim = args.dim // args.n_heads
+        # 定义LLaMA2Attention对象，用于进行多头注意力计算
+        self.attention = Attention(args)
+        # 定义LLaMAMLP对象，用于进行前馈神经网络计算
+        self.feed_forward = MLP(
+            dim=args.dim,
+            hidden_dim=args.hidden_dim,
+            multiple_of=args.multiple_of,
+            dropout=args.dropout,
+        )
+        # 定义层的ID
+        self.layer_id = layer_id
+        # 定义注意力计算的归一化层
+        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
+        # 定义前馈神经网络计算的归一化层
+        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
+
+    def forward(self, x, freqs_cos, freqs_sin):
+        # 前向传播函数
+        # 首先，输入x经过注意力归一化层，然后进行注意力计算，结果与输入x相加得到h
+        # 然后，h经过前馈神经网络归一化层，然后进行前馈神经网络计算，结果与h相加得到输出
+        h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
+        out = h + self.feed_forward.forward(self.ffn_norm(h))
+        return out
+
+class Transformer(nn.Module):
+    last_loss: Optional[torch.Tensor]
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        # 初始化模型参数
+        self.args = args
+        # 词汇表大小
+        self.vocab_size = args.vocab_size
+        # 层数
+        self.n_layers = args.n_layers
+
+        # 词嵌入层
+        self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)
+        # Dropout层
+        self.dropout = nn.Dropout(args.dropout)
+        # Decoder层
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(args.n_layers):
+            self.layers.append(DecoderLayer(layer_id, args))
+        # 归一化层
+        self.norm = RMSNorm(args.dim, eps=args.norm_eps)
+        # 输出层
+        self.output = nn.Linear(args.dim, args.vocab_size, bias=False)
+
+        # 将词嵌入层的权重与输出层的权重共享
+        self.tok_embeddings.weight = self.output.weight 
+
+        # 预计算相对位置嵌入的频率
+        freqs_cos, freqs_sin = precompute_freqs_cis(self.args.dim // self.args.n_heads, self.args.max_seq_len)
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+        # 初始化所有权重
+        self.apply(self._init_weights)
+        # 对残差投影进行特殊的缩放初始化
+        for pn, p in self.named_parameters():
+            if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * args.n_layers))
+
+        # 初始化最后一次前向传播的损失属性
+        self.last_loss = None
+
+    def _init_weights(self, module):
+        # 初始化权重的函数
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+    
+    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None) -> torch.Tensor:
+        # 前向传播函数
+        _bsz, seqlen = tokens.shape
+        # 通过词嵌入层和Dropout层
+        h = self.tok_embeddings(tokens)
+        h = self.dropout(h)
+        # 获取相对位置嵌入的频率
+        freqs_cos = self.freqs_cos[:seqlen]
+        freqs_sin = self.freqs_sin[:seqlen]
+
+        # 通过Decoder层
+        for layer in self.layers:
+            h = layer(h, freqs_cos, freqs_sin)
+        # 通过归一化层
+        h = self.norm(h)
+
+        if targets is not None:
+            # 如果给定了目标，计算损失
+            logits = self.output(h)
+            self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # 推理时的小优化：只对最后一个位置的输出进行前向传播
+            logits = self.output(h[:, [-1], :]) 
+            self.last_loss = None
+
+        return logits
+    
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # 获取所有需要更新的参数
+        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
+        
+        # 将参数分为需要权重衰减和不需要权重衰减的两组
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        
+        # 打印参数数量信息
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        
+        # 根据设备类型选择使用标准 AdamW 或其融合版本
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+    
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ 估计模型的 FLOPs 利用率 (MFU) 单位：A100 bfloat16 的峰值 FLOPS """
+        # 计算每次迭代的 FLOPs 数量（参考 PaLM 论文的附录 B）
+        # PaLM: Scaling Language Modeling with Pathways: https://arxiv.org/abs/2204.02311
+        N = sum(p.numel() for p in self.parameters())
+        cfg = self.args
+        L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.dim//cfg.n_heads, cfg.max_seq_len
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        
+        # 将 FLOPs 吞吐量表示为 A100 bfloat16 峰值 FLOPS 的比例
+        flops_achieved = flops_per_iter * (1.0/dt) # 每秒计算的 FLOPs
+        flops_promised = 312e12 # A100 GPU bfloat16 的峰值 FLOPS 为 312 TFLOPS
+        mfu = flops_achieved / flops_promised
+        return mfu
+    
+    @torch.inference_mode()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。
+        在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。
+        """
+        for _ in range(max_new_tokens):
+            # 如果序列上下文过长，截断它到最大长度
+            idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
+            
+            # 前向传播获取序列中最后一个位置的 logits
+            logits = self(idx_cond)
+            logits = logits[:, -1, :] # 只保留最后一个时间步的输出
+            
+            if temperature == 0.0:
+                # 选择最有可能的索引
+                _, idx_next = torch.topk(logits, k=1, dim=-1)
+            else:
+                # 缩放 logits 并应用 softmax
+                logits = logits / temperature
+                if top_k is not None:
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = -float('Inf')
+                probs = F.softmax(logits, dim=-1)
+                idx_next = torch.multinomial(probs, num_samples=1)
+            
+            # 将采样的索引添加到序列中并继续
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx
+
+if __name__ == '__main__':
+    args = ModelArgs()
+    # LLaMA2Model.forward 接受两个参数，tokens和targets，其中tokens是输入的张量, 应为int类型
+    x = torch.randint(0, 32000, (1, 50)) # [bs, seq_len]
+    # 实例化LLaMA2Model
+    model = Transformer(args=args)
+    # 计算model的全部参数
+    num_params = sum(p.numel() for p in model.parameters())
+    print('Number of parameters:', num_params)
+
+    out = model(x)
+    print(out.shape) # [batch_size, 1, vocab_size]

docs/chapter5/code/preprocess.py

@@ -0,0 +1,194 @@
+import glob
+import json
+import os
+import random
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+
+import numpy as np
+import sentencepiece as spm
+import torch
+import torch.distributed as dist
+from tqdm import tqdm
+
+from tokenizer import Tokenizer
+
+DATA_CACHE_DIR = 'data'
+TOKENIZER_MODEL = "./data/tok4096.model"
+
+
+# 定义分片处理函数
+def process_shard(args, vocab_size, tokenizer_model_path):
+    """
+    处理数据分片，将其中的文本进行分词并保存为二进制文件。
+    
+    参数:
+    args: tuple, 包含分片ID和分片文件名
+    vocab_size: int, 词汇表大小，用于决定输出文件存储路径
+    """
+    # 提取分片ID和文件名
+    shard_id, shard = args
+
+    # 初始化分词器
+    enc = Tokenizer(tokenizer_model_path)
+    
+    # 打开并读取当前分片的JSON文件
+    with open(shard, "r") as f:
+        data = json.load(f)
+    
+    # 用于保存所有的分词后的token
+    all_tokens = []
+    
+    # 遍历每一个例子，tqdm显示进度条
+    for example in tqdm(data, position=shard_id):
+        # 提取故事文本，并去除首尾空白字符
+        text = example["story"]
+        text = text.strip()  # 去掉首尾空白字符
+        
+        # 对文本进行编码，使用BOS（开始标志）但不使用EOS（结束标志）
+        tokens = enc.encode(text, bos=True, eos=False)
+        # 将当前文本的token添加到总token列表
+        all_tokens.extend(tokens)
+    
+    # 将所有的token转换为uint16类型的NumPy数组
+    all_tokens = np.array(all_tokens, dtype=np.uint16)
+    
+    # 根据词汇表大小确定输出文件名
+    if vocab_size == 0:
+        # 如果词汇表大小为0，使用默认的Llama 2分词器，将文件保存到原路径
+        tokenized_filename = shard.replace(".json", ".bin")
+    else:
+        # 如果有指定词汇表大小，保存到新目录`tok{vocab_size}`下
+        bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")
+        shard_basename = os.path.basename(shard)
+        bin_basename = shard_basename.replace(".json", ".bin")
+        tokenized_filename = os.path.join(bin_dir, bin_basename)
+    
+    # 将token以二进制形式保存
+    with open(tokenized_filename, "wb") as f:
+        f.write(all_tokens.tobytes())
+    
+    # 计算平均序列长度（以BOS标记`1`分隔的序列）
+    avg_seq_len = all_tokens.size / ((all_tokens == 1).sum())
+    print(f"Saved {tokenized_filename}, average seqlen: {avg_seq_len:.2f}")
+
+
+# 定义预处理函数，用于对多个数据分片进行批量处理
+def pretokenize(vocab_size):
+    """
+    预处理所有的数据分片，并将分词后的数据保存为二进制文件。
+    
+    参数:
+    vocab_size: int, 词汇表大小，用于决定输出文件存储路径
+    """
+    # 数据所在目录
+    data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
+    
+    # 获取所有JSON文件的文件名列表，并按字典序排序
+    shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
+    
+    # 如果词汇表大小大于0，则创建对应的保存目录
+    if vocab_size > 0:
+        bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")
+        os.makedirs(bin_dir, exist_ok=True)
+
+    # 使用partial函数将vocab_size绑定到process_shard函数
+    fun = partial(process_shard, vocab_size=vocab_size, tokenizer_model_path=TOKENIZER_MODEL)
+    
+    # 使用进程池并行处理每个分片
+    with ProcessPoolExecutor() as executor:
+        executor.map(fun, enumerate(shard_filenames))
+    
+    print("Done.")
+
+
+class PretokDataset(torch.utils.data.IterableDataset):
+    """从磁盘加载已预处理的分词数据，并将其以 PyTorch 张量的形式返回。"""
+
+    def __init__(self, split, max_seq_len, vocab_size, vocab_source):
+        """
+        初始化数据集。
+
+        参数:
+        split: str, 数据集的分割方式（'train' 或 'test'）。
+        max_seq_len: int, 最大序列长度，用于生成输入输出序列。
+        vocab_size: int, 词汇表的大小。
+        vocab_source: str, 词汇表的来源（'llama2' 或 'custom'）。
+        """
+        super().__init__()
+        self.split = split  # 数据集划分（训练集或测试集）
+        self.max_seq_len = max_seq_len  # 最大序列长度
+        self.vocab_size = vocab_size  # 词汇表大小
+        self.vocab_source = vocab_source  # 词汇表来源
+
+    def __iter__(self):
+        """
+        返回迭代器，按批次加载数据并生成模型输入/输出。
+        """
+        # 获取DataLoader的worker信息（用于并行数据加载）
+        worker_info = torch.utils.data.get_worker_info()
+        worker_id = worker_info.id if worker_info else 0  # worker ID
+        # 获取分布式训练的rank信息（用于多GPU训练）
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        # 基于worker_id和rank生成唯一的随机数种子，确保数据在每个worker和rank之间是唯一的
+        seed = 42 + worker_id + 1337 * rank
+        rng = random.Random(seed)
+        print(f"Created a PretokDataset with rng seed {seed}")
+
+        # 根据词汇表来源决定数据路径
+        if self.vocab_source == "llama2":
+            # 如果使用 Llama 2 词汇表，.bin 文件和 .json 文件在同一目录下
+            bin_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
+            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
+        elif self.vocab_source == "custom":
+            # 如果使用自定义词汇表，.bin 文件在 tok{N} 目录下
+            bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{self.vocab_size}")
+            shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
+
+        # 根据数据集划分使用不同的分片文件
+        # 训练集使用所有分片文件，测试集只使用第一个分片
+        shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1]
+        assert len(shard_filenames) > 0, f"在 {bin_dir} 中未找到任何 .bin 文件"
+
+        while True:
+            # 随机打乱分片文件
+            rng.shuffle(shard_filenames)
+            for shard in shard_filenames:
+                # 使用 memmap 读取文件，使得数据留在磁盘上，减少内存占用
+                m = np.memmap(shard, dtype=np.uint16, mode="r")
+                # 计算该分片中的批次数量
+                num_batches = len(m) // self.max_seq_len
+                num_batches -= 1  # 去掉最后一个不完整的批次
+                assert num_batches > 0, "这个分片文件太小了？请检查。"
+                # 随机打乱批次索引
+                ixs = list(range(num_batches))
+                rng.shuffle(ixs)
+                # 对每个批次生成输入 x 和目标输出 y
+                for ix in ixs:
+                    start = ix * self.max_seq_len  # 批次起始索引
+                    end = start + self.max_seq_len + 1  # 批次结束索引
+                    # 将数据转换为 NumPy 数组并拷贝到 RAM 中
+                    chunk = torch.from_numpy((m[start:end]).astype(np.int64))
+                    # 模型输入 x 是当前批次的前 max_seq_len 个词元
+                    x = chunk[:-1]
+                    # 模型输出 y 是下一个词元
+                    y = chunk[1:]
+                    # 生成 x, y 对
+                    yield x, y
+
+
+class Task:
+    @staticmethod
+    def iter_batches(batch_size, device, num_workers=0, **dataset_kwargs):
+        ds = PretokDataset(**dataset_kwargs)
+        dl = torch.utils.data.DataLoader(
+            ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers
+        )
+        for x, y in dl:
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+            yield x, y
+
+
+if __name__ == "__main__":
+    pretokenize(vocab_size=4096)

docs/chapter5/code/requirements.txt

@@ -0,0 +1,5 @@
+numpy==1.23.5
+Requests==2.31.0
+sentencepiece==0.1.99
+torch==2.0.1
+tqdm==4.64.1

docs/chapter5/code/sample.py

@@ -0,0 +1,104 @@
+import os
+import pickle
+from contextlib import nullcontext
+import torch
+from model import ModelArgs, Transformer
+from tokenizer import Tokenizer
+import argparse
+
+class TextGenerator:
+    def __init__(self, 
+                 checkpoint='output/ckpt.pt',  # 模型检查点路径
+                 tokenizer_model_path='tok4096.model',  # 分词器模型路径
+                 seed=1337,  # 随机种子，确保可重复性
+                 device=None,  # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU
+                 dtype="float32"):  # 数据类型，默认为 float32，可以选择 float16 或 bfloat16
+        """
+        初始化 TextGenerator 类，加载模型、设置设备和分词器等。
+        """
+        # 模型加载配置
+        self.checkpoint = checkpoint  # 保存的模型检查点路径
+        self.tokenizer_model_path = tokenizer_model_path  # 分词器模型文件路径
+        self.seed = seed  # 随机数种子，用于生成的可重复性
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')  # 根据硬件条件选择设备
+        self.dtype = dtype  # 模型的浮点数类型
+        self.device_type = 'cuda' if 'cuda' in self.device else 'cpu'  # 判断当前设备是否为 CUDA
+        
+        # 设置随机种子，确保生成的可重复性
+        torch.manual_seed(seed)  # 设置 CPU 随机种子
+        torch.cuda.manual_seed(seed)  # 设置 CUDA 随机种子
+        torch.backends.cuda.matmul.allow_tf32 = True  # 允许 CUDA 使用 TF32 精度进行矩阵乘法运算
+        torch.backends.cudnn.allow_tf32 = True  # 允许 cuDNN 使用 TF32 精度加速
+        
+        # 根据 dtype 选择适当的自动混合精度上下文
+        ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
+        self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype)
+        
+        # 加载模型检查点文件
+        checkpoint_dict = torch.load(self.checkpoint, map_location=self.device)  # 加载模型参数
+        gptconf = ModelArgs(**checkpoint_dict['model_args'])  # 初始化模型参数
+        self.model = Transformer(gptconf)  # 实例化 Transformer 模型
+        state_dict = checkpoint_dict['model']  # 获取模型状态字典
+        
+        # 去除状态字典中的不必要前缀
+        unwanted_prefix = '_orig_mod.'  # 这个前缀在保存时可能被添加，现在要去除它
+        for k, v in list(state_dict.items()):
+            if k.startswith(unwanted_prefix):
+                state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)  # 去除不必要的前缀
+        
+        # 加载模型参数到模型中
+        self.model.load_state_dict(state_dict, strict=False)
+        # 计算模型参数量
+        num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        print(f"Model has {num_params} parameters.")
+        # 设置模型为评估模式（evaluation mode），防止训练模式下的 dropout 等操作影响结果
+        self.model.eval()
+        # 将模型放置到正确的设备上（GPU 或 CPU）
+        self.model.to(self.device)
+        # 初始化分词器
+        self.tokenizer = Tokenizer(tokenizer_model=self.tokenizer_model_path)  # 根据指定的路径加载分词器
+
+    def sample(self, 
+               start="Hello!",  # 生成文本的起始提示词，可以是任意字符串
+               num_samples=3,  # 生成样本的数量，默认生成 3 个样本
+               max_new_tokens=256,  # 每个样本生成的最大 token 数，默认最多生成 256 个 token
+               temperature=1.0,  # 控制生成的随机性，1.0 为标准，值越大越随机
+               top_k=300):  # 保留概率最高的 top_k 个 token，限制生成时的选择范围
+        """
+        根据给定的起始文本生成样本。
+        
+        :param start: 生成文本的起始提示词
+        :param num_samples: 要生成的文本样本数
+        :param max_new_tokens: 每个样本生成的最大 token 数
+        :param temperature: 控制生成的随机性，值越小生成越确定，值越大生成越随机
+        :param top_k: 限制生成时选择的 token 范围
+        :return: 生成的文本样本列表
+        """
+        # 如果 start 是以 'FILE:' 开头，表示从文件中读取起始文本
+        if start.startswith('FILE:'):
+            with open(start[5:], 'r', encoding='utf-8') as f:
+                start = f.read()  # 读取文件内容作为起始文本
+        
+        # 将起始文本编码为 token id 序列
+        start_ids = self.tokenizer.encode(start, bos=True, eos=False)  # bos=True 表示加上句首标记，eos=False 表示不加句尾标记
+        x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...])  # 将编码后的 token id 转为 PyTorch 张量
+        
+        generated_texts = []  # 用于保存生成的文本样本
+        with torch.no_grad():  # 禁用梯度计算，提升效率
+            with self.ctx:  # 进入自动混合精度的上下文（如果是 GPU 并使用 float16 时）
+                for k in range(num_samples):  # 循环生成指定数量的样本
+                    y = self.model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)  # 生成文本
+                    generated_texts.append(self.tokenizer.decode(y[0].tolist()))  # 解码生成的 token 序列为可读文本
+        
+        return generated_texts  # 返回生成的文本样本
+
+# 示例使用
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt", type=str, default="One day, Lily met a Shoggoth")
+    args = parser.parse_args()
+
+    generator = TextGenerator()  # 初始化生成器
+    samples = generator.sample(start=args.prompt, num_samples=3, max_new_tokens=256)  # 生成 3 个样本
+    for i, sample in enumerate(samples):
+        print(f"\nSample {i+1}:\n{sample}\n{'-'*20}")  # 打印生成的样本并用分隔线分割

docs/chapter5/code/tokenizer.py

@@ -0,0 +1,68 @@
+import os
+import struct
+from sentencepiece import SentencePieceProcessor
+from typing import List
+
+TOKENIZER_MODEL = "./data/tok4096.model"
+
+class Tokenizer:
+    def __init__(self, tokenizer_model=None):
+        """
+        初始化分词器。加载预训练的SentencePiece模型，并设置一些特殊的token ID。
+
+        参数:
+        tokenizer_model: str, 可选，分词器模型的路径，如果不指定则使用默认路径 TOKENIZER_MODEL。
+        """
+        # 如果提供了分词器模型路径，使用该路径；否则使用默认模型路径
+        model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL
+        # 确保模型文件存在
+        assert os.path.isfile(model_path), model_path
+
+        # 加载 SentencePiece 模型
+        self.sp_model = SentencePieceProcessor(model_file=model_path)
+        self.model_path = model_path
+
+        # 获取分词器的特殊token和词汇表大小
+        self.n_words: int = self.sp_model.vocab_size()  # 词汇表大小
+        self.bos_id: int = self.sp_model.bos_id()       # 句子开头 (BOS) 的ID
+        self.eos_id: int = self.sp_model.eos_id()       # 句子结尾 (EOS) 的ID
+        self.pad_id: int = self.sp_model.pad_id()       # 填充 (PAD) 的ID
+
+        # 验证分词器词汇表大小是否正确
+        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
+
+    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
+        """
+        将字符串编码为词元ID列表。可以选择是否添加句子开头 (BOS) 和句子结尾 (EOS) 标记。
+
+        参数:
+        s: str, 要编码的字符串。
+        bos: bool, 是否在编码的词元列表前添加 BOS 标记。
+        eos: bool, 是否在编码的词元列表末尾添加 EOS 标记。
+
+        返回:
+        List[int]: 编码后的词元ID列表。
+        """
+        # 确保输入是字符串类型
+        assert type(s) is str
+        # 使用SentencePiece将字符串编码为词元ID
+        t = self.sp_model.encode(s)
+        # 如果需要BOS标记，将其添加到词元列表开头
+        if bos:
+            t = [self.bos_id] + t
+        # 如果需要EOS标记，将其添加到词元列表末尾
+        if eos:
+            t = t + [self.eos_id]
+        return t
+
+    def decode(self, t: List[int]) -> str:
+        """
+        将词元ID列表解码为字符串。
+
+        参数:
+        t: List[int], 词元ID列表。
+
+        返回:
+        str: 解码后的字符串。
+        """
+        return self.sp_model.decode(t)

docs/chapter5/code/train.py

@@ -0,0 +1,257 @@
+import math
+import os
+import time
+from contextlib import nullcontext
+from datetime import datetime
+from functools import partial
+
+import torch
+from model import Transformer, ModelArgs
+from preprocess import Task
+
+# -----------------------------------------------------------------------------
+# I/O 配置，用于定义输出目录和训练时的日志记录与评估设置
+out_dir = "output"  # 模型输出保存路径
+eval_interval = 2000  # 评估间隔步数
+log_interval = 1  # 日志记录间隔步数
+eval_iters = 100  # 每次评估时迭代的步数
+eval_only = False  # 如果为True，脚本在第一次评估后立即退出
+always_save_checkpoint = False  # 如果为True，在每次评估后总是保存检查点
+init_from = "scratch"  # 可以选择从头开始训练（'scratch'）或从已有的检查点恢复（'resume'）
+
+# 数据配置
+batch_size = 8  # 每个微批次的样本数量，如果使用梯度累积，实际批次大小将更大
+max_seq_len = 256  # 最大序列长度
+vocab_size = 4096  # 自定义词汇表大小
+
+# 模型配置
+dim = 288  # 模型的隐藏层维度
+n_layers = 8  # Transformer的层数
+n_heads = 8  # 注意力头的数量
+n_kv_heads = 4  # 模型分组
+multiple_of = 32  # 在某些层的维度必须是该数的倍数
+dropout = 0.0  # Dropout概率
+
+# AdamW优化器配置
+gradient_accumulation_steps = 4  # 梯度累积步数，用于模拟更大的批次
+learning_rate = 5e-4  # 最大学习率
+max_iters = 100000  # 总的训练迭代次数
+weight_decay = 1e-1  # 权重衰减系数
+beta1 = 0.9  # AdamW优化器的β1参数
+beta2 = 0.95  # AdamW优化器的β2参数
+grad_clip = 1.0  # 梯度裁剪阈值，0表示不裁剪
+
+# 学习率衰减配置
+decay_lr = True  # 是否启用学习率衰减
+warmup_iters = 1000  # 学习率预热的步数
+
+# 系统设置
+device = "cuda:0"  # 设备选择：'cpu'，'cuda'，'cuda:0'等
+dtype = "bfloat16"  # 数据类型：'float32'，'bfloat16'，'float16'
+
+# -----------------------------------------------------------------------------
+# 获取配置参数的键值对，便于后续的日志记录
+config_keys = [
+    k
+    for k, v in globals().items()
+    if not k.startswith("_") and isinstance(v, (int, float, bool, str))
+]
+config = {k: globals()[k] for k in config_keys}  # 保存配置到字典中，便于日志记录
+# -----------------------------------------------------------------------------
+
+# 固定一些超参数的默认值
+lr_decay_iters = max_iters  # 学习率衰减步数，设置为等于最大迭代步数
+min_lr = 0.0  # 最小学习率，建议为学习率的十分之一
+vocab_source = 'custom'  # 词汇表来源
+master_process = True  # 用于区分主进程
+seed_offset = 0  # 随机种子偏移量
+ddp_world_size = 1  # 分布式数据并行的世界大小
+tokens_per_iter = batch_size * max_seq_len  # 每次迭代处理的token数
+
+# 设置随机种子，确保可重复性
+torch.manual_seed(1337 + seed_offset)
+torch.backends.cuda.matmul.allow_tf32 = True  # 允许在matmul上使用tf32
+torch.backends.cudnn.allow_tf32 = True  # 允许在cudnn上使用tf32
+device_type = "cuda" if "cuda" in device else "cpu"  # 用于自动选择设备类型
+ptdtype = torch.float16  # 设置训练时使用的数据类型
+
+# 混合精度训练相关
+ctx = (
+    nullcontext()
+    if device_type == "cpu"
+    else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+)
+
+# 为特定任务设置批次迭代器 iter_batches
+iter_batches = partial(
+    Task.iter_batches,  # 调用 Task 类中的 iter_batches 方法
+    batch_size=batch_size,  # 每个批次的样本数量
+    max_seq_len=max_seq_len,  # 每个序列的最大长度
+    vocab_size=vocab_size,  # 词汇表大小
+    vocab_source=vocab_source,  # 词汇表来源（如 llama2 或 custom）
+    device=device,  # 运行模型的设备（如 GPU 或 CPU）
+    num_workers=0,  # 用于数据加载的 worker 数量，0 表示在主线程中加载
+)
+
+# 训练迭代数初始化
+iter_num = 0  # 记录当前迭代数
+
+# 验证集上的最好损失初始值设置为一个极大值，用于后续模型验证时对比更新
+best_val_loss = 1e9  # 设置初始的最佳验证损失为非常大的值，以便在训练中更新
+
+# 模型初始化参数设置
+model_args = dict(
+    dim=dim,  # 模型的隐藏层维度
+    n_layers=n_layers,  # Transformer 的层数
+    n_heads=n_heads,  # 多头注意力机制中的头数
+    n_kv_heads=n_kv_heads,  # 分组数（可能是用于并行化或其他优化目的）
+    vocab_size=vocab_size,  # 词汇表大小
+    multiple_of=multiple_of,  # 用于调整某些维度的参数，确保其为特定数的倍数
+    max_seq_len=max_seq_len,  # 最大序列长度
+    dropout=dropout,  # dropout 概率，用于防止过拟合
+)
+
+# ===========================================================
+# 模型初始化
+gptconf = ModelArgs(**model_args)
+model = Transformer(gptconf)
+
+
+model.to(device)
+
+# 初始化 GradScaler，用于自动混合精度训练（AMP）
+# 如果 enabled=False，表示禁用混合精度，scaler 将不起作用
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == "float16"))
+
+# 优化器初始化，调用模型的 configure_optimizers 方法
+optimizer = model.configure_optimizers(
+    weight_decay,  # 权重衰减（L2 正则化）
+    learning_rate,  # 学习率
+    (beta1, beta2),  # Adam 优化器中的 beta1 和 beta2 参数
+    device_type  # 当前训练设备（如 GPU 或 CPU）
+)
+
+# 定义评估损失的流程
+@torch.no_grad()  # 使用 no_grad 装饰器，确保在评估过程中不计算梯度，从而节省内存
+def estimate_loss():
+    out = {}  # 用于存储训练集和验证集上的平均损失
+    model.eval()  # 将模型设置为评估模式，这会影响 dropout 和 batchnorm 等层的行为
+    for split in ["train", "val"]:  # 分别对训练集和验证集进行评估
+        batch_iter = iter_batches(split=split)  # 获取对应数据集的批次迭代器
+        losses = torch.zeros(eval_iters)  # 初始化一个张量用于存储多次迭代的损失，放在 CPU 上
+        for k in range(eval_iters):  # 进行多次迭代以计算平均损失
+            X, Y = next(batch_iter)  # 从迭代器中获取下一个批次的输入数据 X 和标签 Y
+            with ctx:  # 上下文管理器，可以是 torch.autocast()，用于自动混合精度训练
+                logits = model(X, Y)  # 前向传播，计算模型的输出
+                loss = raw_model.last_loss  # 从模型中获取损失值
+            losses[k] = loss.item()  # 将损失值转换为 Python 标量并存储在 losses 张量中
+        out[split] = losses.mean()  # 计算当前数据集上的平均损失并保存到字典中
+    model.train()  # 恢复模型为训练模式
+    return out  # 返回包含训练集和验证集平均损失的字典
+
+# 定义学习率调度函数
+def get_lr(it):
+    """
+    根据当前的训练迭代步数 it 返回当前的学习率值。
+    学习率调整策略包括线性预热、余弦退火和最小学习率限制。
+    """
+    # 1) 线性预热阶段，在 warmup_iters 之前，学习率线性增加到目标学习率
+    if it < warmup_iters:
+        return learning_rate * it / warmup_iters  # 预热阶段，学习率线性增长
+
+    # 2) 如果迭代步数超过 lr_decay_iters，返回最小学习率 min_lr
+    if it > lr_decay_iters:
+        return min_lr  # 训练进入尾声时，学习率达到最小值并保持不变
+
+    # 3) 余弦退火阶段，在 warmup_iters 和 lr_decay_iters 之间，学习率逐渐降低
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    assert 0 <= decay_ratio <= 1  # 确保衰减比在合法范围内
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # 余弦函数计算衰减系数，范围为0到1
+    return min_lr + coeff * (learning_rate - min_lr)  # 根据衰减系数调整学习率
+
+# 初始化训练数据的迭代器
+train_batch_iter = iter_batches(split="train")
+X, Y = next(train_batch_iter)  # 获取第一个批次的数据
+t0 = time.time()  # 记录开始时间
+local_iter_num = 0  # 本进程中的迭代次数
+raw_model = model  # 如果使用了分布式数据并行 (DDP)，需要解包模型
+running_mfu = -1.0  # 初始化模型浮点运算利用率
+
+os.makedirs(out_dir, exist_ok=True)
+
+while True:
+    # 或许当前step的学习率
+    lr = get_lr(iter_num) if decay_lr else learning_rate
+    # 更新优化器中的学习率
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+    # 在指定的评估间隔进行模型评估和保存检查点
+    if iter_num % eval_interval == 0 and master_process:
+        losses = estimate_loss()  # 评估当前模型在训练集和验证集上的损失
+        print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
+
+        # 如果验证损失降低，或者设置为始终保存检查点，则保存模型
+        if losses["val"] < best_val_loss or always_save_checkpoint:
+            best_val_loss = losses["val"]
+            if iter_num > 0:
+                # 创建检查点字典，包含模型状态、优化器状态和其他信息
+                checkpoint = {
+                    "model": raw_model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "model_args": model_args,
+                    "iter_num": iter_num,
+                    "best_val_loss": best_val_loss,
+                    "config": config,
+                }
+                print(f"saving checkpoint to {out_dir}")
+                # 保存检查点到指定目录
+                torch.save(checkpoint, os.path.join(out_dir, "ckpt.pt"))
+    # 如果只进行评估且已经完成第一次迭代，则退出循环
+    if iter_num == 0 and eval_only:
+        break
+
+    # 前向和反向传播过程，支持梯度累积
+    for micro_step in range(gradient_accumulation_steps):
+
+        with ctx:  # 混合精度训练的上下文管理器
+            logits = model(X, Y)  # 前向传播，计算模型输出
+            loss = raw_model.last_loss  # 获取模型的损失值
+            loss = loss / gradient_accumulation_steps  # 平均损失以支持梯度累积
+
+        X, Y = next(train_batch_iter)  # 获取下一个批次的数据
+        # 反向传播，计算梯度
+        scaler.scale(loss).backward()
+    # 梯度处理阶段
+    if grad_clip != 0.0:
+        # 取消梯度缩放以进行梯度裁剪
+        scaler.unscale_(optimizer)
+        # 对梯度进行裁剪，防止梯度爆炸
+        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+    # 更新优化器和梯度缩放器（用于混合精度训练）
+    scaler.step(optimizer)
+    scaler.update()
+    # 清空优化器的梯度，释放显存
+    optimizer.zero_grad(set_to_none=True)
+
+    # 计时和日志记录
+    t1 = time.time()
+    dt = t1 - t0  # 计算一次迭代所需时间
+    t0 = t1
+    if iter_num % log_interval == 0 and master_process:
+        # 获取当前损失值，并根据梯度累积步骤进行调整
+        lossf = loss.item() * gradient_accumulation_steps
+        if local_iter_num >= 5:  # 让训练循环先运行几个迭代再计算模型利用率
+            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
+            # 使用滑动平均更新模型浮点运算利用率（MFU）
+            running_mfu = mfu if running_mfu == -1.0 else 0.9 * running_mfu + 0.1 * mfu
+        print(
+            f"{iter_num} | loss {lossf:.4f} | lr {lr:e} | {dt*1000:.2f}ms | mfu {running_mfu*100:.2f}%"
+            # mfu 表示模型浮点运算利用率
+        )
+    iter_num += 1  # 全局迭代次数自增
+    local_iter_num += 1  # 本地迭代次数自增
+
+    # 终止条件，达到最大迭代次数则退出循环
+    if iter_num > max_iters:
+        break

docs/chapter5/code/train_vocab.py

@@ -0,0 +1,147 @@
+import glob
+import json
+import os
+from tqdm import tqdm
+import requests
+import sentencepiece as spm
+import argparse
+
+DATA_CACHE_DIR = 'data'
+
+def download_file(url: str, fname: str, chunk_size=1024):
+    """发送HTTP GET请求以流式方式获取文件"""
+    resp = requests.get(url, stream=True)
+    
+    # 获取文件的总大小（以字节为单位），默认为0如果没有提供'content-length'头信息
+    total = int(resp.headers.get("content-length", 0))
+    
+    # 以写二进制模式打开一个文件以保存下载的内容
+    with open(fname, "wb") as file, tqdm(
+        desc=fname,           # 进度条前面的描述信息（通常是文件名）
+        total=total,          # 总的字节数，用于设置进度条的总长度
+        unit="iB",            # 进度条的单位，'iB'代表二进制字节
+        unit_scale=True,      # 启用单位缩放，如KB、MB等
+        unit_divisor=1024,    # 设置单位换算的除数，这里为1024
+    ) as bar:
+        # 逐块读取响应内容并写入文件
+        for data in resp.iter_content(chunk_size=chunk_size):
+            size = file.write(data)  # 写入数据块到文件
+            bar.update(size)         # 更新进度条
+
+def download():
+    """在DATA_CACHE_DIR中创建目录，如果目录不存在则创建"""
+    os.makedirs(DATA_CACHE_DIR, exist_ok=True)
+
+    # 定义TinyStories数据集的下载URL和保存的文件名
+    data_url = "https://www.modelscope.cn/datasets/AI-ModelScope/TinyStories/resolve/master/TinyStories_all_data.tar.gz"
+    data_filename = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data.tar.gz")
+    
+    # 检查数据集是否已经下载，如果没有下载则进行下载
+    if not os.path.exists(data_filename):
+        print(f"Downloading {data_url} to {data_filename}...")
+        download_file(data_url, data_filename)  # 使用之前定义的download_file函数进行下载
+    else:
+        print(f"{data_filename} already exists, skipping download...")
+
+    # 定义解压缩后的数据目录
+    data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
+    
+    # 检查数据目录是否存在，如果不存在则解压缩数据集
+    if not os.path.exists(data_dir):
+        os.makedirs(data_dir, exist_ok=True)  # 创建数据目录
+        print(f"Unpacking {data_filename}...")
+        os.system(f"tar -xzf {data_filename} -C {data_dir}")  # 使用系统命令解压缩.tar.gz文件
+    else:
+        print(f"{data_dir} already exists, skipping unpacking...")
+
+    # 查找解压后的所有JSON文件，排序后获取文件名列表
+    shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
+    
+    # 打开第一个JSON文件并读取内容
+    with open(shard_filenames[0], "r") as f:
+        data = json.load(f)  # 将JSON文件内容加载到变量data中
+    
+    print("Download done.")  # 下载完成信息
+    print(f"Number of shards: {len(shard_filenames)}")  # 打印解压后数据分片的数量
+    print(f"Example story:\n{data[0]}")  # 打印第一个分片中的一个示例故事
+
+def load_text_from_files(path):
+    path_list = glob.glob(path)
+    text_data = []
+    for file_path in path_list:
+        with open(file_path, 'r', encoding='utf-8') as file:
+            text_data.extend(file.readlines())
+    return text_data
+
+def batch_iterator(text_data, batch_size=648):
+    for i in range(0, len(text_data), batch_size):
+        yield text_data[i:i + batch_size]
+
+def train_vocab(vocab_size: int=32000, num_shards: int=20):
+    """
+    vocab_size: int, 词汇表的大小，决定分词器的词汇量。
+    num_shards: int, 用于加快词汇表训练的效率，指定要处理的分片数量。
+    """
+    # 确保词汇表大小为正数
+    assert vocab_size > 0, "Vocab size must be positive"
+
+    # SentencePiece 模型的前缀路径，将用于保存分词器
+    prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")
+
+    # 1) 将多个分片中的文本导出为单个文本文件 tiny.txt
+    tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt")
+    data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
+    shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
+
+    # 创建 tiny.txt 文件并写入指定数量的分片中的文本
+    print(f"Writing temporary file {tiny_file} with {num_shards} shards...")
+    with open(tiny_file, "w", encoding="utf-8") as of:
+        # 遍历前 num_shards 个分片
+        for shard in tqdm(shard_filenames[:num_shards]):
+            with open(shard, "r") as f:
+                data = json.load(f)  # 读取分片中的JSON数据
+            # 遍历每个例子，将其中的故事文本写入 tiny.txt 文件
+            for example in data:
+                text = example["story"]
+                text = text.strip()  # 去除文本首尾的空白字符
+                of.write(text + "\n")  # 每个文本写入一行
+
+    # 输出生成的 tiny.txt 文件的大小
+    print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB")
+
+    # 2) 使用 SentencePiece 训练分词器
+    print("Will now train the vocab...")
+    spm.SentencePieceTrainer.train(
+        input=tiny_file,         # 输入文件为之前生成的 tiny.txt
+        model_prefix=prefix,     # 模型前缀路径
+        model_type="bpe",        # 使用 Byte-Pair Encoding (BPE) 训练分词器
+        vocab_size=vocab_size,   # 词汇表大小
+        self_test_sample_size=0, # 自测样本大小设置为 0
+        input_format="text",     # 输入文件格式为纯文本
+        character_coverage=1.0,  # 覆盖所有字符（包括非常见字符）
+        num_threads=os.cpu_count(),  # 使用 CPU 的线程数
+        split_digits=True,       # 拆分数字
+        allow_whitespace_only_pieces=True,  # 允许仅由空格组成的词元
+        byte_fallback=True,      # 启用字节级回退
+        unk_surface=r" \342\201\207 ",  # UNK token 表示未知字符的方式
+        normalization_rule_name="identity"  # 使用“identity”归一化规则
+    )
+
+    # 3) 可选的清理操作，询问用户是否删除临时文件 tiny.txt
+    dec = input(f"Delete the temporary file {tiny_file}? [y/N] ")
+    if dec.lower() == "y":
+        os.remove(tiny_file)  # 删除临时文件
+        print(f"Deleted {tiny_file}")
+
+    # 输出模型保存的路径
+    print(f"Trained tokenizer is in {prefix}.model")
+    print("Done.")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--download", type=bool, default=True, help="download the dataset")
+    parser.add_argument("--vocab_size", type=int, default=4096, help="vocab size")
+    args = parser.parse_args()
+    if args.download:
+        download()
+    train_vocab(args.vocab_size)