noao-chat-5-训练四阶段数据报告

# 数据源配置
BASE_URL = "https://huggingface.co/datasets/karpathy/fineweb-edu-100b-shuffle/resolve/main"
MAX_SHARD = 1822  # 共 1823 个 shard (shard_00000 到 shard_01822)

# 从 dataset.py
def parquets_iter_batched(split, start=0, step=1):
    """
    - split: "train" 或 "val"
    - start/step: 用于 DDP 分布式训练
    """
    parquet_paths = list_parquet_files()
    
    if split == "train":
        parquet_paths = parquet_paths[:-1]  # 前 1822 个文件
    else:  # val
        parquet_paths = parquet_paths[-1:]  # 最后 1 个文件
    
    for filepath in parquet_paths:
        pf = pq.ParquetFile(filepath)
        for rg_idx in range(start, pf.num_row_groups, step):
            rg = pf.read_row_group(rg_idx)
            texts = rg.column('text').to_pylist()
            yield texts

# 从 loss_eval.py
def evaluate_bpb(model, val_loader, eval_steps, token_bytes):
    total_loss = 0.0
    total_tokens = 0
    
    for _ in range(eval_steps):
        inputs, targets = next(val_loader)  # (B, T)
        loss = model(inputs, targets)  # CrossEntropy loss
        
        # 统计非 padding token 的数量
        num_tokens = (targets != -1).sum()
        total_loss += loss * num_tokens
        total_tokens += num_tokens
    
    avg_loss = total_loss / total_tokens  # 平均 loss
    
    # token_bytes: (vocab_size,) 每个 token 对应的字节数
    bytes_per_token = token_bytes[targets[targets != -1]].float().mean()
    
    # BPB = Loss 转换为 bits，再除以每 token 的字节数
    bpb = avg_loss / math.log(2) / bytes_per_token
    
    return bpb

# 从 base_eval.py
def evaluate_model(model, tokenizer, device, max_per_task=500):
    results = {}
    
    # 评估 9 个任务
    for task_name, task_config in core_tasks.items():
        accuracy = evaluate_task(
            model, 
            tokenizer, 
            task_config, 
            max_examples=max_per_task
        )
        results[task_name] = accuracy
    
    # Centered normalization（减去基线模型的表现）
    centered_results = {}
    for task_name, accuracy in results.items():
        baseline = BASELINES[task_name]  # 基线模型（如 random guess）
        centered_results[task_name] = accuracy - baseline
    
    # CORE Metric = 所有 centered 结果的平均值
    core_metric = sum(centered_results.values()) / len(centered_results)
    
    return {
        "core_metric": core_metric,
        "centered_results": centered_results,
        "raw_results": results
    }

# base_train.py 默认配置
depth = 20                    # 模型深度
max_seq_len = 2048           # 最大序列长度
device_batch_size = 32       # 每张卡的 batch size
total_batch_size = 524288    # 总 batch size（tokens）
target_param_data_ratio = 20 # Chinchilla ratio（数据:参数 = 20:1）

# 学习率
embedding_lr = 0.2           # Embedding 层（AdamW）
unembedding_lr = 0.004       # Unembedding 层（AdamW）
matrix_lr = 0.02             # 其他矩阵参数（Muon）
weight_decay = 0.0

# 学习率调度
warmup_ratio = 0.0           # 预热阶段占比
warmdown_ratio = 0.2         # 衰减阶段占比
final_lr_frac = 0.0          # 最终学习率倍数

num_params = 模型参数量（例如 120M）
target_tokens = 20 * num_params  # Chinchilla ratio
num_iterations = target_tokens // total_batch_size

# 例如：120M 参数的模型
# target_tokens = 20 * 120M = 2.4B tokens
# num_iterations = 2.4B / 524288 ≈ 4577 steps

# 从 mid_train.py
train_dataset = TaskMixture([
    SmolTalk(split="train"),                 # 460K 对话
    MMLU(subset="auxiliary_train", split="train"),  # 100K 问题
    GSM8K(subset="main", split="train"),     # 8K 问题
    CustomJSON(filepath=identity_conversations_filepath),  # 1K 对话
    CustomJSON(filepath=identity_conversations_filepath),  # 1K 对话（2 epochs）
    SimpleSpelling(size=200000, split="train"),  # 200K 对话
    SpellingBee(size=80000, split="train"),  # 80K 对话
])
# 总计：460K + 100K + 8K + 1K + 1K + 200K + 80K = 850K 对话

{
  "messages": [
    {
      "role": "user",
      "content": "Multiple Choice question: What is the chemical symbol for gold?\n- Gold=A\n- Au=B\n- Ag=C\n- Fe=D\n\nRespond only with the letter of the correct answer."
    },
    {
      "role": "assistant",
      "content": "B"
    }
  ]
}

# TaskMixture 实现（从 tasks/common.py）
class TaskMixture(Task):
    def __init__(self, tasks):
        self.tasks = tasks
        self.lengths = [len(task) for task in tasks]
        
        # 构建索引映射
        self.index_map = []
        for task_idx, task_length in enumerate(self.lengths):
            for local_idx in range(task_length):
                self.index_map.append((task_idx, local_idx))
        
        # 确定性 shuffle（seed=42）
        rng = random.Random(42)
        rng.shuffle(self.index_map)
    
    def get_example(self, index):
        task_idx, local_idx = self.index_map[index]
        return self.tasks[task_idx][local_idx]

# 从 mid_train.py
val_dataset = TaskMixture([
    SmolTalk(split="test"),           # 24K 对话
    MMLU(subset="all", split="test", stop=5200),  # 5.2K 问题（从 14K 中取）
    GSM8K(subset="main", split="test", stop=420),  # 420 问题（从 1.3K 中取）
])
# 总计：24K + 5.2K + 0.42K ≈ 30K 对话

# 从 mid_train.py
eval_every = 150              # 每 150 步评估一次
eval_tokens = 20 * 524288     # 评估 10M tokens

# mid_train.py 默认配置
max_seq_len = 2048
device_batch_size = 32
total_batch_size = 524288    # 与 Base 训练相同

# 学习率（与 Base 训练相同）
unembedding_lr = 0.004
embedding_lr = 0.2
matrix_lr = 0.02
weight_decay = 0.0
init_lr_frac = 1.0           # 初始 LR 倍数

# 学习率调度（与 Base 不同！）
def get_lr_multiplier(progress):
    # 前 80% 不衰减，后 20% 线性衰减到 0
    return 1 if progress < 0.8 else 1 - (progress - 0.8) / 0.2

# 从 chat_sft.py
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(filepath=identity_conversations_filepath),  # 1K 对话
    SimpleSpelling(size=300, split="train"),    # 300 对话
    SpellingBee(size=300, split="train"),       # 300 对话
])
# 总计：2.3K + 1.1K + 8K + 10K + 1K + 0.3K + 0.3K = 23K 对话

{
  "messages": [
    {
      "role": "user",
      "content": "Multiple Choice question: Which of the following is a renewable resource?\n- Coal=A\n- Oil=B\n- Solar energy=C\n- Natural gas=D\n\nRespond only with the letter of the correct answer."
    },
    {
      "role": "assistant",
      "content": "C"
    }
  ]
}

# 从 tokenizer.py
def render_conversation(self, conversation):
    """
    返回：
    - ids: List[int]，完整的 token 序列
    - mask: List[int]，0/1 标记（1=计算loss，0=不计算）
    """
    ids = []
    mask = []
    
    # 添加 BOS token
    ids.append(self.encode_special("<|bos|>"))
    mask.append(0)  # BOS 不参与 loss
    
    for message in conversation["messages"]:
        role = message["role"]
        content = message["content"]
        
        if role == "user":
            # User message: 全部 mask=0
            ids.extend(self.encode_special("<|user_start|>"))
            ids.extend(self.encode(content))
            ids.extend(self.encode_special("<|user_end|>"))
            mask.extend([0] * (len(ids) - len(mask)))
        
        elif role == "assistant":
            # Assistant message: 只有 content 部分 mask=1
            ids.append(self.encode_special("<|assistant_start|>"))
            mask.append(0)  # start token 不参与 loss
            
            content_ids = self.encode(content)
            ids.extend(content_ids)
            mask.extend([1] * len(content_ids))  # ← 关键！只有这里是 1
            
            ids.append(self.encode_special("<|assistant_end|>"))
            mask.append(1)  # end token 参与 loss
    
    return ids, mask

# 从 chat_sft.py
def sft_data_generator(dataset, batch_size):
    batch = []
    for conversation in dataset:
        ids, mask = tokenizer.render_conversation(conversation)
        batch.append((ids, mask))
        
        if len(batch) == batch_size:
            # Collate 成张量
            inputs = ...   # (B, T)
            targets = ...  # (B, T)
            
            # 关键：将 mask=0 的位置设为 -1（ignore index）
            for i, (ids, mask) in enumerate(batch):
                row_targets = ids[1:]  # shift right
                mask_tensor = mask[1:]
                row_targets[mask_tensor == 0] = -1  # ← 设为 ignore index
                targets[i] = row_targets
            
            yield inputs, targets

# PyTorch 内部
loss = F.cross_entropy(
    logits.view(-1, vocab_size),  # (B*T, V)
    targets.view(-1),              # (B*T,)
    ignore_index=-1                # ← targets=-1 的位置不计算 loss
)

# 从 chat_sft.py
val_ds = SmolTalk(split="test")  # 24K 对话

# 从 chat_sft.py
eval_every = 100
eval_steps = 100

# 评估过程
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()

# 从 chat_sft.py
eval_metrics_every = 200
eval_metrics_max_problems = 1024

metrics = {
    "mmlu_acc": run_chat_eval("MMLU", ...),
    "arc_easy_acc": run_chat_eval("ARC-Easy", ...),
    "arc_challenge_acc": run_chat_eval("ARC-Challenge", ...),
    "gsm8k_maj1_acc": run_chat_eval("GSM8K", ..., num_samples=1),
    "gsm8k_maj16_acc": run_chat_eval("GSM8K", ..., num_samples=16),
}

# chat_sft.py 默认配置
device_batch_size = 4         # 小 batch size（避免 OOM）
target_examples_per_step = 32 # 每步处理 32 个对话
num_epochs = 1                # 1 epoch

# 学习率（非常小！）
unembedding_lr = 0.004
embedding_lr = 0.2
matrix_lr = 0.02
init_lr_frac = 0.02          # ← 关键！只用 2% 的 LR
# 实际 LR = 基础 LR × 0.02
# 例如：matrix_lr = 0.02 × 0.02 = 0.0004

weight_decay = 0.0

# 学习率调度（线性衰减到 0）
def get_lr_multiplier(it):
    return 1.0 - it / num_iterations

dataset_size = 23000
target_examples_per_step = 32
num_epochs = 1

num_iterations = (dataset_size // target_examples_per_step) * num_epochs
# = (23000 // 32) * 1
# = 718 steps

# 从 chat_rl.py
train_task = GSM8K(subset="main", split="train")
# 7473 个数学问题

"""
GRPO = Group Relative Policy Optimization

简化版：
1. 删除 trust region（无 KL 散度约束）
2. On-policy（无需 PPO ratio+clip）
3. Token-level 优势归一化（GAPO 风格）
4. 只用 (r - mu) 作为 advantage（不除以 sigma）
"""

# 从 chat_rl.py
examples_per_step = 16  # 每步处理 16 个问题
num_samples = 16        # 每个问题生成 16 个回答

@torch.no_grad()
def get_batch():
    # 遍历训练集
    for example_idx in range(ddp_rank, len(train_task), ddp_world_size):
        
        # 1. 获取问题
        conversation = train_task[example_idx]
        
        # 2. Tokenize（保留 <|assistant_start|>，删除后面的内容）
        tokens = tokenizer.render_for_completion(conversation)
        prefix_length = len(tokens)
        
        # 3. 生成 16 个回答（批量生成，避免 OOM）
        model.eval()
        generated_token_sequences = []
        masks = []
        
        for sampling_step in range(num_samples // device_batch_size):
            seed = hash((step, example_idx, sampling_step)) & 0x7FFFFFFF
            
            sequences, masks_batch = engine.generate_batch(
                tokens,
                num_samples=device_batch_size,
                max_tokens=max_new_tokens,
                temperature=1.0,  # ← 高温度，鼓励探索
                top_k=50,
                seed=seed
            )
            
            generated_token_sequences.extend(sequences)
            masks.extend(masks_batch)
        
        # 4. 计算每个回答的奖励
        rewards = []
        for sample_tokens in generated_token_sequences:
            generated_tokens = sample_tokens[prefix_length:]
            generated_text = tokenizer.decode(generated_tokens)
            
            # 奖励函数：答案是否正确
            reward = train_task.reward(conversation, generated_text)
            rewards.append(reward)
        
        # 5. Padding 使所有序列等长
        max_length = max(len(seq) for seq in generated_token_sequences)
        padded_sequences = [
            seq + [pad_token] * (max_length - len(seq))
            for seq in generated_token_sequences
        ]
        padded_masks = [
            mask + [0] * (max_length - len(mask))
            for mask in masks
        ]
        
        # 6. 转换为张量
        ids = torch.tensor(padded_sequences, device=device)  # (16, T)
        mask_ids = torch.tensor(padded_masks, device=device)  # (16, T)
        
        inputs = ids[:, :-1]   # (16, T-1)
        targets = ids[:, 1:].clone()
        targets[mask_ids[:, 1:] == 0] = -1  # mask out
        
        rewards = torch.tensor(rewards, device=device)  # (16,)
        
        # 7. 计算 advantage（只减均值，不除标准差）
        mu = rewards.mean()
        advantages = rewards - mu  # (16,)
        
        # 8. Yield batch
        yield sequences, inputs, targets, rewards, advantages

# 从 tasks/gsm8k.py
def reward(self, problem, completion):
    """
    返回：1.0（正确）或 0.0（错误）
    """
    # 1. 从 completion 中提取答案
    predicted_answer = self.extract_answer(completion)
    
    # 2. 获取正确答案
    correct_answer = problem["answer"]
    
    # 3. 比较（需要处理数值格式：逗号、小数点等）
    return 1.0 if self.answers_match(predicted_answer, correct_answer) else 0.0

def extract_answer(self, completion):
    """
    从生成的文本中提取答案
    支持多种格式：
    - "The answer is 42"
    - "#### 42"
    - "<tool>calculator(...)</tool> = 42"
    """
    # 正则表达式匹配
    patterns = [
        r"#### (\d+)",
        r"The answer is (\d+)",
        r"= (\d+)",
    ]
    for pattern in patterns:
        match = re.search(pattern, completion)
        if match:
            return match.group(1)
    return None

# 从 chat_rl.py
# Training loop
for step in range(num_steps):
    
    # 1. 获取 batch（包含 rollouts）
    sequences, inputs, targets, rewards, advantages = next(get_batch())
    # inputs: (16, T-1)
    # targets: (16, T-1)，mask 后
    # rewards: (16,)
    # advantages: (16,)
    
    # 2. Forward pass（计算 log probs）
    model.train()
    logits = model(inputs)  # (16, T-1, V)
    
    log_probs = F.log_softmax(logits, dim=-1)  # (16, T-1, V)
    
    # 3. 收集 targets 位置的 log probs
    # log_probs_taken: (16, T-1)
    log_probs_taken = log_probs.gather(
        dim=-1,
        index=targets.unsqueeze(-1).clamp(min=0)
    ).squeeze(-1)
    
    # 4. Mask out padding 和 prompt
    mask = (targets != -1).float()  # (16, T-1)
    log_probs_taken = log_probs_taken * mask
    
    # 5. Policy Gradient Loss
    # L = -mean(log_prob * advantage)
    advantages_expanded = advantages.unsqueeze(-1)  # (16, 1)
    
    # Token-level loss
    token_losses = -log_probs_taken * advantages_expanded  # (16, T-1)
    
    # 只对非 mask 的 token 求均值
    loss = (token_losses * mask).sum() / mask.sum()
    
    # 6. Backward + Optimizer step
    loss.backward()
    
    # Gradient accumulation
    if (step + 1) % grad_accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 从 chat_rl.py
val_task = GSM8K(subset="main", split="test")
# 1319 个问题

# 从 chat_rl.py
def run_gsm8k_eval(
    task,
    tokenizer,
    engine,
    max_examples=400,
    num_samples=16,
    temperature=1.0,
):
    correct_counts = {1: 0, 4: 0, 16: 0}
    total = 0
    
    for problem in task[:max_examples]:
        conversation = problem
        tokens = tokenizer.render_for_completion(conversation)
        
        # 生成 num_samples 个回答
        sequences, _ = engine.generate_batch(
            tokens,
            num_samples=num_samples,
            max_tokens=256,
            temperature=temperature,
            top_k=50,
        )
        
        # 评估每个回答
        is_correct = []
        for seq in sequences:
            generated_text = tokenizer.decode(seq[len(tokens):])
            reward = task.reward(conversation, generated_text)
            is_correct.append(reward == 1.0)
        
        # 计算 Pass@k（k = 1, 4, 16）
        correct_counts[1] += int(is_correct[0])  # Pass@1 = 第 1 次正确
        correct_counts[4] += int(any(is_correct[:4]))  # Pass@4
        correct_counts[16] += int(any(is_correct[:16]))  # Pass@16
        
        total += 1
    
    return {
        "pass@1": correct_counts[1] / total,
        "pass@4": correct_counts[4] / total,
        "pass@16": correct_counts[16] / total,
    }

eval_every = 60              # 每 60 步评估一次
eval_examples = 400          # 使用 400 个测试样本
num_samples = 16             # 每个问题生成 16 次
temperature = 1.0            # 高温度，保持探索性

# chat_rl.py 默认配置
device_batch_size = 8        # 每张卡的 batch size
examples_per_step = 16       # 每步处理 16 个问题
num_samples = 16             # 每个问题生成 16 个回答
max_new_tokens = 256         # 最大生成长度
temperature = 1.0            # 采样温度（高温度，鼓励探索）
top_k = 50                   # Top-k 采样

# 学习率（与 SFT 相似）
unembedding_lr = 0.004
embedding_lr = 0.2
matrix_lr = 0.02
weight_decay = 0.0
init_lr_frac = 0.05          # 5% 的基础 LR

num_epochs = 1               # 1 epoch over GSM8K

dataset_size = 7473
examples_per_step = 16
num_epochs = 1

num_steps = (dataset_size // examples_per_step) * num_epochs
# = (7473 // 16) * 1
# = 467 steps

# User 的输入：mask = 0（不计算 loss）
# Assistant 的回复：mask = 1（计算 loss）
targets[mask == 0] = -1  # ignore_index