{ "slug": "allgather-ring-chunk-steps", "title": "手撕 Ring AllGather 环形步进块收发跟踪 (Ring AllGather Steps)", "difficulty": "Hard", "category": "AIInfra", "tags": [ "RingAllGather", "环形步进", "块传输索引" ], "description_md": "### 题目描述\n在包含 $P$ 个节点的 Ring AllGather 中，共需进行 $P - 1$ 步环形中继。设卡 $r$ 在第 $s$ 步 ($s = 0, \\dots, P-2$) 发送其持有的块号，接收来自左邻居的块号：第 $s$ 步发送块号 $(r - s) \\pmod P$，接收块号 $(r - s - 1) \\pmod P$。返回卡 $rank$ 在整个过程中的 `(send_chunks, recv_chunks)` 列表。\n\n### 函数签名\n```python\ndef allgather_steps(\n world_size: int,\n rank: int,\n) -> Tuple[List[int], List[int]]:\n```\n\n### 示例\n```python\n# 输入\nworld_size = 4\nrank = 1\n\n# 输出 (tuple)\n([1, 0, 3], [0, 3, 2])\n```\n\n### 参数规格与约束\n- **`world_size`**：`int`\n- **`rank`**：`int`\n- **返回值**：`Tuple[List[int], List[int]]`，输出规格 `Tuple` (`List` (长度 `3`), `List` (长度 `3`))\n- **精度容差**：`atol=0.0001, rtol=0.0001`", "entry_point": "allgather_steps", "starter_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import Optional, List, Tuple, Dict, Any\n\ndef allgather_steps(\n world_size: int,\n rank: int,\n) -> Tuple[List[int], List[int]]:\n \"\"\"\n 手撕 Ring AllGather 环形步进块收发跟踪 (Ring AllGather Steps)\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, Dict, Any\n\ndef allgather_steps(world_size: int, rank: int) -> Tuple[List[int], List[int]]:\n send_list = []\n recv_list = []\n p = world_size\n for s in range(p - 1):\n send_chunk = (rank - s) % p\n recv_chunk = (rank - s - 1) % p\n send_list.append(send_chunk)\n recv_list.append(recv_chunk)\n return send_list, recv_list\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 ('4卡拓扑下第1卡的AllGather步骤', 4, 1),\n ('数值缩放与幅值扰动验证', 8, 1),\n ('多维批次与参数泛化验证', 16, 1),\n ('边界极值与稳定性验证', 2, 1),\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**手撕 Ring AllGather 环形步进块收发跟踪 (Ring AllGather Steps)**（`allgather_steps`）聚焦于 **AIInfra** 中的关键计算逻辑：\n\n在包含 $P$ 个节点的 Ring AllGather 中，共需进行 $P - 1$ 步环形中继。设卡 $r$ 在第 $s$ 步 ($s = 0, \\dots, P-2$) 发送其持有的块号，接收来自左邻居的块号：第 $s$ 步发送块号 $(r - s) \\pmod P$，接收块号 $(r - s - 1) \\pmod P$。返回卡 $rank$ 在整个过程中的 `(send_chunks, recv_chunks)` 列表。\n\n在实现中，函数接收参数 `world_size` (`int`)、`rank` (`int`)，目标是计算出类型与形状严格符合 `Tuple[List[int], List[int]]` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`world_size` 规格为 `int` (示例值: `4`)，`rank` 规格为 `int` (示例值: `1`)。\n- **期望输出规格**：`Tuple` (`List` (长度 `3`), `List` (长度 `3`))。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `send_list`：通过表达式 `[]` 完成张量或数值变换。\n2. 计算中间量 `recv_list`：通过表达式 `[]` 完成张量或数值变换。\n3. 计算中间量 `p`：通过表达式 `world_size` 完成张量或数值变换。\n4. 迭代循环 `for s in range(p - 1)`：逐步递推更新状态或聚合分块结果。\n5. 返回最终结果：输出 `send_list, recv_list`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(T \\cdot N)$，其中 $T$ 为迭代步数、$N$ 为单步处理的张量元素总数。\n- **空间复杂度**：$O(N)$，用于维护迭代状态张量与返回结果。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，必须按公式位置准确加上 `eps`（例如在平方根内 `sqrt(var + eps)` 还是分母外 `+ eps`），防止全零或常量输入触发 `NaN` / `Inf`。\n- **整数整除与返回类型契约**：涉及块索引、分片大小或显存字节换算时，需严格区分 Python 浮点除法 `/` 与整数整除 `//`，确保返回的索引、步数或字节数类型与函数签名一致。\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, Dict, Any\n\ndef allgather_steps(world_size: int, rank: int) -> Tuple[List[int], List[int]]:\n send_list = []\n recv_list = []\n p = world_size\n for s in range(p - 1):\n send_chunk = (rank - s) % p\n recv_chunk = (rank - s - 1) % p\n send_list.append(send_chunk)\n recv_list.append(recv_chunk)\n return send_list, recv_list\n```", "difficulty_breakdown": { "concept": 4, "implementation": 4, "math_derivation": 4, "boundary_complexity": 4 }, "test_cases": [ { "name": "4卡拓扑下第1卡的AllGather步骤", "input_desc": "真实参数契约用例 (world_size, rank)", "params": [ { "name": "world_size", "type": "int", "value": "4" }, { "name": "rank", "type": "int", "value": "1" } ], "raw_inputs": { "world_size": "4", "rank": "1" } }, { "name": "数值缩放与幅值扰动验证", "input_desc": "真实参数契约用例 (world_size, rank)", "params": [ { "name": "world_size", "type": "int", "value": "8" }, { "name": "rank", "type": "int", "value": "1" } ], "raw_inputs": { "world_size": "8", "rank": "1" } }, { "name": "多维批次与参数泛化验证", "input_desc": "真实参数契约用例 (world_size, rank)", "params": [ { "name": "world_size", "type": "int", "value": "16" }, { "name": "rank", "type": "int", "value": "1" } ], "raw_inputs": { "world_size": "16", "rank": "1" } } ] }