Hugging Face's logo

Spaces:
dongyanggg
/
codestar-backend
Running on Zero

App Files Community
codestar-backend / data /problems /adamw-optimizer-step.json
github-actions
Sync backend from GitHub 47e8067
5f7afef
Raw History Blame Contribute Delete
14.7 kB
{
"slug": "adamw-optimizer-step",
"title": "手撕 AdamW 解耦权重衰减单步 (AdamW Step)",
"difficulty": "Medium",
"category": "Optim",
"tags": [
"常见手撕",
"AdamW",
"权重衰减",
"Transformer"
],
"description_md": "### 题目描述\n将 L2 权重衰减直接作用于参数本身，解耦动量估计：$\\displaystyle \\theta_t = (1 - \\eta \\lambda) \\theta_{t-1} - \\eta \\frac{\\hat{m}_t}{\\sqrt{\\hat{v}_t} + \\epsilon}$。\n\n### 函数签名\n```python\ndef adamw_step(\n param: torch.Tensor,\n grad: torch.Tensor,\n m: torch.Tensor,\n v: torch.Tensor,\n t: int,\n lr: float = 1e-3,\n weight_decay: float = 0.01,\n beta1: float = 0.9,\n beta2: float = 0.999,\n eps: float = 1e-8,\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n```\n\n### 示例\n```python\n# 输入\nparam = [1.0]\ngrad = [0.1]\nm = [0.0]\nv = [0.0]\nt = 1\nlr = 0.001\nweight_decay = 0.01\nbeta1 = 0.9\nbeta2 = 0.999\neps = 1e-8\n\n# 输出\n(\n [0.999], # shape=[1], dtype=float32\n [0.01], # shape=[1], dtype=float32\n [0.0] # shape=[1], dtype=float32\n)\n```\n\n### 参数规格与约束\n- **`param`**：`torch.Tensor`，测试规格 `shape=[1], dtype=float32`\n- **`grad`**：`torch.Tensor`，测试规格 `shape=[1], dtype=float32`\n- **`m`**：`torch.Tensor`，测试规格 `shape=[1], dtype=float32`\n- **`v`**：`torch.Tensor`，测试规格 `shape=[1], dtype=float32`\n- **`t`**：`int`\n- **`lr`**：`float`（默认 `1e-3`）\n- **`weight_decay`**：`float`（默认 `0.01`）\n- **`beta1`**：`float`（默认 `0.9`）\n- **`beta2`**：`float`（默认 `0.999`）\n- **`eps`**：`float`（默认 `1e-8`）\n- **返回值**：`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`，输出规格 `Tuple` (`torch.Tensor(shape=[1], dtype=float32)`, `torch.Tensor(shape=[1], dtype=float32)`, `torch.Tensor(shape=[1], dtype=float32)`)\n- **精度容差**：`atol=0.0001, rtol=0.0001`",
"entry_point": "adamw_step",
"starter_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import Optional, List, Tuple\n\ndef adamw_step(\n param: torch.Tensor,\n grad: torch.Tensor,\n m: torch.Tensor,\n v: torch.Tensor,\n t: int,\n lr: float = 1e-3,\n weight_decay: float = 0.01,\n beta1: float = 0.9,\n beta2: float = 0.999,\n eps: float = 1e-8,\n) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"\n 手撕 AdamW 解耦权重衰减单步 (AdamW Step)\n \"\"\"\n # 请在此编写代码\n pass\n",
"reference_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import Optional, List, Tuple\n\ndef adamw_step(param: torch.Tensor, grad: torch.Tensor, m: torch.Tensor, v: torch.Tensor, t: int, lr: float = 1e-3, weight_decay: float = 0.01, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n new_m = beta1 * m + (1.0 - beta1) * grad\n new_v = beta2 * v + (1.0 - beta2) * (grad ** 2)\n m_hat = new_m / (1.0 - beta1 ** t)\n v_hat = new_v / (1.0 - beta2 ** t)\n new_p = param * (1.0 - lr * weight_decay) - lr * m_hat / (torch.sqrt(v_hat) + eps)\n return new_p, new_m, new_v\n",
"forbidden_modules": [
"os",
"sys",
"subprocess",
"shutil",
"socket",
"requests"
],
"forbidden_calls": [],
"time_limit": 3.0,
"memory_limit_mb": 512,
"atol": 0.0001,
"rtol": 0.0001,
"test_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\n\ncases = [\n ('AdamW单步', torch.tensor([1.0], dtype=torch.float32), torch.tensor([0.1], dtype=torch.float32), torch.tensor([0.0], dtype=torch.float32), torch.tensor([0.0], dtype=torch.float32), 1, 0.001, 0.01),\n ('多批次扩展泛化验证', torch.tensor([1.0, 1.0], dtype=torch.float32), torch.tensor([0.1, 0.15], dtype=torch.float32), torch.tensor([0.0, 0.0], dtype=torch.float32), torch.tensor([0.0, 0.0], dtype=torch.float32), 1, 0.001, 0.01),\n ('多维批次与参数泛化验证', torch.tensor([1.0], dtype=torch.float32), torch.tensor([0.15], dtype=torch.float32), torch.tensor([0.0], dtype=torch.float32), torch.tensor([0.0], dtype=torch.float32), 1, 0.0008, 0.008),\n ('边界极值与稳定性验证', torch.tensor([1.0], dtype=torch.float32), torch.tensor([0.05], dtype=torch.float32), torch.tensor([0.0], dtype=torch.float32), torch.tensor([0.0], dtype=torch.float32), 1, 0.0012, 0.012),\n]\nfor idx, (name, *args) in enumerate(cases, 1):\n t0 = time.time()\n try:\n exp = ref_obj(*args)\n act = user_obj(*args)\n match, err, max_diff = compare_outputs(act, exp, atol=atol, rtol=rtol)\n passed, msg = match, err\n except Exception as e:\n passed, msg = False, f\"Runtime Exception: {str(e)}\"\n act, exp, max_diff = None, None, None\n cases_results.append({\n \"case_id\": idx, \"case_name\": name, \"passed\": passed,\n \"time_ms\": round((time.time() - t0) * 1000, 2), \"message\": msg,\n \"user_output_summary\": summarize_obj(act) if act is not None else None,\n \"expected_summary\": summarize_obj(exp) if exp is not None else None,\n \"max_diff\": max_diff,\n })\n",
"is_featured": true,
"editorial_md": "#### 💡 核心算法原理与数学推导\n**手撕 AdamW 解耦权重衰减单步 (AdamW Step)**（`adamw_step`）聚焦于 **Optim** 中的关键计算逻辑：\n\n将 L2 权重衰减直接作用于参数本身，解耦动量估计：$\\displaystyle \\theta_t = (1 - \\eta \\lambda) \\theta_{t-1} - \\eta \\frac{\\hat{m}_t}{\\sqrt{\\hat{v}_t} + \\epsilon}$。\n\n在实现中，函数接收参数 `param` (`torch.Tensor`)、`grad` (`torch.Tensor`)、`m` (`torch.Tensor`)、`v` (`torch.Tensor`)、`t` (`int`)、`lr` (`float`)、`weight_decay` (`float`)、`beta1` (`float`)、`beta2` (`float`)、`eps` (`float`)，目标是计算出类型与形状严格符合 `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`param` 规格为 `torch.Tensor(shape=[1], dtype=float32)`，`grad` 规格为 `torch.Tensor(shape=[1], dtype=float32)`，`m` 规格为 `torch.Tensor(shape=[1], dtype=float32)`，`v` 规格为 `torch.Tensor(shape=[1], dtype=float32)`，`t` 规格为 `int` (示例值: `1`)，`lr` 规格为 `float` (示例值: `0.001`)，`weight_decay` 规格为 `float` (示例值: `0.01`)。\n- **期望输出规格**：`Tuple` (`torch.Tensor(shape=[1], dtype=float32)`, `torch.Tensor(shape=[1], dtype=float32)`, `torch.Tensor(shape=[1], dtype=float32)`)。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `new_m`：通过表达式 `beta1 * m + (1.0 - beta1) * grad` 完成张量或数值变换。\n2. 计算中间量 `new_v`：通过表达式 `beta2 * v + (1.0 - beta2) * (grad ** 2)` 完成张量或数值变换。\n3. 计算中间量 `m_hat`：通过表达式 `new_m / (1.0 - beta1 ** t)` 完成张量或数值变换。\n4. 计算中间量 `v_hat`：通过表达式 `new_v / (1.0 - beta2 ** t)` 完成张量或数值变换。\n5. 计算中间量 `new_p`：通过表达式 `param * (1.0 - lr * weight_decay) - lr * m_hat / (torch.sqrt(v_hat) + eps)` 完成张量或数值变换。\n6. 返回最终结果：输出 `new_p, new_m, new_v`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(N)$（线性于输入张量总元素个数 $\\prod d_i$），所有运算均为向量化逐元素算子或沿指定维度的归约计算。\n- **空间复杂度**：$O(N)$，需要与输入或输出张量同规模的显存空间保存计算结果。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，必须按公式位置准确加上 `eps`（例如在平方根内 `sqrt(var + eps)` 还是分母外 `+ eps`），防止全零或常量输入触发 `NaN` / `Inf`。\n- **边界输入鲁棒性**：注意单样本（`Batch=1`）、单步序列（`Seq=1`）、零值输入或极端超参数下的计算行为，不要在代码中硬编码特定测试用例的维度常数。\n- **纯向量化算子优先**：尽量使用 PyTorch 原生向量化张量算子替代低效的 Python 逐元素循环，以兼顾代码简洁性与执行性能。\n\n#### 💻 官方标准参考实现\n```python\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import Optional, List, Tuple\n\ndef adamw_step(param: torch.Tensor, grad: torch.Tensor, m: torch.Tensor, v: torch.Tensor, t: int, lr: float = 1e-3, weight_decay: float = 0.01, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n new_m = beta1 * m + (1.0 - beta1) * grad\n new_v = beta2 * v + (1.0 - beta2) * (grad ** 2)\n m_hat = new_m / (1.0 - beta1 ** t)\n v_hat = new_v / (1.0 - beta2 ** t)\n new_p = param * (1.0 - lr * weight_decay) - lr * m_hat / (torch.sqrt(v_hat) + eps)\n return new_p, new_m, new_v\n```",
"difficulty_breakdown": {
"concept": 3,
"implementation": 3,
"math_derivation": 3,
"boundary_complexity": 3
},
"test_cases": [
{
"name": "AdamW单步",
"input_desc": "真实参数契约用例 (param, grad, m, v, t, lr, weight_decay, beta1, beta2, eps)",
"params": [
{
"name": "param",
"type": "torch.Tensor",
"value": "torch.tensor([1.0], dtype=torch.float32)"
},
{
"name": "grad",
"type": "torch.Tensor",
"value": "torch.tensor([0.1], dtype=torch.float32)"
},
{
"name": "m",
"type": "torch.Tensor",
"value": "torch.tensor([0.0], dtype=torch.float32)"
},
{
"name": "v",
"type": "torch.Tensor",
"value": "torch.tensor([0.0], dtype=torch.float32)"
},
{
"name": "t",
"type": "int",
"value": "1"
},
{
"name": "lr",
"type": "float",
"value": "0.001"
},
{
"name": "weight_decay",
"type": "float",
"value": "0.01"
},
{
"name": "beta1",
"type": "float",
"value": "0.9"
},
{
"name": "beta2",
"type": "float",
"value": "0.999"
},
{
"name": "eps",
"type": "float",
"value": "1e-8"
}
],
"raw_inputs": {
"param": "torch.tensor([1.0], dtype=torch.float32)",
"grad": "torch.tensor([0.1], dtype=torch.float32)",
"m": "torch.tensor([0.0], dtype=torch.float32)",
"v": "torch.tensor([0.0], dtype=torch.float32)",
"t": "1",
"lr": "0.001",
"weight_decay": "0.01",
"beta1": "0.9",
"beta2": "0.999",
"eps": "1e-8"
}
},
{
"name": "多批次扩展泛化验证",
"input_desc": "真实参数契约用例 (param, grad, m, v, t, lr, weight_decay, beta1, beta2, eps)",
"params": [
{
"name": "param",
"type": "torch.Tensor",
"value": "torch.tensor([1.0, 1.0], dtype=torch.float32)"
},
{
"name": "grad",
"type": "torch.Tensor",
"value": "torch.tensor([0.1, 0.15], dtype=torch.float32)"
},
{
"name": "m",
"type": "torch.Tensor",
"value": "torch.tensor([0.0, 0.0], dtype=torch.float32)"
},
{
"name": "v",
"type": "torch.Tensor",
"value": "torch.tensor([0.0, 0.0], dtype=torch.float32)"
},
{
"name": "t",
"type": "int",
"value": "1"
},
{
"name": "lr",
"type": "float",
"value": "0.001"
},
{
"name": "weight_decay",
"type": "float",
"value": "0.01"
},
{
"name": "beta1",
"type": "float",
"value": "0.9"
},
{
"name": "beta2",
"type": "float",
"value": "0.999"
},
{
"name": "eps",
"type": "float",
"value": "1e-8"
}
],
"raw_inputs": {
"param": "torch.tensor([1.0, 1.0], dtype=torch.float32)",
"grad": "torch.tensor([0.1, 0.15], dtype=torch.float32)",
"m": "torch.tensor([0.0, 0.0], dtype=torch.float32)",
"v": "torch.tensor([0.0, 0.0], dtype=torch.float32)",
"t": "1",
"lr": "0.001",
"weight_decay": "0.01",
"beta1": "0.9",
"beta2": "0.999",
"eps": "1e-8"
}
},
{
"name": "多维批次与参数泛化验证",
"input_desc": "真实参数契约用例 (param, grad, m, v, t, lr, weight_decay, beta1, beta2, eps)",
"params": [
{
"name": "param",
"type": "torch.Tensor",
"value": "torch.tensor([1.0], dtype=torch.float32)"
},
{
"name": "grad",
"type": "torch.Tensor",
"value": "torch.tensor([0.15], dtype=torch.float32)"
},
{
"name": "m",
"type": "torch.Tensor",
"value": "torch.tensor([0.0], dtype=torch.float32)"
},
{
"name": "v",
"type": "torch.Tensor",
"value": "torch.tensor([0.0], dtype=torch.float32)"
},
{
"name": "t",
"type": "int",
"value": "1"
},
{
"name": "lr",
"type": "float",
"value": "0.0008"
},
{
"name": "weight_decay",
"type": "float",
"value": "0.008"
},
{
"name": "beta1",
"type": "float",
"value": "0.9"
},
{
"name": "beta2",
"type": "float",
"value": "0.999"
},
{
"name": "eps",
"type": "float",
"value": "1e-8"
}
],
"raw_inputs": {
"param": "torch.tensor([1.0], dtype=torch.float32)",
"grad": "torch.tensor([0.15], dtype=torch.float32)",
"m": "torch.tensor([0.0], dtype=torch.float32)",
"v": "torch.tensor([0.0], dtype=torch.float32)",
"t": "1",
"lr": "0.0008",
"weight_decay": "0.008",
"beta1": "0.9",
"beta2": "0.999",
"eps": "1e-8"
}
}
]
}