noao-chat-3-sft阶段训练

# 训练数据混合 (总共 ~23K 对话)
train_ds = TaskMixture([
    ARC(subset="ARC-Easy", split="train"),      # 2.3K 简单推理
    ARC(subset="ARC-Challenge", split="train"), # 1.1K 困难推理
    GSM8K(subset="main", split="train"),        # 8K 数学题
    SmolTalk(split="train", stop=10_000),       # 10K 对话
    CustomJSON(identity_file),                  # 1K 身份对话
    SimpleSpelling(size=300, split="train"),    # 300 拼写
    SpellingBee(size=300, split="train"),       # 300 字母计数
])

# 所有 token 都计算损失
tokens = [<|bos|>, <|user_start|>, "What", "is", "the", "capital", 
          "of", "France", "?", <|user_end|>, <|assistant_start|>, 
          "The", "capital", "of", "France", "is", "Paris", ".", 
          <|assistant_end|>]

# 模型需要学习：
# - 预测 <|user_start|> 之后是 "What"  ← 为什么要学这个？
# - 预测 "What" 之后是 "is"            ← 用户怎么说话不重要！
# - 预测 "?" 之后是 <|user_end|>        ← 这是固定格式

tokens = [<|bos|>, <|user_start|>, "What", "is", ..., <|user_end|>, 
          <|assistant_start|>, "The", "capital", ..., <|assistant_end|>]

mask   = [0,       0,              0,      0,   ..., 0,
          0,                       1,      1,        ..., 1]
          ↑                        ↑
          不训练                   训练！

def render_conversation(self, conversation):
    ids, mask = [], []
    
    def add_tokens(token_ids, mask_val):
        """辅助函数：添加 token 和对应的 mask"""
        ids.extend(token_ids)
        mask.extend([mask_val] * len(token_ids))
    
    # 开始标记
    add_tokens(bos, 0)  # <|bos|> 不训练
    
    for message in messages:
        if message["role"] == "user":
            # 用户消息：全部不训练
            add_tokens(user_start, 0)
            add_tokens(user_message_ids, 0)  # ← mask=0
            add_tokens(user_end, 0)
            
        elif message["role"] == "assistant":
            # 助手消息：开始标记不训练，内容训练
            add_tokens(assistant_start, 0)  # ← 不训练开始标记
            
            if isinstance(content, str):
                # 纯文本回复
                add_tokens(content_ids, 1)  # ← mask=1，训练！
                
            elif isinstance(content, list):
                # 包含工具调用的回复
                for part in content:
                    if part["type"] == "text":
                        add_tokens(part_ids, 1)  # ← 文本部分训练
                    elif part["type"] == "python":
                        add_tokens(python_start, 1)
                        add_tokens(code_ids, 1)  # ← 代码部分训练
                        add_tokens(python_end, 1)
                    elif part["type"] == "python_output":
                        add_tokens(output_start, 0)
                        add_tokens(output_ids, 0)  # ← 输出不训练
                        add_tokens(output_end, 0)
            
            add_tokens(assistant_end, 1)  # ← 结束标记训练（学会停止）
    
    return ids, mask

def collate_and_yield(batch):
    nrows = len(batch)
    ncols = max(len(ids) for ids, mask in batch) - 1
    
    # 初始化
    inputs = torch.full((nrows, ncols), pad_token_id)
    targets = torch.full((nrows, ncols), -1)  # ← -1 是关键！
    
    for i, (ids, mask) in enumerate(batch):
        n = len(ids)
        ids_tensor = torch.tensor(ids)
        
        # 输入和目标
        inputs[i, :n-1] = ids_tensor[:-1]
        targets[i, :n-1] = ids_tensor[1:]
        
        # 应用 mask：mask=0 的位置设为 -1
        mask_tensor = torch.tensor(mask[1:])
        targets[i, :n-1][mask_tensor == 0] = -1  # ← 屏蔽
    
    return inputs, targets

# 在模型的 forward 方法中
def forward(self, inputs, targets):
    logits = self.lm_head(x)  # (B, T, vocab_size)
    
    # 交叉熵损失，自动忽略 target=-1 的位置
    loss = F.cross_entropy(
        logits.view(-1, vocab_size),
        targets.view(-1),
        ignore_index=-1  # ← 忽略 -1
    )
    return loss

# inputs (2, 8)
[[101, 202, 303, 404, 505, 606, 707, 808],
 [909, 1010, 1111, 1212, 1313, 1414, 1515, 1616]]

# targets (2, 8) - 注意 -1 的位置
[[202, 303, 404, -1,  -1,  606, 707, 808],
 [1010, -1,  -1,  -1, 1313, 1414, 1515, -1]]
  ↑                   ↑
  训练               不训练

# 计算损失时：
# - 位置 [0,0]: 预测 202，计算损失 ✅
# - 位置 [0,3]: 目标是 -1，跳过 ❌
# - 位置 [1,1]: 目标是 -1，跳过 ❌

# SFT 超参数
sft_config = {
    # 模型来源
    "source": "mid",  # 从 Mid 模型开始
    "model_tag": None,
    "step": None,
    
    # 训练规模
    "device_batch_size": 4,           # 每卡批次（小）
    "target_examples_per_step": 32,   # 目标批次大小
    "num_epochs": 1,                  # 训练轮数
    "num_iterations": -1,             # 自动计算
    
    # 学习率（层级化）
    "unembedding_lr": 0.004,  # 输出层（最小）
    "embedding_lr": 0.2,       # 输入层（最大）
    "matrix_lr": 0.02,         # 中间层
    "weight_decay": 0.0,
    "init_lr_frac": 0.02,      # 从 2% 开始！
    
    # 评估
    "eval_every": 100,
    "eval_steps": 100,
    "eval_metrics_every": 200,
}

# 实际学习率
actual_lr = base_lr * init_lr_frac

# 例如：
embedding_lr = 0.2 * 0.02 = 0.004   # 实际从 0.004 开始
matrix_lr = 0.02 * 0.02 = 0.0004    # 实际从 0.0004 开始

def get_lr_multiplier(it):
    # 简单的线性衰减
    lrm = 1.0 - it / num_iterations
    return lrm

# 预计算迭代次数
num_iterations = (len(train_ds) // target_examples_per_step) * num_epochs
# 例如：(23000 // 32) * 1 ≈ 718 步

for step in range(num_iterations):
    last_step = (step == num_iterations)
    
    # ===== 评估验证损失 =====
    if last_step or step % eval_every == 0:
        model.eval()
        val_loader = build_val_loader()
        losses = []
        for _ in range(eval_steps):
            val_inputs, val_targets = next(val_loader)
            with torch.no_grad():
                loss = model(val_inputs, val_targets)
            losses.append(loss)
        val_loss = torch.stack(losses).mean()
        print(f"Step {step} | Val loss: {val_loss:.6f}")
        model.train()
    
    # ===== 评估任务指标 =====
    if last_step or (step > 0 and step % eval_metrics_every == 0):
        model.eval()
        # MMLU 准确率
        mmlu_acc = run_chat_eval("MMLU", model, ...)
        # ARC-Easy 准确率
        arc_acc = run_chat_eval("ARC-Easy", model, ...)
        print(f"Step {step} | MMLU: {mmlu_acc:.4f}, ARC: {arc_acc:.4f}")
        model.train()
    
    if last_step:
        break
    
    # ===== 计算梯度 =====
    num_tokens = 0
    for micro_step in range(grad_accum_steps):
        train_inputs, train_targets = next(train_loader)
        
        with autocast_ctx:
            loss = model(train_inputs, train_targets)
        
        loss = loss / grad_accum_steps
        loss.backward()
        
        # 统计有效 token 数量
        num_tokens += (train_targets >= 0).sum()
    
    # ===== 更新学习率 =====
    lrm = get_lr_multiplier(step)
    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()
    
    # ===== 日志 =====
    print(f"Step {step}/{num_iterations} | loss: {loss:.6f} | "
          f"lrm: {lrm:.4f} | num_tokens: {num_tokens}")

# ===== Micro-step 1 =====
train_inputs, train_targets = next(train_loader)
# inputs:  (4, 2048) - int32
# targets: (4, 2048) - int64，包含 -1

# 前向传播
x = model.transformer.wte(train_inputs)  # (4, 2048, 768)
# ... Transformer blocks ...
logits = model.lm_head(x)  # (4, 2048, 32000)

# 损失计算
loss = F.cross_entropy(
    logits.view(-1, 32000),   # (8192, 32000)
    train_targets.view(-1),   # (8192,)
    ignore_index=-1
)
# loss: 标量，例如 2.345

# 归一化并反向传播
loss = loss / 8  # 2.345 / 8 = 0.293
loss.backward()  # 累积梯度

# ===== Micro-step 2-8 =====
# 重复上述过程，梯度持续累积

# ===== 更新参数 =====
for opt in optimizers:
    opt.step()  # 使用累积的梯度更新
model.zero_grad()  # 清零，准备下一步

def evaluate_validation_loss():
    model.eval()
    val_loader = build_val_loader()
    losses = []
    
    for _ in range(eval_steps):  # 例如 100 步
        val_inputs, val_targets = next(val_loader)
        with torch.no_grad():
            loss = model(val_inputs, val_targets)
        losses.append(loss)
    
    val_loss = torch.stack(losses).mean()
    return val_loss

val_ds = SmolTalk(split="test")  # 24K 测试对话

def run_chat_eval(task_name, model, tokenizer, engine, 
                  batch_size=8, max_problems=1024):
    """
    评估模型在选择题任务上的表现
    """
    if task_name == "MMLU":
        task = MMLU(subset="all", split="test")
    elif task_name == "ARC-Easy":
        task = ARC(subset="ARC-Easy", split="test")
    
    correct = 0
    total = 0
    
    for idx in range(min(max_problems, len(task))):
        # 获取问题
        conversation = task[idx]
        # conversation["letters"] = ["A", "B", "C", "D"]
        
        # 生成回答
        tokens = tokenizer.render_for_completion(conversation)
        
        with torch.no_grad():
            # 限制输出只能是 A/B/C/D
            assistant_response = engine.generate_restricted(
                tokens,
                allowed_tokens=conversation["letters"]
            )
        
        # 评估正确性
        is_correct = task.evaluate(conversation, assistant_response)
        correct += is_correct
        total += 1
    
    accuracy = correct / total
    return accuracy

# 问题：模型可能生成 "The answer is B" 而不是 "B"
# 解决：强制只能输出 A/B/C/D 中的一个

def generate_restricted(tokens, allowed_tokens):
    # 获取 logits
    logits = model(tokens)[:, -1, :]  # (1, vocab_size)
    
    # 找到允许的 token IDs
    allowed_ids = [tokenizer.encode_special(t) for t in allowed_tokens]
    # 例如：[65, 66, 67, 68] 对应 A, B, C, D
    
    # 只保留这些位置的 logits
    restricted_logits = logits[:, allowed_ids]  # (1, 4)
    
    # 选择最大的
    best_idx = restricted_logits.argmax()
    return allowed_tokens[best_idx]  # "A" or "B" or "C" or "D"

# ARC (AI2 Reasoning Challenge)
# 类似 MMLU，也是选择题

arc_easy_acc = run_chat_eval("ARC-Easy", model, ...)
arc_challenge_acc = run_chat_eval("ARC-Challenge", model, ...)

# 验证损失：频繁评估
eval_every = 100  # 每 100 步

# 任务准确率：不太频繁
eval_metrics_every = 200  # 每 200 步

# 原因：
# - 验证损失快速（只是前向传播）
# - 任务准确率慢（需要生成 + 解析）

# ===== Base/Mid Training =====
loss = F.cross_entropy(
    logits.view(-1, vocab_size),
    targets.view(-1)
    # 所有位置都计算损失
)

# ===== SFT =====
loss = F.cross_entropy(
    logits.view(-1, vocab_size),
    targets.view(-1),
    ignore_index=-1  # ← 关键差异！
)
# targets 中 mask=0 的位置是 -1，被忽略