
Karpathy 在 README 中引用了一段半讽刺半预言的"2026 年 3 月"的话:
曾经,前沿 AI 研究是由"人肉计算机"在吃饭、睡觉、找乐子之间,偶尔通过"声音波互联"在"组会"仪式中同步完成的。那个时代早已过去。研究现在完全是自主 AI 代理集群在天空中的计算集群巨型结构上运行的领域。
很有趣。但也看看机器学习研究今天实际是如何工作的。研究者有一个想法。他们实现它。他们启动训练。他们等待。他们检查结果。他们又有一个想法。他们去吃午饭。他们回来。他们尝试别的东西。实际的"思考"部分可能只占整个周期的 20%。其余是机械性的:编辑代码、照看运行、上下文切换。
autoresearch 提出的问题是:如果你把人类从内循环中完全移除会怎样?
不是从外 循环中移除。仍然需要有人决定研究什么问题、优化什么指标、施加什么约束。这就是 program.md。这是人类的工作。但是想法→实现→运行→评估→保留/丢弃这个乏味的循环呢?那是可自动化的。那是代理擅长的。
设计通过一系列精心设计的约束使这成为可能:
每个约束都服务于同一个目的:使实验循环完全可通过单一的保留/丢弃决策实现自动化。没有歧义。不需要人类判断。代理读取一个数字,与前一个数字比较,然后行动。
有了这个框架,让我们看看每段代码如何实现这一愿景。
1、系统架构概览
现在我们理解了为什么 系统要这样构建,让我们看看这三个文件如何相互关联:
关注点的分离是外科手术式的:
prepare.py是宪法:实验的不可变法则(时间预算、评估指标、数据加载)。代理不能触碰它。train.py是游乐场:从模型架构到优化器超参数,一切都可以修改。program.md是管理层:人类编写的指令集,将通用 LLM 变成专注的机器学习研究者。
2、第一部分:prepare.py,不可变的基础
这个文件有两个用途:一次性数据准备(由人类运行一次),以及在每次实验中被 train.py 导入的运行时工具。说实话,真正的约束就在这里。让我们逐块分解。
2.1 定义整个实验的三个数字
MAX_SEQ_LEN = 2048 # 上下文长度
TIME_BUDGET = 300 # 训练时间预算(秒)(5 分钟)
EVAL_TOKENS = 40 * 524288 # 验证评估的 token 数量
为什么这很重要: 这三个数字定义了整个实验方案。
-
MAX_SEQ_LEN = 2048:上下文窗口。代理创建的每个模型变体都必须在这个固定序列长度内运行,确保公平比较。 -
TIME_BUDGET = 300:5 分钟的壁钟预算是 autoresearch 中最重要的设计决策。它使实验依赖于硬件但独立于时间 。H100 在 5 分钟内会处理比 RTX 4090 更多的 token,但两者都恰好获得 5 分钟。代理为你的硬件优化。 -
EVAL_TOKENS = 40 * 524288:大约 2000 万 token 用于验证。足够大以保证统计可靠性,又足够小以不影响训练时间。CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "autoresearch")
DATA_DIR = os.path.join(CACHE_DIR, "data")
TOKENIZER_DIR = os.path.join(CACHE_DIR, "tokenizer")
BASE_URL = "https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle/resolve/main"
MAX_SHARD = 6542
VAL_SHARD = MAX_SHARD # 固定的验证分片 (shard_06542)
VOCAB_SIZE = 8192
数据集架构: 数据以 6543 个 parquet 分片的形式存储在 HuggingFace 上,来自 climbmix-400b-shuffle 数据集。最后一个分片 (shard_06542) 永久固定为验证集;它永远不会出现在训练数据中。
VOCAB_SIZE = 8192 故意设置得很小。更小的词表意味着更小的嵌入表,这意味着在时间预算内更快的迭代。这是为实验速度优化的,而不是为生产规模的语言建模。
2.2 原子下载模式
def download_single_shard(index):
"""下载一个 parquet 分片,带重试。成功返回 True。"""
filename = f"shard_{index:05d}.parquet"
filepath = os.path.join(DATA_DIR, filename)
if os.path.exists(filepath):
return True
url = f"{BASE_URL}/{filename}"
max_attempts = 5
for attempt in range(1, max_attempts + 1):
try:
response = requests.get(url, stream=True, timeout=30)
response.raise_for_status()
temp_path = filepath + ".tmp"
with open(temp_path, "wb") as f:
for chunk in response.iter_content(chunk_size=1024 * 1024):
if chunk:
f.write(chunk)
os.rename(temp_path, filepath)
return True
except (requests.RequestException, IOError) as e:
# ... 指数退避重试 ...
if attempt < max_attempts:
time.sleep(2 ** attempt)
return False
先写入临时文件再重命名的模式是这里的安全机制。 注意 temp_path = filepath + ".tmp" 后面跟着 os.rename(temp_path, filepath)。如果下载在中途崩溃,你不会得到一个看起来"已下载"但实际损坏的 .parquet 文件。在大多数文件系统上,os.rename 是原子操作;文件要么完全写入,要么不存在。
指数退避 (time.sleep(2 ** attempt)) 是教科书式的:重试之间延迟 2 秒、4 秒、8 秒、16 秒。
def download_data(num_shards, download_workers=8):
"""下载训练分片 + 固定的验证分片。"""
os.makedirs(DATA_DIR, exist_ok=True)
num_train = min(num_shards, MAX_SHARD)
ids = list(range(num_train))
if VAL_SHARD not in ids:
ids.append(VAL_SHARD)
# ...
workers = max(1, min(download_workers, needed))
with Pool(processes=workers) as pool:
results = pool.map(download_single_shard, ids)
带验证保证的并行下载: 即使你只下载 10 个训练分片,验证分片 (shard_06542) 总是 被包含在内。使用 8 个工作进程的 multiprocessing.Pool 并行化 I/O 密集型的下载。
2.3 为什么用 rustbpe 而不是 HuggingFace Tokenizers?
def train_tokenizer():
tokenizer = rustbpe.Tokenizer()
vocab_size_no_special = VOCAB_SIZE - len(SPECIAL_TOKENS)
tokenizer.train_from_iterator(
text_iterator(), vocab_size_no_special, pattern=SPLIT_PATTERN
)
为什么用 rustbpe 而不是标准的 tokenizer 库? 速度。rustbpe 是一个 Rust 实现的 BPE,训练速度显著快于纯 Python 替代方案。tokenizer 在最多 10 亿个字符上训练(通过 text_iterator(max_chars=1_000_000_000))。
分词模式是 GPT-4 风格,但有一个关键区别:
SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
\p{N}{1,2} 的调整微妙但重要。 与 GPT-4 的 {1,3} 不同,这将数字 token 限制为最多 2 位数字。这迫使 tokenizer 将大数字分解为更小的块;对于无法在很少见的 3 位数字 token 上浪费条目的小词表来说更好。
然后将 tokenizer 包装成 tiktoken.Encoding 以实现兼容性并进行 pickle 序列化:
enc = tiktoken.Encoding(
name="rustbpe",
pat_str=pattern,
mergeable_ranks=mergeable_ranks,
special_tokens=special_tokens,
)
with open(tokenizer_pkl, "wb") as f:
pickle.dump(enc, f)
token_bytes 查找表对评估至关重要:
token_bytes_list = []
for token_id in range(enc.n_vocab):
token_str = enc.decode([token_id])
if token_str in special_set:
token_bytes_list.append(0)
else:
token_bytes_list.append(len(token_str.encode("utf-8")))
token_bytes_tensor = torch.tensor(token_bytes_list, dtype=torch.int32)
这预计算了每个 token 代表多少 UTF-8 字节。 对于 bits-per-byte 指标必不可少。你需要知道每个 token 的字节长度才能将每 token 损失转换为每字节损失。
2.4 不浪费任何 Token 的数据加载器
哦!这是整个仓库中最复杂的代码之一:
def make_dataloader(tokenizer, B, T, split, buffer_size=1000):
"""
BOS 对齐的数据加载器,使用最佳适配打包。
每行以 BOS 开头。文档使用最佳适配打包以最小化裁剪。
100% 利用率(无填充)。
"""
assert split in ["train", "val"]
row_capacity = T + 1 # +1 因为目标偏移了 1
# ...
为什么是 T + 1****? 对于语言建模,如果输入是 tokens [0, 1, 2, ..., T-1],目标是 [1, 2, 3, ..., T]。你需要一个额外的 token 来创建最终目标。所以每行持有 T + 1 = 2049 个 token,然后分割成 inputs = row[:-1] 和 targets = row[1:]。
打包算法是一种最佳适配递减装箱策略:
while pos < row_capacity:
while len(doc_buffer) < buffer_size:
refill_buffer()
remaining = row_capacity - pos
# 找到能完全放下的最大文档
best_idx = -1
best_len = 0
for i, doc in enumerate(doc_buffer):
doc_len = len(doc)
if doc_len <= remaining and doc_len > best_len:
best_idx = i
best_len = doc_len
if best_idx >= 0:
doc = doc_buffer.pop(best_idx)
row_buffer[row_idx, pos:pos + len(doc)] = torch.tensor(doc)
pos += len(doc)
else:
# 没有文档能放下:裁剪最短的来填满剩余空间
shortest_idx = min(range(len(doc_buffer)),
key=lambda i: len(doc_buffer[i]))
doc = doc_buffer.pop(shortest_idx)
row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining])
pos += remaining
逻辑很优雅:
- 尝试放入最大的能完全放下的文档(最佳适配)。
- 如果没有能放下的,取最短 的文档并裁剪它来精确填充。
为什么要裁剪最短的? 因为裁剪短文档比裁剪长文档浪费更少的 token。如果你有一个 50 token 的文档和一个 5000 token 的文档,你需要 40 个 token 来填满一行,裁剪 50 token 的文档浪费 10 个 token。裁剪 5000 token 的文档浪费 4960 个 token。
零填充,100% 利用率。 每行的每个位置都包含一个真实的 token。这比用零填充短文档的朴素方法有显著的效率优势。
预分配的锁存内存缓冲区是另一个性能细节:
cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=True)
gpu_buffer = torch.empty(2 * B * T, dtype=torch.long, device="cuda")
# ...
gpu_buffer.copy_(cpu_buffer, non_blocking=True)
锁存内存通过 non_blocking=True 实现异步 CPU 到 GPU 传输,使数据加载与计算重叠。
2.5 为什么用 Bits Per Byte 而不是 Perplexity?
@torch.no_grad()
def evaluate_bpb(model, tokenizer, batch_size):
token_bytes = get_token_bytes(device="cuda")
val_loader = make_dataloader(tokenizer, batch_size, MAX_SEQ_LEN, "val")
steps = EVAL_TOKENS // (batch_size * MAX_SEQ_LEN)
total_nats = 0.0
total_bytes = 0
for _ in range(steps):
x, y, _ = next(val_loader)
loss_flat = model(x, y, reduction='none').view(-1)
y_flat = y.view(-1)
nbytes = token_bytes[y_flat]
mask = nbytes > 0
total_nats += (loss_flat * mask).sum().item()
total_bytes += nbytes.sum().item()
return total_nats / (math.log(2) * total_bytes)
为什么用 bits-per-byte 而不是 perplexity 或交叉熵损失?
标准交叉熵损失依赖于词表大小。如果代理以某种方式改变模型架构导致输出分布变化,原始损失值就不具有可比性。BPB 将所有东西归一化到每字节基础上:
BPB = total_nats / (ln(2) * total_bytes)
这是信息论压缩率:模型每字节文本需要多少比特?越低越好。一个完美的英文文本压缩器大约能达到 1.0 BPB。这个指标与词表大小无关,所以代理可以自由地实验架构变更。
mask = nbytes > 0 这一行将特殊 token 排除在指标之外。 BOS、保留 token;它们的字节长度为 0,不应贡献到压缩测量中。
reduction='none' 是关键。 它不是在所有 token 上平均损失(那是 reduction='mean' 会做的),而是返回每 token 损失。这允许按字节长度加权求和,这在数学上不同于平均。
3、train.py,代理唯一修改的文件
那是基础。现在让我们看看真正的工作在哪里发生。
这个文件包含完整的 GPT 实现、混合 Muon+AdamW 优化器和训练循环。这不是普通的训练脚本。这里的一切都是代理可以重写的。
3.1 Flash Attention 3 内核加载
from kernels import get_kernel
cap = torch.cuda.get_device_capability()
# varunneal 的 FA3 仅支持 Hopper,在非 Hopper GPU 上使用 kernels-community
repo = "varunneal/flash-attention-3" if cap == (9, 0) else "kernels-community/flash-attn3"
fa3 = get_kernel(repo).flash_attn_interface
导入时的 GPU 能力检测。 CUDA 计算能力 (9, 0) 是 NVIDIA Hopper 架构(H100、H200)。代码根据你是否在 Hopper 上加载不同的 Flash Attention 3 内核。kernels 库动态从 GitHub 仓库获取并编译 CUDA 内核;这种模式保持代码库小巧,同时支持在不同硬件上的优化注意力计算。
3.2 GPT 架构:每一层可视化
配置
@dataclass
class GPTConfig:
sequence_len: int = 2048
vocab_size: int = 32768
n_layer: int = 12
n_head: int = 6
n_kv_head: int = 6
n_embd: int = 768
window_pattern: str = "SSSL"
window_pattern = "SSSL" 是滑动窗口注意力模式。S = 短窗口(一半上下文 = 1024),L = 长/全窗口(2048)。模式在层间重复:层 0-2 使用短窗口,层 3 使用全注意力,层 4-6 使用短窗口,层 7 使用全注意力,等等。最后一层无论模式如何都使用全注意力。 这确保模型能够在最终预测时关注完整上下文。
一行归一化。就这些。
def norm(x):
return F.rms_norm(x, (x.size(-1),))
一行。没有可学习参数(没有 gamma/beta)。RMSNorm 比 LayerNorm 更便宜,因为它跳过了均值减法步骤。对于这种无参数版本,它就是:x / sqrt(mean(x^2) + eps)。
值嵌入(ResFormer)
这是更奇特的组件之一:
def has_ve(layer_idx, n_layer):
"""返回该层是否应该有值嵌入(交替,最后一层总是包含)。"""
return layer_idx % 2 == (n_layer - 1) % 2
值嵌入是 ResFormer 论文中的一种技术。想法是:让注意力机制直接访问输入 token 嵌入作为额外的"值"信号,独立于通过残差流流动的上下文表示。
# 在 CausalSelfAttention.forward 中:
if ve is not None:
ve = ve.view(B, T, self.n_kv_head, self.head_dim)
gate = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
v = v + gate.unsqueeze(-1) * ve
门控机制很微妙。 2 * sigmoid(gate) 产生 [0, 2] 范围内的值。在初始化时,门权重为零,所以 sigmoid(0) = 0.5,2 * 0.5 = 1.0。一个中性的起点。模型可以学习抑制值嵌入(门→0)或放大它(门→2)。门只使用输入的前 32 个通道(x[..., :self.ve_gate_channels]),这是一个刻意的瓶颈以防止过拟合。
值嵌入只应用于交替的层(不是每一层),减少了参数数量同时仍然提供好处。
旋转位置嵌入(RoPE)
def apply_rotary_emb(x, cos, sin):
assert x.ndim == 4
d = x.shape[3] // 2
x1, x2 = x[..., :d], x[..., d:]
y1 = x1 * cos + x2 * sin
y2 = x1 * (-sin) + x2 * cos
return torch.cat([y1, y2], 3)
标准 RoPE 实现。关键洞察:RoPE 不是添加位置信息,而是以位置依赖的方式旋转 查询和键向量。两个距离 d 的 token 总是有相同的旋转差,无论绝对位置如何。这给模型提供了相对位置感知。
RoPE 后的 QK-Norm:
q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
q, k = norm(q), norm(k)
在应用旋转嵌入后归一化 Q 和 K 是一种稳定性技术。 它防止点积注意力分数变得过大,这会在 softmax 中引起数值问题,特别是在较长序列时。
为什么用 ReluSquared 而不是 GELU 或 SwiGLU?
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
def forward(self, x):
x = self.c_fc(x)
x = F.relu(x).square() # <-- ReluSquared
x = self.c_proj(x)
return x
问题是: relu(x).square() 不是显而易见的选择。ReluSquared(也称为 Squared ReLU)被 Primer 论文证明在语言建模方面与 GELU 竞争或更好,同时计算更便宜。它也更简单。没有像 SwiGLU 那样的门控机制,这意味着不需要额外的线性层,这意味着更少的参数和在固定时间预算内每步更快的计算。
扩展比率是 4x(标准),并且没有任何偏置。这是减少参数数量的另一个简化。
前向传播:残差 Lambda 和 Logit 软截断
def forward(self, idx, targets=None, reduction='mean'):
B, T = idx.size()
cos_sin = self.cos[:, :T], self.sin[:, :T]
x = self.transformer.wte(idx)
x = norm(x)
x0 = x # <-- 保存初始嵌入
for i, block in enumerate(self.transformer.h):
x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 # <-- 残差混合
ve = self.value_embeds[str(i)](idx) if str(i) in self.value_embeds else None
x = block(x, ve, cos_sin, self.window_sizes[i])
x = norm(x)
softcap = 15
logits = self.lm_head(x)
logits = logits.float()
logits = softcap * torch.tanh(logits / softcap) # <-- 软截断
残差 Lambda (resid_lambdas 和 x0_lambdas):不是标准的 x = x + block(x) 残差连接,而是使用可学习的逐层缩放:
x = resid_lambda[i] * x + x0_lambda[i] * x0
这受到 DeepNet 和相关工作的启发。x0 是初始嵌入;它在每一层提供一个"来自底部的跳跃连接"。模型可以学习平衡在每个深度使用多少累积的残差流与原始输入嵌入。
在初始化时:resid_lambdas = 1.0,x0_lambdas = 0.1。所以初始行为大致是 x = x + 0.1 * x0。主要是标准残差加上少量原始嵌入的注入。
Logit 软截断将幅度限制在 [-15, 15]****。 没有这个,logits 在训练过程中会变得非常大,导致交叉熵损失中的数值不稳定。Gemma 2 和其他模型使用这种技术。logits.float() 在软截断前转换为 float32 至关重要;bfloat16 中的 tanh 会损失太多精度。
3.3 权重初始化:为什么输出投影从零开始
@torch.no_grad()
def init_weights(self):
# 嵌入:宽初始化
torch.nn.init.normal_(self.transformer.wte.weight, mean=0.0, std=1.0)
# LM Head:微小初始化
torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
# Transformer 矩阵:均匀分布,按维度缩放
s = 3**0.5 * n_embd**-0.5
for block in self.transformer.h:
torch.nn.init.uniform_(block.attn.c_q.weight, -s, s)
torch.nn.init.uniform_(block.attn.c_k.weight, -s, s)
torch.nn.init.uniform_(block.attn.c_v.weight, -s, s)
torch.nn.init.zeros_(block.attn.c_proj.weight) # <-- 零初始化!
torch.nn.init.uniform_(block.mlp.c_fc.weight, -s, s)
torch.nn.init.zeros_(block.mlp.c_proj.weight) # <-- 零初始化!
输出投影的零初始化模式是这里的关键洞察。 在初始化时,注意力和 MLP 中的 c_proj 输出恰好为零。这意味着每个 transformer 块开始时是一个恒等函数;残差流无损通过。然后模型在训练过程中逐渐学习通过这些投影注入信息。这是深度 transformer 的一个著名稳定性技术。
缩放因子 s = sqrt(3) * n_embd^(-0.5) 来自均匀分布。如果你想要方差 1/n_embd,均匀边界应该是 ±sqrt(3/n_embd)。
3.4 两个优化器,一个模型:混合 MuonAdamW
但是等等。这里才是事情变得真正有趣的地方。代码库对不同类型的参数使用两个不同的优化器。不是因为有人无法决定。是因为数学要求如此:
为什么两个优化器?
Muon (来自 Muon 优化器论文)通过正交化梯度更新来工作。它使用"极分解快速算法"将动量投影到最近的正交矩阵上。这对于2D 权重矩阵(注意力投影、MLP 层)效果很好,因为正交性对矩阵有清晰的几何意义。
但是对于1D 参数(嵌入、标量),正交性没有意义。所以那些使用标准 AdamW。
极分解快速算法:无需 SVD 的正交化
polar_express_coeffs = [
(8.156554524902461, -22.48329292557795, 15.878769915207462),
(4.042929935166739, -2.808917465908714, 0.5000178451051316),
# ... 更多 3 组系数
]
# 在 muon_step_fused 内部:
X = g.bfloat16()
X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
if g.size(-2) > g.size(-1):
for a, b, c in polar_express_coeffs[:ns_steps]:
A = X.mT @ X
B = b * A + c * (A @ A)
X = a * X + X @ B
else:
for a, b, c in polar_express_coeffs[:ns_steps]:
A = X @ X.mT
B = b * A + c * (A @ A)
X = a * X + B @ X
这里发生了什么?矩阵 M 的极分解给出 M = UP,其中 U 是正交的,P 是半正定的。找到 U 通常很昂贵(需要 SVD)。"极分解快速算法"是一种多项式近似,通过迭代收敛到正交因子。
每次迭代应用:X = a*X + X @ (b * X^T X + c * (X^T X)^2)。这是 Padé 型有理近似到矩阵符号函数,收敛到极因子。
ns_steps=5 参数控制运行多少次迭代。更多迭代 = 更接近真正的正交投影,但更多计算。5 步是最佳点。
高矩阵 vs 宽矩阵的分支 (if g.size(-2) > g.size(-1))是一种计算优化。对于高矩阵,计算 X^T X(小方阵)更便宜。对于宽矩阵,X X^T 更便宜。
NorMuon:Muon 的类似 Adam 的二阶矩
正交化后,Muon 步骤应用 NorMuon 方差归一化:
v_mean = g.float().square().mean(dim=red_dim, keepdim=True)
second_momentum_buffer.lerp_(v_mean, 1 - beta2)
step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt()
这类似于 Adam 中的二阶矩。 它估计梯度的方差并相应地缩放更新。这防止任何单个维度主导更新。
谨慎权重衰减
mask = (g * stacked_params) >= 0
stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)
这不是标准权重衰减。 标准权重衰减无条件应用 p = p - lr * wd * p。谨慎权重衰减只在梯度和参数指向相同方向时 (g * p >= 0)应用衰减。直觉是:如果梯度想把参数推得更远离零(同号),那么权重衰减应该抵抗这一点。但如果梯度想缩小参数(异号),权重衰减就不应该与梯度对抗;它们已经一致了。
3.5 调度基于壁钟时间,而不是步数
def get_lr_multiplier(progress):
if progress < WARMUP_RATIO:
return progress / WARMUP_RATIO if WARMUP_RATIO > 0 else 1.0
elif progress < 1.0 - WARMDOWN_RATIO:
return 1.0
else:
cooldown = (1.0 - progress) / WARMDOWN_RATIO
return cooldown * 1.0 + (1 - cooldown) * FINAL_LR_FRAC
注意: WARMUP_RATIO = 0.0****。 没有预热!模型立即以全学习率开始。训练的整个后半段(WARMDOWN_RATIO = 0.5)是到 FINAL_LR_FRAC = 0.0 的线性冷却。
关键是,progress = total_training_time / TIME_BUDGET。调度基于壁钟时间,而不是步数。这意味着调度自动适应代理架构选择产生的任何步吞吐量。
3.6 训练循环(以及每次暂停节省 500ms 的 GC 技巧)
t_start_training = time.time()
step = 0
while True:
torch.cuda.synchronize()
t0 = time.time()
for micro_step in range(grad_accum_steps):
with autocast_ctx:
loss = model(x, y)
train_loss = loss.detach()
loss = loss / grad_accum_steps
loss.backward()
x, y, epoch = next(train_loader)
# 进度和调度
progress = min(total_training_time / TIME_BUDGET, 1.0)
# ... 更新学习率 ...
optimizer.step()
model.zero_grad(set_to_none=True)
梯度累积 允许大的有效批量大小(TOTAL_BATCH_SIZE = 2^19 ≈ 524K tokens),即使设备只能一次容纳 DEVICE_BATCH_SIZE = 128 个序列。计算:grad_accum_steps = 524288 / (128 * 2048) = 2。
GC 管理是一个生产技巧:
if step == 0:
gc.collect()
gc.freeze()
gc.disable()
elif (step + 1) % 5000 == 0:
gc.collect()
Python 的垃圾收集器在运行时会导致约 500ms 的暂停。在 5 分钟的训练运行中,每步约 300ms,损失 500ms 给 GC 是很显著的。解决方案:第一步之后,收集垃圾,冻结所有现有对象(这样 GC 就忽略它们),然后完全禁用自动 GC。只每 5000 步手动收集一次。
预热步跳过:
if step > 10:
total_training_time += dt
if step > 10 and total_training_time >= TIME_BUDGET:
break
前 10 步被排除在时间预算之外。这考虑了 PyTorch 编译时间(当 torch.compile 首次跟踪模型时)、CUDA 内核启动和其他一次性预热成本。没有这个,不同的模型架构会因为编译成本不同而不公平地获得不同的实际训练时间。
4、program.md,取代研究团队的 114 行
现在。这是仓库中最迷人的文件。它不是代码;它是一组将通用 AI 代理变成专注机器学习研究者的指令。
4.1 实验循环
program.md 中的关键约束:
代理可以做什么: 修改 train.py。一切都是公平游戏:架构、优化器、超参数、批量大小、模型大小。
代理不能做什么:
- 修改
prepare.py(评估是神圣不可侵犯的) - 安装新包
- 修改评估框架
简洁性标准特别有趣:
在其他条件相同时,更简单更好。增加丑陋复杂性的小改进不值得。相反,移除某些东西并获得相同或更好的结果是一个很好的结果。
这反映了优秀人类研究者的思维方式。代理被明确告知在指标改进的同时重视代码简洁性。
"永不停止"指令:
一旦实验循环开始,不要停下来问人类是否应该继续......循环一直运行到人类中断你为止。
这就是让 autoresearch 成为自主 的原因。代理每小时运行约 12 个实验,一晚上约 100 个。人类醒来时会看到一个充满实验日志的 results.tsv。
4.2 结果跟踪格式
commit val_bpb memory_gb status description
a1b2c3d 0.997900 44.0 keep baseline
b2c3d4e 0.993200 44.2 keep increase LR to 0.04
c3d4e5f 1.005000 44.0 discard switch to GeLU activation
d4e5f6g 0.000000 0.0 crash double model width (OOM)
简单。制表符分隔。每个实验一行。commit 列创建了每个实验的完整 git 历史;你可以 git checkout 任何 commit 来查看代理尝试了什么。
5、analysis.ipynb,你醒来时看到的东西
分析 notebook 是你醒来时查看的东西。它读取 results.tsv 并生成可视化:
notebook 计算:
- 保留率:改进了指标的实验比例
- 运行最小值:随时间推移的最佳结果"前沿"
- 每实验增量:每个保留的实验相对于之前最佳的改进程度
- 最佳命中排名:哪些变化产生了最大的改进
6、结束语
我们读了每个函数。让我们退一步看大局。
人类一直是瓶颈
传统机器学习研究中的瓶颈一直是人类。睡觉、吃饭、在项目间上下文切换、忘记运行下一个实验。Autoresearch 完全将人类从内循环中移除。
在一次夜间运行(约 8 小时)中,系统可以执行约 100 个实验。人类研究者做同样的工作(想出想法、实现它、运行训练、分析结果)一天可能只能完成 3-5 个实验,还要休息。我们是认真的吗?
代码中编码的设计原则
- 固定时间预算,而不是固定步数。 这使优化目标变为"你的硬件在 5 分钟内能产生的最佳模型"。不同的硬件自然找到不同的最优解。
- 单文件修改范围。 代理不能破坏评估。它不能损坏数据管道。它只能修改模型和训练代码。这是一个沙盒。
- Git 作为实验日志。 每个实验都是一个 commit。每次回退都是
git reset。完整历史被保留。这正是版本控制被设计来做的事。 - BPB 作为北极星指标。 与词表大小无关,具有信息论基础,不可能在不实际改进模型的情况下作弊。
- 人类写
program.md****,而不是 Python。 元洞察:
最重要的要优化的代码是代理的指令,而不是训练代码本身。人类的工作从研究者转变为研究经理。
真正的实验是 program.md,不是 train.py
让我们回顾 Karpathy 在 README 中的开场白
曾经,前沿 AI 研究是由"肉计算机"完成的......研究现在完全是自主 AI 代理集群在天空中的计算集群巨型结构上运行的领域。
代码本身是适度的。三个文件总共约 1000 行。但它建立的架构**(不可变评估 + 可变模型 + 代理指令 + 自主循环)是一个模板**。扩大计算规模,扩大代理能力,添加更多代理,你会得到越来越像自主研究实验室的东西。
最能说明问题的细节可能是最小的:program.md 被描述为"超级轻量级技能"。它是 114 行 Markdown。这就是自主研究组织的整个管理层。那些迭代这个文件的人,寻找更好的指令、更好的约束、更好的代理奖励信号;他们正在做一种新的工作。他们不是在写神经网络代码。他们在写规范给一个写神经网络代码的实体。