Blame for data/problems/adam-optimizer-memory-bytes.json

{
  "slug": "adam-optimizer-memory-bytes",
  "title": "手撕 Adam 优化器状态显存占用计算",
  "difficulty": "Easy",
  "category": "AIInfra",
  "tags": [
    "优化器显存",
    "Adam",
    "分布式训练"
  ],
  "description_md": "### 题目描述\n标准混合精度训练下，优化器状态需存储 FP32 Master Weights (4B)、一阶动量 $m$ (4B)、二阶动量 $v$ (4B)，每参数共 12 字节。\n\n### 函数签名\n```python\ndef adam_memory_bytes(num_params: int) -> int:\n```\n\n### 示例\n```python\n# 输入\nnum_params = 7000000000\n\n# 输出 (int)\n84000000000\n```\n\n### 参数规格与约束\n- **`num_params`**：`int`\n- **返回值**：`int`，输出规格 `int` (示例值: `84000000000`)\n- **精度容差**：`atol=0.0001, rtol=0.0001`",
  "entry_point": "adam_memory_bytes",
  "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 adam_memory_bytes(num_params: int) -> int:\n    \"\"\"\n    手撕 Adam 优化器状态显存占用计算\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 adam_memory_bytes(num_params: int) -> int:\n    return num_params * 12\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    ('7B 模型 Adam 优化器显存', 7000000000),\n    ('数值缩放与幅值扰动验证', 14000000000),\n    ('多维批次与参数泛化验证', 28000000000),\n    ('边界极值与稳定性验证', 56000000000),\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": false,
  "editorial_md": "#### 💡 核心算法原理与数学推导\n**手撕 Adam 优化器状态显存占用计算**（`adam_memory_bytes`）聚焦于 **AIInfra** 中的关键计算逻辑：\n\n标准混合精度训练下，优化器状态需存储 FP32 Master Weights (4B)、一阶动量 $m$ (4B)、二阶动量 $v$ (4B)，每参数共 12 字节。\n\n在实现中，函数接收参数 `num_params` (`int`)，目标是计算出类型与形状严格符合 `int` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`num_params` 规格为 `int` (示例值: `7000000000`)。\n- **期望输出规格**：`int` (示例值: `84000000000`)。\n\n**具体推导与计算步骤如下：**\n1. 返回最终结果：输出 `num_params * 12`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(1)$，仅涉及标量代数运算或闭式公式求值，执行耗时与模型规模参数数值大小无关。\n- **空间复杂度**：$O(1)$，仅使用常数级标量局部变量，无额外张量显存分配。\n\n#### ⚠️ 关键实现细节与易错点\n- **输入输出类型与形状契约**：请严格按照函数签名返回对应类型（`torch.Tensor`、`Tuple` 或标量数值），避免遗漏批次维度或改变张量的数据类型（`dtype`）。\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 adam_memory_bytes(num_params: int) -> int:\n    return num_params * 12\n```",
  "difficulty_breakdown": {
    "concept": 1,
    "implementation": 1,
    "math_derivation": 2,
    "boundary_complexity": 1
  },
  "test_cases": [
    {
      "name": "7B 模型 Adam 优化器显存",
      "input_desc": "真实参数契约用例 (num_params)",
      "params": [
        {
          "name": "num_params",
          "type": "int",
          "value": "7000000000"
        }
      ],
      "raw_inputs": {
        "num_params": "7000000000"
      }
    },
    {
      "name": "数值缩放与幅值扰动验证",
      "input_desc": "真实参数契约用例 (num_params)",
      "params": [
        {
          "name": "num_params",
          "type": "int",
          "value": "14000000000"
        }
      ],
      "raw_inputs": {
        "num_params": "14000000000"
      }
    },
    {
      "name": "多维批次与参数泛化验证",
      "input_desc": "真实参数契约用例 (num_params)",
      "params": [
        {
          "name": "num_params",
          "type": "int",
          "value": "28000000000"
        }
      ],
      "raw_inputs": {
        "num_params": "28000000000"
      }
    }
  ]
}