noao-chat-2-mid阶段训练

# 训练数据组成（总共 ~850K 条对话）
train_dataset = TaskMixture([
    SmolTalk(split="train"),              # 460K 通用对话
    MMLU(subset="auxiliary_train"),        # 100K 知识问答
    GSM8K(subset="main", split="train"),   # 8K 数学推理
    CustomJSON(identity_file),             # 1K 身份对话 × 2 轮
    SimpleSpelling(size=200000),           # 200K 拼写
    SpellingBee(size=80000),               # 80K 字母计数
])

def tokenizing_distributed_data_loader(B, T, split):
    # 读取 Parquet 文件
    pf = pq.ParquetFile(filepath)
    batch = rg.column('text').to_pylist()
    
    # Tokenize
    token_lists = tokenizer.encode(batch, prepend=bos_token)
    
    # 拼接成连续流
    for tokens in token_lists:
        token_buffer.extend(tokens)
    
    # 取出固定长度
    needed = B * T + 1
    tokens = [token_buffer.popleft() for _ in range(needed)]
    
    # 切分
    inputs = tokens[:-1].view(B, T)
    targets = tokens[1:].view(B, T)

def mid_data_generator(split):
    dataset = train_dataset  # TaskMixture
    
    # 遍历数据集
    for cursor in range(ddp_rank, len(dataset), ddp_world_size):
        # 获取一个对话
        conversation = dataset[cursor]
        
        # Tokenize（保留结构）
        ids, mask = tokenizer.render_conversation(conversation)
        # ids: [<|bos|>, <|user_start|>, ..., <|assistant_end|>]
        
        # 累积到缓冲区
        token_buffer.extend(ids)
        
        # 当缓冲区足够时，切分批次
        if len(token_buffer) >= needed_tokens:
            tokens = [token_buffer.popleft() for _ in range(needed_tokens)]
            inputs = tokens[:-1].view(B, T)
            targets = tokens[1:].view(B, T)
            yield inputs, targets

# 错误做法：依次训练每个任务
for task in tasks:
    for example in task:
        train(example)

# 结果：前期都在训练 SmolTalk，后期才见到 GSM8K
# 模型会遗忘早期学到的东西！

class TaskMixture:
    def __init__(self, tasks):
        # 1. 构建索引映射
        self.index_map = []
        for task_idx, task in enumerate(tasks):
            for local_idx in range(len(task)):
                self.index_map.append((task_idx, local_idx))
        
        # 2. 打乱索引（确定性）
        rng = random.Random(42)
        rng.shuffle(self.index_map)
        
        # 现在访问顺序是：
        # [SmolTalk_123, MMLU_45, GSM8K_7, SmolTalk_456, ...]

# Base Training: 明确的迭代次数
num_iterations = 10000  # 确定
for step in range(num_iterations):
    ...

# Mid Training: 基于数据集大小
dataset_size = 850000  # 对话数量
# 每个 step 消耗多少对话？不确定！（因为对话长度不一）
# 所以用进度百分比

# 全局变量
last_step = False       # 是否到达最后一步
approx_progress = 0.0   # 0 → 1

# 在数据生成器中更新
def mid_data_generator(split):
    cursor = ddp_rank
    it = 0
    
    while True:
        # 处理一个对话
        conversation = dataset[cursor]
        cursor += ddp_world_size
        
        # 更新进度
        if cursor >= dataset_size:
            cursor -= dataset_size  # 回绕（下一个 epoch）
            if split == "train":
                last_step = True  # 标记为最后一步
        
        # 计算近似进度
        approx_progress = cursor / dataset_size
        
        yield inputs, targets

# 问题：不同 GPU 可能处理速度不同
# GPU 0: last_step = True
# GPU 1: last_step = False  ← 还没完成

# 解决：同步 last_step
if ddp:
    last_step_tensor = torch.tensor(last_step, device=device)
    dist.all_reduce(last_step_tensor, op=dist.ReduceOp.MAX)
    # 只要有一个 GPU 到达终点，所有 GPU 都停止
    last_step = bool(last_step_tensor.item())

# Base Training: 基于步数
def get_lr_multiplier(step):
    if step < warmup_iters:
        return step / warmup_iters
    elif step <= num_iterations - warmdown_iters:
        return 1.0
    else:
        return linear_decay(...)

# Mid Training: 基于进度百分比
def get_lr_multiplier(progress):
    # progress: 0.0 → 1.0
    if progress < 0.8:
        return 1.0  # 前 80% 保持不变
    else:
        # 后 20% 线性衰减到 0
        return 1 - (progress - 0.8) / 0.2

# 预取第一批数据
x, y = next(train_loader)

while True:
    # === 评估阶段 ===
    if step % eval_every == 0:
        model.eval()
        val_bpb = evaluate_bpb(...)
        model.train()
    
    # === 保存检查点 ===
    if last_step:
        save_checkpoint(...)
        break
    
    # === 训练阶段 ===
    synchronize()
    t0 = time.time()
    
    # 梯度累积
    for micro_step in range(grad_accum_steps):
        with autocast_ctx:
            loss = model(x, y)
        
        loss = loss / grad_accum_steps
        loss.backward()
        
        # 异步预取下一批
        x, y = next(train_loader)
        progress = max(progress, approx_progress)
    
    # 更新学习率
    lrm = get_lr_multiplier(progress)
    for opt in optimizers:
        for group in opt.param_groups:
            group["lr"] = group["initial_lr"] * lrm
    
    # 更新参数
    for opt in optimizers:
        opt.step()
    model.zero_grad()
    
    synchronize()
    t1 = time.time()
    
    # === 日志记录 ===
    step += 1
    dt = t1 - t0
    print(f"step {step} ({progress*100:.2f}%) | loss: {loss:.6f} ...")

# 输入：结构化对话
conversation = {
    "messages": [
        {"role": "user", "content": "What is 2+2?"},
        {"role": "assistant", "content": "4"}
    ]
}

# 输出：token IDs + mask
ids, mask = tokenizer.render_conversation(conversation)

# ids: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
# mask: [0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0]

# 展开看：
# token_id | token_str           | mask | 说明
# ---------|---------------------|------|------
# 1        | <|bos|>             | 0    | 文档开始
# 2        | <|user_start|>      | 0    | 用户开始
# 3-7      | "What is 2+2?"      | 0    | 用户消息（不训练）
# 8        | <|user_end|>        | 0    | 用户结束
# 9        | <|assistant_start|> | 0    | 助手开始
# 10-11    | "4"                 | 1    | 助手回复（训练！）
# 12       | <|assistant_end|>   | 1    | 助手结束（训练）

# 多个对话拼接
conversation_1: 200 tokens
conversation_2: 150 tokens
conversation_3: 180 tokens
...

# 累积到缓冲区
token_buffer: [tok1, tok2, ..., tok_N]  # N 很大

# 切分批次
needed = 32 * 2048 + 1 = 65537

tokens = token_buffer[:65537]
# [tok_1, tok_2, ..., tok_65537]

# 重塑
inputs = tokens[:-1].view(32, 2048)   # (32, 2048)
targets = tokens[1:].view(32, 2048)   # (32, 2048)

# 输入
inputs: (32, 2048)  # int32

# Embedding
x = model.transformer.wte(inputs)
# (32, 2048) → (32, 2048, 768)

# Transformer Blocks
for block in model.transformer.h:
    x = block(x, ...)
# (32, 2048, 768) → (32, 2048, 768)

# LM Head
logits = model.lm_head(x)
# (32, 2048, 768) @ (768, 32000) → (32, 2048, 32000)

# 损失计算
loss = F.cross_entropy(
    logits.view(-1, 32000),   # (65536, 32000)
    targets.view(-1)          # (65536,)
)
# loss: 标量

# ===== Base Training =====
# 从 Parquet 读取
pf = pq.ParquetFile(filepath)
batch = rg.column('text').to_pylist()

# 简单 tokenize
token_lists = tokenizer.encode(batch, prepend=bos_token)

# 无结构，直接拼接
for tokens in token_lists:
    token_buffer.extend(tokens)

# ===== Mid Training =====
# 从 Task 对象读取
conversation = dataset[cursor]

# 结构化 tokenize
ids, mask = tokenizer.render_conversation(conversation)
# 包含 <|user_start|>, <|assistant_start|> 等特殊 token

# 保留对话结构
token_buffer.extend(ids)

# ===== Base Training =====
# 从零开始
with torch.device("meta"):
    model_config = GPTConfig(...)
    model = GPT(model_config)
model.to_empty(device=device)
model.init_weights()  # 随机初始化

# ===== Mid Training =====
# 从 Base 模型加载
model, tokenizer, meta = load_model(
    "base",  # 加载 base 模型
    device,
    phase="train",
    model_tag=model_tag,
    step=step
)
# 参数已经训练过，继续训练

# ===== Base Training =====
for step in range(num_iterations):  # 固定次数
    loss = model(x, y)
    loss.backward()
    optimizer.step()

# ===== Mid Training =====
while True:  # 直到遍历完数据集
    loss = model(x, y)
    loss.backward()
    optimizer.step()
    
    if last_step:  # 数据集遍历完
        break

# ===== Base Training =====
base_config = {
    "total_batch_size": 524288,      # 512K tokens/step
    "device_batch_size": 32,
    "num_iterations": 50000,         # 明确指定
    "warmup_ratio": 0.0,             # 无 warmup
    "warmdown_ratio": 0.2,           # 最后 20% 衰减
    "embedding_lr": 0.2,
    "matrix_lr": 0.02,
    "unembedding_lr": 0.004,
    "init_lr_frac": 1.0,             # 从 100% 开始
}

# ===== Mid Training =====
mid_config = {
    "total_batch_size": 524288,      # 512K tokens/step (相同)
    "device_batch_size": 32,
    "num_iterations": -1,            # 自动推断
    "warmup_ratio": 0.0,             # 无 warmup
    "warmdown_ratio": 0.2,           # 最后 20% 衰减
    "embedding_lr": 0.2,
    "matrix_lr": 0.02,
    "unembedding_lr": 0.004,
    "init_lr_frac": 1.0,             # 从 100% 开始
}

# 假设 token_buffer 中有以下对话（简化）：

# 对话 1: SmolTalk（200 tokens）
[<|bos|>, <|user_start|>, "How", "are", "you", "?", <|user_end|>,
 <|assistant_start|>, "I'm", "doing", "well", "!", <|assistant_end|>, ...]

# 对话 2: MMLU（50 tokens）
[<|bos|>, <|user_start|>, "What", "is", "2+2", "?", "\n", "A.", "2", ...,
 <|user_end|>, <|assistant_start|>, "D", <|assistant_end|>]

# 对话 3: GSM8K（300 tokens）
[<|bos|>, <|user_start|>, "Weng", "earns", ..., <|user_end|>,
 <|assistant_start|>, "First", ",", "let's", ..., <|python_start|>, "12/60", ...]

# 对话 4: SpellingBee（100 tokens）
[<|bos|>, <|user_start|>, "How", "many", "r", "in", "strawberry", <|user_end|>,
 <|assistant_start|>, <|python_start|>, "'strawberry'.count('r')", ...]

# 从 buffer 中取出 4*512+1 = 2049 个 token
# 这些 token 来自上述对话的混合

tokens = token_buffer[:2049]

inputs = tokens[:-1]   # 前 2048 个
targets = tokens[1:]   # 后 2048 个

# 重塑
inputs = inputs.view(4, 512)
targets = targets.view(4, 512)

# 现在每个样本包含多个对话的片段
# 样本 0: 对话1的一部分 + 对话2的开头
# 样本 1: 对话2的剩余 + 对话3的一部分
# 样本 2: 对话3的中间部分
# 样本 3: 对话3的结尾 + 对话4

# 前向传播
logits = model(inputs)  # (4, 512, 32000)

# 计算损失
loss = F.cross_entropy(
    logits.view(-1, 32000),
    targets.view(-1)
)

# 反向传播
loss.backward()

# 更新参数
optimizer.step()

prompts = [
    # 1. 对话能力
    "Hello! How are you?",
    
    # 2. 知识问答
    "What is the capital of France?",
    
    # 3. 数学推理
    "If I have 5 apples and buy 3 more, how many do I have?",
    
    # 4. 拼写能力
    "How many 'r' are in 'strawberry'?",
]

# 推荐配比
train_dataset = TaskMixture([
    SmolTalk(...),          # 50-60% 通用对话
    Knowledge_QA(...),      # 15-20% 知识问答
    Math_Reasoning(...),    # 1-5% 数学推理
    Tool_Use(...),          # 1-5% 工具使用
    Spelling(...),          # 20-30% 拼写/基础能力
    Identity(...),          # <1% 身份设定（过采样）
])

ids, mask = tokenizer.render_conversation(conversation)
# ids: [token_id1, token_id2, ...]  完整的 token 序列
# mask: [0, 0, 1, 1, 0, ...]        标记哪些 token 需要训练

def render_conversation(self, conversation):
    ids, mask = [], []
    
    def add_tokens(token_ids, mask_val):
        ids.extend(token_ids)
        mask.extend([mask_val] * len(token_ids))
    
    # BOS token: mask=0
    add_tokens(bos, 0)
    
    for message in messages:
        if message["role"] == "user":
            # User 的所有内容：mask=0
            add_tokens(user_start, 0)
            add_tokens(content_ids, 0)  # ← User 输入不训练
            add_tokens(user_end, 0)
        
        elif message["role"] == "assistant":
            add_tokens(assistant_start, 0)  # Start token 不训练
            add_tokens(content_ids, 1)      # ← Assistant 回复训练！
            add_tokens(assistant_end, 1)    # End token 也训练
    
    return ids, mask

# Mid Training (mid_train.py 第 133 行)
ids, _ = tokenizer.render_conversation(conversation)
#      ↑ 直接丢弃 mask！

# SFT Training (chat_sft.py 第 123 行)
ids, mask = tokenizer.render_conversation(doc)
#      ↑ 保留并使用 mask

# Mid Training 的数据处理
conversation = dataset[cursor]
ids, _ = tokenizer.render_conversation(conversation)
token_buffer.extend(ids)  # 直接拼接所有 token

# 生成 inputs, targets（没有 mask）
inputs = tokens[:-1]   # (B, T)
targets = tokens[1:]   # (B, T)

# 计算 loss（所有 token 都参与）
loss = model(inputs, targets)
# CrossEntropyLoss 会在所有 token 上计算梯度

# SFT Training 的数据处理
ids, mask = tokenizer.render_conversation(doc)

# 关键步骤：将 mask=0 的位置设为 -1
targets = ids[1:]
mask_tensor = mask[1:]
targets[mask_tensor == 0] = -1  # ← 设为 ignore_index

# 计算 loss（只在 mask=1 的 token 上）
loss = F.cross_entropy(
    logits.view(-1, vocab_size),
    targets.view(-1),
    ignore_index=-1  # ← targets=-1 的位置不计算 loss
)

# 训练目标：预测下一个 token
输入：<|user_start|>What is 2+
目标：2

输入：<|user_start|>What is 2+2
目标：?

输入：<|user_start|>What is 2+2?
目标：<|user_end|>

TaskMixture([
    SmolTalk,        # 通用对话
    MMLU,            # 多选题
    GSM8K,           # 数学推理
    SpellingBee,     # 字符级任务
])

# 问题：如果训练所有 token
输入：<|user_start|>What is the capital of France?
目标：<|user_end|>

# 风险：模型可能记住了这个问题
# 测试时看到类似问题，直接输出训练数据的答案

# Step 1: Tokenize（返回 mask 但丢弃）
conversation = {
    "messages": [
        {"role": "user", "content": "What is 2+2?"},
        {"role": "assistant", "content": "4"}
    ]
}

ids, _ = tokenizer.render_conversation(conversation)
# ids = [BOS, USER_START, W, h, a, t, ..., USER_END, 
#        ASST_START, 4, ASST_END]

# Step 2: 拼接到 token buffer
token_buffer.extend(ids)

# Step 3: 生成 batch（没有 mask）
inputs = [USER_START, W, h, a, t, ...]  # (B, T)
targets = [W, h, a, t, ..., USER_END, ASST_START, 4, ASST_END]

# Step 4: 计算 loss（所有 token）
logits = model(inputs)  # (B, T, V)
loss = F.cross_entropy(
    logits.view(-1, V),
    targets.view(-1)
)
# 每个 token 都有梯度：
# - W, h, a, t (user 的输入)
# - 4 (assistant 的输出)
# - 所有特殊 token

# Step 1: Tokenize（保留 mask）
ids, mask = tokenizer.render_conversation(conversation)
# ids = [BOS, USER_START, W, h, a, t, ..., USER_END, 
#        ASST_START, 4, ASST_END]
# mask = [0, 0, 0, 0, 0, 0, ..., 0, 
#         0, 1, 1]
#                ↑  ↑  只有这两个是 1

# Step 2: 应用 mask（设置 ignore_index）
targets = ids[1:]
mask_tensor = mask[1:]
targets[mask_tensor == 0] = -1
# targets = [-1, -1, -1, -1, -1, ..., -1, 
#            -1, 4, ASST_END]
#                ↑   ↑  只有这两个保留

# Step 3: 计算 loss（只在 mask=1 的 token 上）
loss = F.cross_entropy(
    logits.view(-1, V),
    targets.view(-1),
    ignore_index=-1  # ← 关键！
)
# 只有 targets != -1 的位置有梯度

# Mask 的生成：nanochat/tokenizer.py 第 258-358 行
def render_conversation(self, conversation):
    """返回 ids 和 mask"""
    ids, mask = [], []
    # ... 详细的 mask 生成逻辑

# Mid Training 丢弃 mask：scripts/mid_train.py 第 133 行
ids, _ = tokenizer.render_conversation(conversation)
#      ↑ 下划线表示丢弃返回值

# SFT Training 使用 mask：scripts/chat_sft.py 第 123 行
ids, mask = tokenizer.render_conversation(doc)
# 后续在 collate_and_yield 中将 mask=0 的位置设为 -1