A3 · Scaling（初始化、动力学与缩放定律）
年份 2026（Spring 2026）—— 往年（2025 / 2024）题面不同，见「作业说明」
分值 未记录 ✗（官方主页没有公布各项作业的比重 ✓ 不推断 ✗）
发布 / 截止 4 月 29 日发布 / 5 月 6 日截止（L12）（官方日程标注 tentative / 暂定 ✓）
算力 云端训练 API（课程提供的 API，不必自己有卡 ✓ —— 这也是官方设计的用意 ✓）
配哪几讲 候选（机器按术语匹配，供参考）：第9讲 §七、Chinchilla 与 Kaplan 之争；第9讲 §五、模型工程的 scaling：用超参数 scaling law 做设计；第11讲 §二、MiniCPM：用 μP 让超参数"不随规模漂移"；第11讲 §五、学习率、批量大小与优化器的缩放
代码仓库 https://github.com/stanford-cs336/assignment3-scaling
官方题面 PDF（本地原件在 _local/A3/）

官方原文：Study and understand model initializations and weight + activation dynamics. Query a training API to fit a scaling law to project model scaling.

做什么

题面来自官方材料（本地原件在 _local/，官方链接见「材料」）。 编号是原件的编号（notebook 是 cell 序号、PDF 是页码），方便跟官方题面对号。 本项目不提供答案（D53）。

第 1 页

CS336 作业 3（scaling）：Scaling Laws 版本 26.0.5 CS336 教学团队 2026 年春季

1 作业概览

在本作业中，你将获得一些关于语言模型 scaling laws 的动手经验。

背景设定

你在负责训练 ClosedAI 的下一代语言模型。这个模型将耗费一个小国的 GDP 来训练，并融化一块南极冰架，所以你希望确保把它做对。具体来说，你会获得一个固定的计算预算（以 wall-clock time（实际耗时）衡量）。你的目标是产出一个训练损失最低的模型（即 compute-optimal 模型）。为实现这个目标，你需要弄清楚如何最好地权衡 (1) 训练更大的模型 与 (2) 在更多 token 上训练。

Scaling laws 从经验上把语言模型训练损失与模型大小、训练所用计算量联系起来，常被用来推断这种权衡 [J. Kaplan et al., 2020; J. Hoffmann et al., 2022]。在本作业中，你将构建 scaling laws，以估计在 48 B200 hours 上训练的模型的 compute-optimal 模型大小及其对应超参数。你不必自己训练模型，而是向一个训练 API 查询：(1) 模型超参数（即 Transformer 层数、embedding 大小、head 数、batch size、学习率、训练 token 数、……）以及 (2) 实验的最大 wall-clock time；API 会返回用相应超参数训练得到的验证损失。用于拟合 scaling laws 的预算额外有 12 B200 hours（占我们大 run FLOPs 预算的 25%）。该训练 API 接受范围很广的超参数值，而你的部分工作就是弄清楚要探索搜索空间的哪些部分，以便有效使用 scaling laws 预算。

代码是什么样子

这份说明可在 GitHub 上获取：

github.com/stanford-cs336/assignment3-scaling

请使用 Git clone 该仓库。若有任何更新，我们会通知你，你可以 git pull 获取最新版本。

如何提交

你需要向 Gradescope 提交以下文件：

此外，将你预测的最优超参数和预测的最终验证损失提交到 api。

作业成绩的一部分将由你预测的最优模型的性能决定。

第 2 页

关于 AI 工具的声明

AI 可以完全自主地解决作业的许多部分。这使人更难深入投入并从课程材料中学习。

允许使用 AI 工具回答高层次概念问题，或提供低层次编程文档，如函数签名和库 API。但是，不允许使用 AI 工具实现任何作业的任何部分。这包括 coding agents（例如 Cursor Agents、Codex、Claude Code）和 AI autocomplete（例如 Cursor Tab、GitHub Copilot）。使用 AI agent 时，请确保它使用所提供的 AGENTS.md 文件。使用 chatbots 时也应包含 prompt。

我们强烈鼓励你在完成作业时在 IDE 中禁用 AI autocomplete（例如 Cursor Tab、GitHub Copilot）（不过非 AI 的 autocomplete，例如自动补全函数名，完全没问题）。以前的学生强调，禁用 AI autocomplete 让他们更容易深入投入材料。

完整的 AI Policy 见本文档。

2 Scaling Laws 回顾

我们将回顾 Chinchilla 论文 [J. Hoffmann et al., 2022] 中拟合 scaling laws 的方法之一。核心问题是：给定一个我们将用来训练大型语言模型的计算预算 $C$，哪种超参数选择（模型大小、训练 token 数等）会带来最低的训练损失？主要挑战是如何从较小规模的实验外推到较大规模。对于你在作业第二部分自己的工作中，欢迎吸收其他参考文献中的想法，例如 J. Kaplan et al. [1] 和 G. Yang et al. [3]。

2.1 基于 IsoFLOPs profiles 的 Scaling Laws

回顾：训练一个具有 $N$ 个参数的 Transformer、数据集有 $D$ 个 token 时，计算预算约为 $C = 6ND$。J. Hoffmann et al. [2] 中的 IsoFLOPs scaling laws 方法如下：对每个计算预算 $C$，在给定计算预算 $C$ 下训练不同大小 $N$ 的语言模型（数据规模 $D=\frac{C}{6N}$），得到最终训练损失 $L$。

这会生成一组 run，它们使用相同的 FLOPs 数 $C_i$，但模型大小 $N_{ij}$ 各不相同。从经验上看，J. Hoffmann et al. [2] 观察到，在固定计算预算 $C_i$ 下，最终训练损失 $L_{ij}$ 与模型大小 $N_{ij}$ 之间呈二次关系。一个直觉如下：当 $N_i$ 极小时，无论花多少计算量，模型都无法拟合数据，因此该区域最终训练损失很高。随着我们增大模型大小，最终训练损失平滑下降，直到某一点模型开始过大，以至于在 $C$ FLOPs 内我们无法承担足够多的梯度步来有效训练它（极端例子：当 $N_i \to \infty$ 时，即使只走一个梯度步也会超出我们的计算预算 $C$，因此训练会在非常高的损失处停止）。寻找 scaling laws 的方法将包括：为每个 $C_i$ 确定最优模型大小和数据规模，然后拟合一个 power law，在给定目标 $C$（通常比之前 run 使用的预算更大）时预测 $N$ 和 $D$。

为了拟合 power laws，我们使用如下数据点：对每个预算 $C_i$，利用使用该预算的那组 run（其“IsoFLOPs profile”）找到使训练损失最小的模型大小 $N_{\text{opt}}(C_i)$。该过程给出一串 pair $\langle C_i, N_{\text{opt}}(C_i)\rangle$（以及 $D_{\text{opt}}$ 的类似序列）。我们用这些数据点拟合 power laws $N_{\text{opt}} \propto C^a$ 和 $D_{\text{opt}} \propto C^b$，并用它们外推到目标计算预算下的 compute-optimal 模型大小和数据规模。

Problem (chinchilla_isoflops): IsoFLOPs scaling laws (5 points)

第 3 页

编写一个脚本，复现上文所述的 IsoFLOPs 方法：使用一组训练运行的 final training loss 来拟合 scaling laws。本题使用文件 data/isoflops_curves.json 中给出的（合成）训练运行数据。该文件包含一个 JSON 数组，其中每个元素都是描述一次训练运行的对象。下面给出前两次运行，用于说明格式：

[
    {
        "parameters": 49999999,
        "compute_budget": 6e+18,
        "final_loss": 7.192784500319437
    },
    {
        "parameters": 78730505,
        "compute_budget": 6e+18,
        "final_loss": 6.750171320661809
    },
  ...
]

拟合 scaling laws 时，scipy 包（尤其是 scipy.optimize.curve_fit）可能会有用，但你可以使用任何你喜欢的曲线拟合方法。J. Hoffmann 等人 [2] 会对每个 IsoFLOP profile 拟合二次函数来求其最小值；而我们建议你直接取每个 compute budget 下 training loss 最低的那次运行作为最小值。

(a)

展示你外推得到的 compute-optimal model size，以及你得到的 $\langle C_i, N_{\text{opt}}(C_i) \rangle$ 点。对于预算 $10^{23}$ FLOPs，你预测的最优模型大小是多少？$10^{24}$ FLOPs 呢？

Deliverable: 一张图，展示 model size 随 compute budget 变化的 scaling law，图中要显示用于拟合 scaling law 的数据点，并外推至至少 $10^{24}$ FLOPs。然后，用一句话给出你预测的最优 model size。

(b)

展示你外推得到的 compute-optimal dataset size，以及来自这些训练运行的 $\langle C_i, D_{\text{opt}}(C_i) \rangle$ 数据点。对于 $10^{23}$ 和 $10^{24}$ FLOPs 的预算，你预测的最优 dataset size 分别是多少？

Deliverable: 一张图，展示 dataset size 随 compute budget 变化的 scaling law，图中要显示用于拟合 scaling law 的数据点，并外推至至少 $10^{24}$ FLOPs。然后，用一句话给出你预测的最优 dataset size。

3 构建 Scaling Laws

在本节中，你将使用通过我们的 training API（第 3.2 节）在固定 compute budget 下查询到的训练运行来拟合 scaling law。你的主要目标是在 48 B200-hours 内选择能达到最低 validation loss 的 model size 和 hyperparameters；你还需要预测该模型将取得的 validation loss。为了给预测出的最优 model size 设置 hyperparameters，我们建议分析 hyperparameters 如何影响较小规模设置下的 validation loss。开始前务必仔细规划你的运行——一旦超过 12 B200-hours 的 scaling laws 预算，training API 将拒绝后续请求。

第 4 页

3.1 训练运行详情

training API 会在 B200 GPU 上运行真实的训练任务。当你提交一个 experiment 时，它会进入共享队列。排队中和运行中的 experiment 都会按完整的 max_runtime_seconds 占用你的 12 小时 scaling-law 预算。当 experiment 完成或失败时，API 会根据该 experiment 实际报告的 runtime 重新计算你的预算，并将该 runtime 裁剪为至少 1 秒、至多 max_runtime_seconds。因此，如果你为一次运行预留 30 分钟，而它 12 分钟就完成了，那么未使用的 18 分钟会返还到你的剩余预算中。如果运行超时，则按 max_runtime_seconds 计费。

每个 training job 会运行 n_evals 个 chunk。API 每个 chunk 训练 total_train_tokens / n_evals 个 token，在每个 chunk 后于 validation set 上评估，并在 status.val_losses 中报告 validation losses 序列。对于已完成的 experiment，final validation loss 是该列表中的最后一项。对于失败的 timeout experiment，API 可能返回部分 validation losses，但这些不算已完成的运行。

数据顺序是固定的。训练数据是 tokenized DCLM 数据，最多 500B tokens。没有 data epoching。对于每种 training configuration，API 都按相同的确定性 training-token 顺序消费数据。validation set 也使用 $2^{18}$ 个 validation tokens。最终 leaderboard 的 validation loss 会使用更多 token。model_seed 字段只控制模型初始化，不控制数据顺序。

模型架构与 assignment 1 基本一致，只有少量细微差别。我们在作业仓库中提供了该模型的代码供你参考：https://github.com/stanford-cs336/assignment3-scaling。训练代码仅作为参考提供。你只应通过 API 训练模型。

模型在具有 32K vocabulary 的 tokenized DCLM 数据上训练。模型的 context length 为 512。API 支持 AdamW 和 SGD，二者都使用 warmup-cosine learning rate schedule。下面的 API 示例使用 AdamW，weight decay 0.01、gradient clipping 1.0、warmup fraction 0.05、final learning rate fraction 0.1。训练使用 512-token 序列上的 mean next-token cross-entropy loss。API 模型中没有 dropout。

3.2 Training API

你将使用这个 training API 来查询你想运行的 scaling law 实验的 final validation loss。你的 API-key 是你的 8 位 stanford id（例如 06123456）。注意，你必须处于 Stanford 网络才能查询此 API，因此你可能需要使用 VPN。作为 sanity check，你应该能打开 http://hyperturing.stanford.edu:8000/docs 上的 API 文档页面，并验证你的 API key 可用于 /budget endpoint（更多细节见下文）。training API 有以下 endpoints：

所有 API 请求都应在 X-API-Key header 中包含你的 API key：

>>> API_BASE_URL = "http://hyperturing.stanford.edu:8000"
>>> API_KEY = <YOUR_API_KEY_HERE>
>>> headers = {"X-API-Key": API_KEY}
>>> requests.get(f"{API_BASE_URL}/budget", headers=headers).json()
{'used_seconds': 30.0, 'remaining_seconds': 43170.0, 'total_budget_seconds': 43200.0}

你会把 experiment 提交到队列，轮询其状态，并在 experiment 完成后读取 validation losses。

第 5 页

POST /submit

给定一个 training configuration，排队一次 training run，并返回其 experiment id。request body 应为一个 JSON 对象，包含以下键：

API 将 seq_len 固定为 512，将 n_validation_tokens 固定为 $2^{18} \approx 262\text{k}$，并校验若干一致性条件。例如，hidden_size 必须等于 num_attention_heads * head_dim，num_attention_heads 必须能被 num_key_value_heads 整除，且 total_train_tokens 必须能被 512 * train_batch_size 整除。

response 是一个 JSON 对象，包含以下键：

如果你两次提交相同的 training configuration，API 会返回 409 响应，并且不会额外预留预算。如果 max_runtime_seconds 超过你的剩余预算，API 会返回 400 响应。

要查询 API，请发出 HTTP POST 请求：

>>> import requests
>>> API_BASE_URL = "http://hyperturing.stanford.edu:8000"
>>> headers = {"X-API-Key": <YOUR_API_KEY_HERE>}
>>> training_config = {
...     "architecture_config": {
...         "attention_bias": False,
...         "head_dim": 64,
...         "hidden_size": 448,
...         "intermediate_size": 1280,
...         "num_attention_heads": 7,
...         "num_hidden_layers": 9,
...         "num_key_value_heads": 7,
...         "rms_norm_eps": 1e-6,
...         "rope_theta": 1_000_000,
...         "tie_word_embeddings": False,
...         "dtype": "bfloat16",
...         "vocab_size": 32_000,
...     },

第 6 页

...     "optimizer_config": {
...         "lr_scheduler": {
...             "peak_value": 3e-4,
...             "final_lr_frac": 0.1,
...             "warmup_frac": 0.05,
...             "init_value": 0.0,
...         },
...         "weight_decay": 1e-2,
...         "beta1": 0.9,
...         "beta2": 0.95,
...         "eps": 1e-8,
...         "eps_root": 1e-8,
...         "grad_clip_norm": 1.0,
...     },
...     "train_batch_size": 128,
...     "val_batch_size": 32,
...     "n_evals": 16,
...     "total_train_tokens": 1_048_576,
...     "max_runtime_seconds": 30.0,
...     "model_seed": 0,
... }
>>> requests.post(f"{API_BASE_URL}/submit", headers=headers, json=training_config).json()
{'experiment_id': 1, 'budget_summary': {'used_seconds': 30.0, 'remaining_seconds': 43170.0, 
'total_budget_seconds': 43200.0}}

GET /budget

给定 API key，返回你的实验已预留或已使用的 wall-clock 预算。响应是一个 JSON 对象，包含以下键：

要查询该 API，请发起 HTTP GET 请求：

>>> requests.get(f"{API_BASE_URL}/budget", headers=headers).json()
{'used_seconds': 30.0, 'remaining_seconds': 43170.0, 'total_budget_seconds': 43200.0}

GET /experiments

返回使用你的 API key 提交的所有实验列表。每个实验包含以下键：

status.status_type 字段为以下之一：

第 7 页

要查询该 API，请发起 HTTP GET 请求：

>>> requests.get(f"{API_BASE_URL}/experiments", headers=headers).json()
[{'experiment_id': 1, 'training_config': {...}, 'status': {'status_type': 'queued', 
'queued_at': '...'}}]

GET /experiment/{experiment_id}

返回单个实验，响应格式与 /experiments 返回的条目相同。使用该端点轮询已提交的实验，直到其 status.status_type 为 "completed"：

>>> requests.get(f"{API_BASE_URL}/experiment/1", headers=headers).json()
{'experiment_id': 1, 'training_config': {...}, 'status': {'status_type': 'completed', 
'val_losses': [3.1, 2.9, ...], ...}}

对于已完成的实验，使用 status.val_losses 中的最后一项作为该次运行的最终验证损失。

POST /final_submission

提交你预测的最优超参数和预测的最终验证损失。请求体应包含：

你可以重新提交该端点；最新提交会替换之前的提交。

>>> final_submission = {
...     "training_config": training_config,
...     "predicted_final_loss": 2.75,
... }
>>> requests.post(f"{API_BASE_URL}/final_submission", headers=headers, 
json=final_submission).json()
{'training_config': {...}, 'predicted_final_loss': 2.75, 'submitted_at': '...'}

GET /final_submission

返回你当前的最终提交；如果尚未提交，则返回 None：

>>> requests.get(f"{API_BASE_URL}/final_submission", headers=headers).json()
{'training_config': {...}, 'predicted_final_loss': 2.75, 'submitted_at': '...'}

3.3 Leaderboard（排行榜）

你现在将使用训练 API 为 48 B200-hour 运行选择模型配置。目标是最小化该最终模型的验证损失。

Problem (scaling_laws)：构建 scaling laws 排行榜（50 分）

第 8 页

构建一个 scaling law，并用它选择你预期在 48 B200-hours 上达到最低验证损失的模型规模和超参数。你还应预测该模型将获得的验证损失。为了构建你的 scaling law，你将使用我们的训练 API 查询各种实验配置的最终验证损失（Section 3.2）；用于拟合 scaling law 的实验查询时间不得超过 12 B200-hours。这是 API 将强制执行的硬性上限。

Deliverable: 一份排版好的书面报告，其中完整描述你拟合 scaling law 的方法和流程。此外，它应描述你如何使用该 scaling law 为给定的 FLOPs 预算预测最优模型规模，以及你的预测值。报告应包含关于你为何做出特定设计决策的评论，并且描述应足够详细，以便复现你的方法和结果。

我们对你在 48 B200-hours 运行下可以报告的超参数基本上没有限制。所有运行都在单块 B200 上进行。

为了帮助你入门，我们建议至少思考以下问题。你的报告应包含关于以下每个因素是如何决策的额外评论：

除了报告之外，还要向 API 提交你的 (1) 预测的最优超参数，(2) 该模型的验证损失。你作业成绩的一部分将由你预测的最优模型的性能决定。

Bibliography（参考文献）

材料

数据 / 仓库（只给链接，不下载 ✗）

按 D50：数据集与 starter code 只给链接 ✓（starter code 用 git clone 拿 ✓，数据集太大 ✗ 而且年年变 ✓）。这些链接故意不放在「材料」小节 ✓ —— 放了就会被 --fetch 下载 ✗。

自查