10GB代码用于GPT-2训练 | AI生成和翻译

Home 2026.06

综合摘要：从10GB代码数据到训练

== 第一阶段：分词处理 ==

目标：将38个Parquet文件（10.2 GB，GitHub代码数据集）转换为nanoGPT二进制格式（.npy分片）。

实际情况：原始脚本使用multiprocessing.Pool(16)。每个工作进程通过pq.read_table()加载整个Parquet文件，然后将词元累积到Python列表中，再转换为NumPy。

致命问题：一个Parquet文件产生约3.69亿个词元。包含3.69亿个整数的Python列表占用约10 GB内存（每个Python整数开销约28字节）。16个工作进程并行运行时，在62 GB内存的机器上总内存需求超过160 GB。内核要么因OOM杀死工作进程，要么导致交换空间过度使用，以至于sshd甚至无法发送SSH横幅——TCP端口22虽能接受连接，但在横幅交换期间挂起。需要物理重启。

解决方案：重写脚本，包含三项改动：

  1. 单进程——完全取消多进程处理
  2. 通过pq.ParquetFile.iter_batches(batch_size=8192)流式读取Parquet，而非加载整个文件
  3. 直接进行NumPy uint16累积——预分配200 MB缓冲区，无需中间Python列表

结果：141个分片，约14.07B个词元，27 GB输出，耗时41分钟，峰值内存约600 MB。

== 第二阶段：训练环境搭建 ==

可用基础设施：

== 第三阶段：模型规模选择 ==

为何选择GPT-2 124M（而非760M或350M）：

760M配置（n_layer=24, n_head=24, n_embd=1536）专为192 GB HBM3的MI300X设计。RTX 4070仅有12 GB。简单计算：

GPT-2 124M（n_layer=12, n_head=12, n_embd=768）：

规模定律考量：14B词元 / 124M参数 ≈ 每个参数113个词元。Chinchilla最优值为每个参数约20个词元，因此从数据角度看实际属于过度训练——这意味着每个参数质量更高。数据利用效果更好。

== 第四阶段：冒烟测试 ==

测试1——batch_size=8, grad_accum=4： 结果：OOM。已占用10.37 GB，还需1.54 GB。评估前向传播（50304个词表 × 1024个位置的logits）是罪魁祸首。

测试2——batch_size=4, grad_accum=8： 结果：成功。10步内损失从10.77降至8.03。约700ms/步，MFU为12.83%。

冒烟测试关键观察：

== 第五阶段：配置决策 ==

batch_size=4（而非8）： 冒烟测试证明8无法容纳。4可容纳且有空间进行前向/反向传播。每个微步骤处理4 × 1024 = 4,096个词元。

gradient_accumulation_steps=8（而非4）： 保持有效批次大小为每步32,768个词元（4096 × 8）。这是GPT-2规模训练的标准有效批次大小——足够大以实现稳定梯度，足够小以保证良好学习动态。

block_size=1024： GPT-2标准值。与模型架构的上下文窗口匹配。无需更改。

learning_rate=6e-4（而非3e-4）： 760M配置使用3e-4。较小模型可承受更高学习率——这是标准规模定律做法。6e-4是GPT-3论文中其125M模型使用的学习率。

min_lr=6e-5： 峰值学习率的10%。标准余弦退火下限。

warmup_iters=2000： 该模型规模的标准值。防止训练初期不稳定。

max_iters=427000： 14B词元 / 每步32,768个词元 ≈ 427,000步。这将一次性消耗整个数据集。对于从头开始的预训练，一个epoch是标准做法（LLM从同一数据的多次遍历中获益有限）。

compile=True： torch.compile融合内核并减少开销。初始编译（约2-3分钟）后预计提升2-3倍。冒烟测试中的警告”没有足够的SM来使用max_autotune_gemm模式”无害——RTX 4070有46个SM，低于自动调优的阈值，但编译仍有帮助。

== 预计训练时间 ==

未使用compile：约700ms/步 × 427,000步 ≈ 83小时（约3.5天） 使用compile（预期2倍提升）：约41小时（约1.7天） 使用compile（乐观3倍提升）：约28小时（约1.2天）

== 需关注事项 ==

Back Donate