Spaces:
Running on Zero
Running on Zero
| { | |
| "slug": "activation-checkpointing-sqrt-memory", | |
| "title": "Activation Checkpointing 分段重计算激活显存开销模型", | |
| "difficulty": "Easy", | |
| "difficulty_breakdown": { | |
| "concept": 2, | |
| "implementation": 2, | |
| "math_derivation": 2, | |
| "boundary_complexity": 2 | |
| }, | |
| "category": "AIInfra", | |
| "tags": [ | |
| "ActivationCheckpointing", | |
| "GradientCheckpointing", | |
| "Memory", | |
| "TorchLeet" | |
| ], | |
| "description_md": "### 题目描述\n计算分段梯度检查点（Segmented Activation Checkpointing, Chen et al., 2016）的峰值激活显存（MB）：共 `num_layers` 层，每 `checkpoint_interval` ($K$) 层设一个检查点（共 $\\lceil L / K \\rceil$ 个检查点边界，每个边界只保存输入张量 `boundary_input_mb`）；在反向传播重计算某一段时，还需同时在显存中展开该段内 $K$ 层的完整中间激活 `K * layer_act_mb`。返回峰值激活显存 $\\lceil L / K \\rceil \\cdot M_{\\text{boundary}} + K \\cdot M_{\\text{layer}}$。\n\n### 函数签名\n```python\ndef checkpointed_activation_memory_mb(\n num_layers: int,\n checkpoint_interval: int,\n layer_act_mb: float,\n boundary_input_mb: float,\n) -> float:\n```\n\n### 示例\n```python\n# 输入\nnum_layers = 36\ncheckpoint_interval = 6\nlayer_act_mb = 100.0\nboundary_input_mb = 10.0\n\n# 输出 (float)\n660.0\n```\n\n### 参数规格与约束\n- **`num_layers`**：`int`\n- **`checkpoint_interval`**：`int`\n- **`layer_act_mb`**：`float`\n- **`boundary_input_mb`**：`float`\n- **返回值**：`float`，输出规格 `float` (示例值: `660.0`)\n- **精度容差**：`atol=0.0001, rtol=0.0001`\n> 题目来源：改编自 [TorchLeet](https://github.com/Exorust/TorchLeet)（MIT License），已改写为函数式评测接口与中文题面。", | |
| "entry_point": "checkpointed_activation_memory_mb", | |
| "starter_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ndef checkpointed_activation_memory_mb(\n num_layers: int,\n checkpoint_interval: int,\n layer_act_mb: float,\n boundary_input_mb: float,\n) -> float:\n # 请在此处编写你的实现代码\n pass\n", | |
| "reference_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ndef checkpointed_activation_memory_mb(num_layers: int, checkpoint_interval: int, layer_act_mb: float, boundary_input_mb: float) -> float:\n num_ckpt = (num_layers + checkpoint_interval - 1) // checkpoint_interval\n return float(num_ckpt * boundary_input_mb + checkpoint_interval * layer_act_mb)\n", | |
| "test_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ncases = [\n (\"36 层每 6 层一检查点 (sqrt(L) 最优)\", 36, 6, 100.0, 10.0),\n (\"每层全检查点 K=1\", 32, 1, 120.0, 15.0),\n (\"不设中间检查点 K=L\", 24, 24, 80.0, 8.0),\n (\"非整除层数 L=30,K=8\", 30, 8, 50.0, 5.0),\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", | |
| "test_cases": [ | |
| { | |
| "name": "36 层每 6 层一检查点 (sqrt(L) 最优)", | |
| "input_desc": "真实参数契约用例 (num_layers, checkpoint_interval, layer_act_mb, boundary_input_mb)", | |
| "params": [ | |
| { | |
| "name": "num_layers", | |
| "type": "int", | |
| "value": "36" | |
| }, | |
| { | |
| "name": "checkpoint_interval", | |
| "type": "int", | |
| "value": "6" | |
| }, | |
| { | |
| "name": "layer_act_mb", | |
| "type": "float", | |
| "value": "100.0" | |
| }, | |
| { | |
| "name": "boundary_input_mb", | |
| "type": "float", | |
| "value": "10.0" | |
| } | |
| ], | |
| "raw_inputs": { | |
| "num_layers": "36", | |
| "checkpoint_interval": "6", | |
| "layer_act_mb": "100.0", | |
| "boundary_input_mb": "10.0" | |
| } | |
| }, | |
| { | |
| "name": "每层全检查点 K=1", | |
| "input_desc": "真实参数契约用例 (num_layers, checkpoint_interval, layer_act_mb, boundary_input_mb)", | |
| "params": [ | |
| { | |
| "name": "num_layers", | |
| "type": "int", | |
| "value": "32" | |
| }, | |
| { | |
| "name": "checkpoint_interval", | |
| "type": "int", | |
| "value": "1" | |
| }, | |
| { | |
| "name": "layer_act_mb", | |
| "type": "float", | |
| "value": "120.0" | |
| }, | |
| { | |
| "name": "boundary_input_mb", | |
| "type": "float", | |
| "value": "15.0" | |
| } | |
| ], | |
| "raw_inputs": { | |
| "num_layers": "32", | |
| "checkpoint_interval": "1", | |
| "layer_act_mb": "120.0", | |
| "boundary_input_mb": "15.0" | |
| } | |
| }, | |
| { | |
| "name": "不设中间检查点 K=L", | |
| "input_desc": "真实参数契约用例 (num_layers, checkpoint_interval, layer_act_mb, boundary_input_mb)", | |
| "params": [ | |
| { | |
| "name": "num_layers", | |
| "type": "int", | |
| "value": "24" | |
| }, | |
| { | |
| "name": "checkpoint_interval", | |
| "type": "int", | |
| "value": "24" | |
| }, | |
| { | |
| "name": "layer_act_mb", | |
| "type": "float", | |
| "value": "80.0" | |
| }, | |
| { | |
| "name": "boundary_input_mb", | |
| "type": "float", | |
| "value": "8.0" | |
| } | |
| ], | |
| "raw_inputs": { | |
| "num_layers": "24", | |
| "checkpoint_interval": "24", | |
| "layer_act_mb": "80.0", | |
| "boundary_input_mb": "8.0" | |
| } | |
| } | |
| ], | |
| "time_limit": 5.0, | |
| "memory_limit_mb": 512, | |
| "atol": 0.0001, | |
| "rtol": 0.0001, | |
| "forbidden_modules": [ | |
| "os", | |
| "sys", | |
| "subprocess", | |
| "socket", | |
| "shutil", | |
| "pathlib", | |
| "ctypes", | |
| "multiprocessing" | |
| ], | |
| "forbidden_calls": [ | |
| "eval", | |
| "exec", | |
| "compile", | |
| "open", | |
| "__import__", | |
| "input", | |
| "breakpoint" | |
| ], | |
| "is_featured": false, | |
| "editorial_md": "#### 💡 核心算法原理与数学推导\n**Activation Checkpointing 分段重计算激活显存开销模型**（`checkpointed_activation_memory_mb`）聚焦于 **AIInfra** 中的关键计算逻辑：\n\n计算分段梯度检查点（Segmented Activation Checkpointing, Chen et al., 2016）的峰值激活显存（MB）：共 `num_layers` 层，每 `checkpoint_interval` ($K$) 层设一个检查点（共 $\\lceil L / K \\rceil$ 个检查点边界，每个边界只保存输入张量 `boundary_input_mb`）；在反向传播重计算某一段时，还需同时在显存中展开该段内 $K$ 层的完整中间激活 `K * layer_act_mb`。返回峰值激活显存 $\\lceil L / K \\rceil \\cdot M_{\\text{boundary}} + K \\cdot M_{\\text{layer}}$。\n\n在实现中，函数接收参数 `num_layers` (`int`)、`checkpoint_interval` (`int`)、`layer_act_mb` (`float`)、`boundary_input_mb` (`float`)，目标是计算出类型与形状严格符合 `float` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`num_layers` 规格为 `int` (示例值: `36`)，`checkpoint_interval` 规格为 `int` (示例值: `6`)，`layer_act_mb` 规格为 `float` (示例值: `100.0`)，`boundary_input_mb` 规格为 `float` (示例值: `10.0`)。\n- **期望输出规格**：`float` (示例值: `660.0`)。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `num_ckpt`：通过表达式 `(num_layers + checkpoint_interval - 1) // checkpoint_interval` 完成张量或数值变换。\n2. 返回最终结果：输出 `float(num_ckpt * boundary_input_mb + checkpoint_interval * layer_act_mb)`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(1)$，仅涉及标量代数运算或闭式公式求值，执行耗时与规模参数大小无关。\n- **空间复杂度**：$O(1)$，仅使用常数级标量局部变量，无额外张量显存分配。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，需加上 `eps` 或执行 `clamp_min` 防止除零产生 `NaN` / `Inf`。\n- **边界输入鲁棒性**：注意单样本（`Batch=1`）、零值输入或极端超参数下的计算行为，禁止硬编码特定测试用例常量。\n- **向量化算子优先**：优先使用 PyTorch 原生向量化张量算子，避免低效的逐元素循环。\n\n#### 💻 官方标准参考实现\n```python\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ndef checkpointed_activation_memory_mb(num_layers: int, checkpoint_interval: int, layer_act_mb: float, boundary_input_mb: float) -> float:\n num_ckpt = (num_layers + checkpoint_interval - 1) // checkpoint_interval\n return float(num_ckpt * boundary_input_mb + checkpoint_interval * layer_act_mb)\n```" | |
| } |