CS336 Assignment 1 后半段记录：实验

已经搭建起了训练循环，接下来做一些实验。训得有多快？有多好？

BPE Tokenizer

之前的实现接口不是很好用，更进一步包装了一下，并麻烦GPT帮我优化了naive实现。

由于 test_tokenizer 中 import resource 这个库只能在linux中使用，决定将核心代码拷贝到linux环境下进行测试。
而为了保证实现的正确性，在原来的windows环境下利用gpt撰写了 test_tokenizer_windows.py 用于替换。
成功通过了所有测试。

在实际的训练过程中大约有如下的日志：

============================================================
Step 1/3: train BPE tokenizer
BPE training configuration:
  input_path: data/TinyStoriesV2-GPT4-train.txt
  vocab_size: 10000
  special_tokens: ['<|endoftext|>']
  progress_every: 100
  heartbeat_seconds: 15
  vocab_out: tokenizer/tinystories_bpe_vocab.pkl
  merges_out: tokenizer/tinystories_bpe_merges.pkl
[progress] read_corpus done in 12.2s, chars=2226845268
[heartbeat] elapsed=15.0s | stage=pretokenize
...
[heartbeat] elapsed=11m00.2s | stage=pretokenize[progress] pretokenize done in 10m48.0s, unique_words=59933, word_instances=536592168
[progress] init_vocab done, size=257

[progress] merge stage started. target_merges=9743
[progress] merges 1/9743 (0.0%) | rate 46.72/s | eta 3m28.5s
[progress] merges 100/9743 (1.0%) | rate 69.60/s | eta 2m18.5s
[progress] merges 200/9743 (2.1%) | rate 111.81/s | eta 1m25.3s
[progress] merges 300/9743 (3.1%) | rate 147.55/s | eta 1m04.0s
[progress] merges 400/9743 (4.1%) | rate 174.47/s | eta 53.5s
[progress] merges 500/9743 (5.1%) | rate 196.80/s | eta 47.0s
[progress] merges 600/9743 (6.2%) | rate 213.98/s | eta 42.7s
[progress] merges 700/9743 (7.2%) | rate 230.18/s | eta 39.3s
[progress] merges 800/9743 (8.2%) | rate 234.73/s | eta 38.1s
[progress] merges 900/9743 (9.2%) | rate 245.64/s | eta 36.0s
[progress] merges 1000/9743 (10.3%) | rate 254.83/s | eta 34.3s
[progress] merges 1100/9743 (11.3%) | rate 263.16/s | eta 32.8s
[progress] merges 1200/9743 (12.3%) | rate 269.57/s | eta 31.7s
[progress] merges 1300/9743 (13.3%) | rate 274.76/s | eta 30.7s
[progress] merges 1400/9743 (14.4%) | rate 277.75/s | eta 30.0s
[progress] merges 1500/9743 (15.4%) | rate 280.87/s | eta 29.3s
[progress] merges 1600/9743 (16.4%) | rate 282.89/s | eta 28.8s
...
[progress] merges 9700/9743 (99.6%) | rate 234.68/s | eta 183ms
[progress] merges 9743/9743 (100.0%) | rate 234.48/s | eta 0ms
[progress] merge stage finished. completed=9743 in 41.6s
[progress] train done in 11m41.9s, merges=9743, final_vocab=10000
Saved tokenizer artifacts:
  vocab_path: tokenizer/tinystories_bpe_vocab.pkl
  merges_path: tokenizer/tinystories_bpe_merges.pkl
  longest_token: b' accomplishment'
  longest_token_len: 15
  total_elapsed: 11m42.2s
============================================================
Step 2/3: tokenize txt -> bin
Tokenization config:
  input_text: data/TinyStoriesV2-GPT4-train.txt
  vocab_pkl: tokenizer/tinystories_bpe_vocab.pkl
  merges_pkl: tokenizer/tinystories_bpe_merges.pkl
  output_bin: data/tinystories_train.bin
  output_meta: data/tinystories_train.bin.meta.json
  special_tokens: ['<|endoftext|>']
Loaded tokenizer in 0.02s, selected dtype=uint16
[count] lines=10,000 tokens=346,490 rate=757,865 tok/s
[count] lines=20,000 tokens=688,564 rate=920,827 tok/s
[count] lines=30,000 tokens=1,030,877 rate=1,050,978 tok/s
...
[count] lines=15,560,000 tokens=539,410,167 rate=1,686,385 tok/s
[count] lines=15,570,000 tokens=539,755,742 rate=1,686,406 tok/s
[count] lines=15,580,000 tokens=540,104,767 rate=1,686,436 tok/s
[count] lines=15,590,000 tokens=540,453,813 rate=1,686,456 tok/s
[count] lines=15,600,000 tokens=540,794,508 rate=1,686,484 tok/s
Count pass done: lines=15,600,057, tokens=540,796,778, elapsed=320.67s
[write] lines=10,000 tokens_written=346,490/540,796,778 rate=799,121 tok/s
[write] lines=20,000 tokens_written=688,564/540,796,778 rate=1,021,996 tok/s
[write] lines=30,000 tokens_written=1,030,877/540,796,778 rate=1,127,432 tok/s
[write] lines=40,000 tokens_written=1,378,058/540,796,778 rate=1,005,091 tok/s
...
[write] lines=150,000 tokens_written=5,191,898/5,461,210 rate=179,838 tok/s Write pass done: output=data/tinystories_val.bin, elapsed=29.08s Metadata written: data/tinystories_val.bin.meta.json ============================================================

Training loop / Generating text

我们要做什么？

5.1 Data Loader
编写DataLoader，从输入序列，生成一系列的batch。
还包括 np.memmap 这样的数据集加载工程优化。

5.2 Checkpointing
编写Checkpointing，将模型存储下来。

5.3 Training Loop
编写训练循环，并且要求具有日志功能（例如接到WanDB上）。

6 Generating Text
使用训练出的模型进行next token prediction，从而实现生成文本。

DataLoader

tokenized data是一个token序列 $x=(x_1,x_2, \dots, x_n)$ ，
DataLoader把它变成batch流，
每个batch包含 $B$ 个 (sequence, next_tokens) ，其中 sequence 和 next_tokens 的长度均为 $m$ 。
例如一个 $B=1,m=3$ 的batch可以为 $([x_2,x_3,x_4],[x_3,x_4,x_5])$ 。用于next token prediction的原始序列是 $[x_2,x_3,x_4]$ ，它需要预测出的 token 是 $[x_3,x_4,x_5]$ ，正好是原始序列整体往后一个的窗口。
具体而言，这个任务是masked的：

看到 x₂，预测 x₃
看到 x₂, x₃，预测 x₄
看到 x₂, x₃, x₄，预测 x₅

def get_batch(
    dataset: np.ndarray,
    batch_size: int,
    context_length: int,
    device: str
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    从数据集采样batch
    
    Args:
        dataset: 1D numpy数组，包含所有token IDs
        batch_size: 批次大小 B
        context_length: 上下文长度 m
        device: PyTorch设备字符串
    
    Returns:
        (inputs, targets): 两个形状为 (batch_size, context_length) 的张量
    """
    dataset_len = len(dataset)
    inputs = torch.empty(batch_size, context_length, dtype=torch.long)
    targets = torch.empty(batch_size, context_length, dtype=torch.long)
    
    for i in range(batch_size):
        start_idx = torch.randint(0, dataset_len-context_length, (1, )).item()
        input_seq = dataset[start_idx: start_idx+context_length]
        input_target = dataset[start_idx+1: start_idx+context_length+1]
        
        inputs[i] = torch.tensor(input_seq, dtype=torch.long)
        targets[i] = torch.tensor(input_target, dtype=torch.long)
    
    inputs = inputs.to(device=device)
    targets = targets.to(device=device)
    return (inputs, targets)

Experiments

我们要做什么？
一个小的数据集（TinyStories）的反复训练，四处模块作用的消融实验，然后是一个更大的数据集的训练。

7.1 How to Run Experiments and Deliverables
确认wanDB等日志系统成功接入。

uv sync --upgrade-package wandb

7.2 TinyStories
给定初始参数，尝试在TinyStories训练，并寻找下述参数的好的初始值。
learning rate, learning rate warmup, other AdamW hyperparameters (β1, β2, ϵ), and weight decay.
然后进行文本生成。

7.3 Ablations and architecture modification
做RMSNorm，pre-norm，RoPE，SwiGLU的消融实验。

7.4 Running on OpenWebText
使用更大的数据集，训练语言模型。

7.5 Your own modification + leaderboard
提供了一些其他的工程实践参考（比如权重共享），以及打榜。

第一次实验

我当前做了什么？

√完成了训练循环的最小化测试。
√在5090环境中配置了cs336和minitorch两个环境，并正在上传数据，用于训练。
√启动第一个默认训练，tokenize正常，但默认的pytorch版本在5090有点小问题。需要改到2.7.1。

--skip-bpe 用于跳过tokenizer训练,
--skip-tokenize 用于跳过将tinystories从txt转为bin的过程.

我当前做了什么？

√正在准备minitorch和cs336的数据。
√跑起来cs336的第一次模型训练。

训练指令如下：（我提前跑过bpe和数据转二进制了，所以skip）

bash scripts/run_tinystories_train.sh --use-wandb --skip-bpe --skip-tokenize

第一次模型训练大致如下。

检查一下，文档要求的token budget是 327,680,000 ，符合我们的 batch size * context_length * steps = 64 * 256 * 20000 。这个budget大约是题目所示的参数量（17M）的20倍，符合Chinchilla定律。
检查一下，发现模型大小大约是22M，这与文档中提到的约17M有所差别。（简单思考一下，应该是LM head和embedding层没有共享权重（测试要求），那么按照估算公式，差个 Vd = 10000 * 512 刚好是5M。完美！）
检查一下，learning rate的变化是先急剧上升到 6e-4 ，然后先缓、再快、再缓地下降到 6e-5 。这符合cosine学习率调度的想法：最开始lr急剧上升而不是直接设置为 6e-4 是因为一开始模型随机初始化、对数据分布没有什么了解，如果贸然采取大学习率，可能使得它陷入不稳定的初始状态；再之后保持一段时间的大学习率，快速探索搜索空间，使得模型快速收敛；再快速转到小学习率，在最优解附近稳定且精细地探索。

============================================================
Step 3/3: train model
train_bin: data/tinystories_train.bin
val_bin: data/tinystories_val.bin
steps: 20000
warmup: 400
save_dir: runs/tinystories_base

Configuration:
train_data: data/tinystories_train.bin
val_data: data/tinystories_val.bin
data_dtype: uint16
vocab_size: 10000
context_length: 256
d_model: 512
num_heads: 16
d_ff: 1344
num_layers: 4
rope_theta: 10000.0
optimizer: custom_adamw
learning_rate: 0.0006
min_learning_rate: 6e-05
warmup_iters: 400
total_iters: 20000
beta1: 0.9
beta2: 0.95
eps: 1e-08
weight_decay: 0.1
grad_clip: 1.0
batch_size: 64
device: auto
seed: 1337
log_interval: 50
eval_interval: 500
eval_iters: 50
save_interval: 1000
save_dir: runs/tinystories_base
resume: None
wandb: True
wandb_project: cs336-assignment1
wandb_entity: None
wandb_run_name: tinystories_base_bs64_ctx256_20260310_000840
wandb_mode: online

Loading train data from: data/tinystories_train.bin
Loading val data from: data/tinystories_val.bin
Device: cuda
Trainable parameters: 22,696,448

wandb: Currently logged in as: 10pi (10pi-fudan-university-school-of-management) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
wandb: WARNING Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to https://wandb.me/wandb-init.
wandb: Tracking run with wandb version 0.21.2
wandb: Run data is saved locally in runs/tinystories_base/wandb/run-20260310_000927-d0vkurgn
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run tinystories_base_bs64_ctx256_20260310_000840
wandb: ⭐️ View project at https://wandb.ai/10pi-fudan-university-school-of-management/cs336-assignment1
wandb: 🚀 View run at https://wandb.ai/10pi-fudan-university-school-of-management/cs336-assignment1/runs/d0vkurgn

Starting training...
step 50 | train_loss 8.5434 | lr 7.500e-05 | tok/s 94,349
...
step   19950 | train_loss 1.4084 | lr 6.001e-05 | tok/s 266,695
step   20000 | train_loss 1.3999 | lr 6.000e-05 | tok/s 266,221
step   20000 | val_loss 1.4049
  saved new best checkpoint to runs/tinystories_base/best.pt (val_loss=1.4049)
  checkpoint saved: runs/tinystories_base/step_00020000.pt
wandb:                                                                                
wandb: 
wandb: Run history:
wandb:                    step ▁▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▅▅▅▆▆▆▆▆▆▆▇▇▇▇▇▇█████
wandb:              train/loss ██▇▆▅▄▄▄▄▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb:                train/lr ███████▇▇▇▆▆▆▆▆▆▅▅▅▅▅▄▄▄▄▄▃▃▃▃▃▂▂▂▁▁▁▁▁▁
wandb: train/tokens_per_second ▇▇▇█████▁██████████▂▃█████▃██▄██████████
wandb:                val/loss █▅▄▄▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb: 
wandb: Run summary:
wandb:                    step 20000
wandb:              train/loss 1.39992
wandb:                train/lr 6e-05
wandb: train/tokens_per_second 266221.25455
wandb:                val/loss 1.40486
wandb: 
wandb: 🚀 View run tinystories_base_bs64_ctx256_20260310_000840 at: https://wandb.ai/10pi-fudan-university-school-of-management/cs336-assignment1/runs/d0vkurgn
wandb: ⭐️ View project at: https://wandb.ai/10pi-fudan-university-school-of-management/cs336-assignment1
wandb: Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
wandb: Find logs at: runs/tinystories_base/wandb/run-20260310_000927-d0vkurgn/logs
Training complete in 1320.5s. Final checkpoint: runs/tinystories_base/final.pt
Best validation loss 1.4049 at step 20000.
============================================================
Done.
Tokenizer:
  tokenizer/tinystories_bpe_vocab.pkl
  tokenizer/tinystories_bpe_merges.pkl
Tokenized data:
  data/tinystories_train.bin
  data/tinystories_val.bin
Checkpoints:
  runs/tinystories_base

再看看图表，发现loss按照幂律（Scaling Law）下降得很健康，lr也是很合理的cosine，先快后慢半个周期。
我在5090的分词器训练大约11分钟，不包含分词器在内模型训练大约运行了22分钟，整体符合文档中30-40分钟的预期。