Part 9 — local_map 扩展（可选）

仅当首个需要 q/k/v 走 head-shard 路径的模型（如某些 attention 变种、GatedDeltaNet 的 TP 切分等）迁移时才做。与主线解耦，可独立推进，也可推迟。

1. 目标

为 Module.parallelize 增加 torchtitan 风格的 local_map 支持 —— 把 sharded DTensor 输入解包为 local 张量、调用纯 local 函数（如 SDPA / attention kernel）、再包回 DTensor。

2. 任务边界

3. 核心设计

3.1 跨后端 local_map 包装

# hyper_parallel/core/dtensor/local_map.py
from hyper_parallel.platform import get_platform
from hyper_parallel.platform.platform import PlatformType

platform = get_platform()


def local_map(
    func,
    out_placements,
    in_placements,
    in_grad_placements=None,
    device_mesh=None,
    *,
    redistribute_inputs=True,
):
    """Cross-backend wrapper around torch's experimental local_map."""
    if platform.platform_type == PlatformType.PYTORCH:
        # pylint: disable=C0415
        from torch.distributed.tensor.experimental import local_map as torch_local_map
        return torch_local_map(
            func,
            out_placements=out_placements,
            in_placements=in_placements,
            in_grad_placements=in_grad_placements,
            device_mesh=device_mesh,
            redistribute_inputs=redistribute_inputs,
        )
    if platform.platform_type == PlatformType.MINDSPORE:
        raise NotImplementedError(
            "local_map on MindSpore backend is pending. "
            "Currently the inner-attention TP path requires PyTorch backend."
        )
    raise RuntimeError(f"Unknown platform: {platform.platform_type}")

3.2 Module.parallelize 启用 local_map

把 M1 中的占位逻辑：

if sc.local_map is not None:
    raise NotImplementedError("local_map will be added in M9")

替换为：

if sc.local_map is not None:
    from hyper_parallel.core.dtensor.local_map import local_map
    inner_fn = self._unwrap_inner_callable(sc.local_map.callable_path)
    wrapped = local_map(
        inner_fn,
        out_placements=[
            resolve_placements(p, tp_mesh.mesh_dim_names) for p in sc.local_map.out_placements
        ],
        in_placements=[
            resolve_placements(p, tp_mesh.mesh_dim_names) for p in sc.local_map.in_placements
        ],
        device_mesh=tp_mesh,
        redistribute_inputs=True,
    )
    self._bind_inner_callable(sc.local_map.callable_path, wrapped)

3.3 set_gqa_inner_attention_local_map 助手

def set_gqa_inner_attention_local_map(
    inner_attn_cfg, *, return_lse: bool = False,
):
    """Mark inner attention kernel (e.g. F.scaled_dot_product_attention)
    as a local-map region: q/k/v come in as Shard(2) on head dim,
    SDPA runs on local heads, output comes back as Shard(2).
    """
    inner_attn_cfg.sharding_config = ShardingConfig(
        local_map=LocalMapConfig(
            callable_path="inner_sdpa",
            in_placements=(
                {MeshAxisName.TP: Shard(2)},          # q
                {MeshAxisName.TP: Shard(2)},          # k
                {MeshAxisName.TP: Shard(2)},          # v
            ),
            out_placements=(
                ({MeshAxisName.TP: Shard(2)},)        # attn_output
                if not return_lse
                else (
                    {MeshAxisName.TP: Shard(2)},
                    {MeshAxisName.TP: Shard(2)},     # lse
                )
            ),
        )
    )

4. 与 torchtitan 接口差异说明

5. 开发步骤

6. 验证标准

新建 tests/torch/ut/dtensor/test_local_map.py 和 tests/torch/st/local_map_attention/：

通过门槛：

• 4 个测试全绿。
• M9 启用后，已有 M6 / M8 模型行为不变（local_map 字段默认 None，路径无差异）。

7. 工期 & 依赖