Tiny-universe-taks2-从零预训练一个tiny-llama

本次笔记内容全部参考datawhale官方给出的tiny-llm教程欢迎感兴趣的小伙伴去给个datawhale贡献者们一个star，本篇笔记主要记录手动训练一个简单的大语言模型tiny-llm全流程心得。

Tiny-llama

训练一个简单的大语言模型主要有以下步骤：

1. 训练Tokenizer：这里使用 TinyStory来进行训练 ，它是一个由GPT-3.5和GPT-4生成的小型故事数据集，符合我们探索的需求。
2. 数据预处理：将文本序列转化成模型可以理解的数字序列
3. 训练模型：训练一个与Llama一致的decoder only Transformer模型
4. 使用模型生成文本：使用训练好的模型推理，验证模型可行性

Step1 训练Tokenizer

LLaMA2 的词表大小为 32000，但由于 TinyStory 数据集较小，词汇量有限，在这里将词表大小设置为 4,096。训练完成后，我们得到的 Tokenizer 能够将文本转换为数字序列，也可以将数字序列还原为文本。
这里给出下载数据并训练tokenizer的脚本，datawhale助教给出的注释基本已经很清楚了，我自己也在里面做了一些补充。
运行python train_vocab.py --download True --vocab_size 4096即可开始训练。

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请求以流式方式获取文件，这里提供了一个较为通用的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)

当data文件夹下存在以下两个文件表示tokenizer训练成功

Step2 数据预处理

使用 python preprocess.py对数据进行处理，这步主要是将训练数据编码成模型能够理解数字序列。

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)

训练模型

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

使用模型生成文本

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}")  # 打印生成的样本并用分隔线分割