新增API local_map #71804

zty-king · 2025-03-20T06:18:56Z

PR Category

Auto Parallel

PR Types

Others

Description

新增API local_map

1、相关背景

在分布式训练场景中,经常需要将分布式张量(dist_tensor)传递给为仅仅能处理普通张量(dense_tensor)或者必须以本地视角处理本地张量的函数。为了简化这个过程,需要提供一个工具函数来处理分布式张量到普通张量的转换,以及将函数处理的结果重新加上分布式属性。local_map API 就是为了解决这个问题而设计的。

2、功能目标

local_map 函数的主要功能是允许用户将分布式张量传递给为普通张量编写的函数。它实现了以下目标:

本地分片提取: 从分布式张量中提取本地分片数据
函数应用: 将用户函数应用于提取出的本地分片
结果包装: 根据指定的切分标记方式将结果重新包装为分布式张量
分布验证: 验证输入输出是否符合要求
自动reshard: 支持在需要时自动对输入张量进行reshard操作

3、意义

为 Paddle 分布式训练提供更便捷的张量处理方式,使得用户可以轻松地在分布式环境中复用为普通张量编写的函数。

4、常见使用场景

带 mask 的 loss 计算:需要在每张卡上独立计算 masked token 的 loss
MoE (混合专家模型)相关计算:

aux_loss 计算:基于每张卡上专家分配到的局部 token 数进行计算
z_loss 计算:对每张卡上的 logits 独立计算 z_loss
张量 reshape 操作:在局部维度上进行 shape 变换

其他需要保持局部计算语义的场景

5、local_map相比较LocalLayer的优化点

从结构上看，local_map是一个函数装饰器，直接包装普通函数，无需像LocalLayer那样继承Layer类去使用，无需管理Layer的状态，使用起来更方便，逻辑更清晰，相对用户使用比较友好。
支持混合输入(分布式张量、普通张量、以及一些函数运算中必须要使用到非tensor的数值参数)，使用更灵活，适应性更强。
支持自动reshard逻辑，可以对输入的分布式张量进行批量的reshard，无需用户手动多次调用reshard。
支持自动推导process_mesh。
支持动态图和静态图的自适应。
可以直接处理任意Python函数，不限于Layer中的forward方法。
支持混合输出(分布式张量、普通张量、以及一些非tensor的数值)，更灵活的输出处理方式。

paddle-bot · 2025-03-20T06:19:03Z

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

From00 · 2025-03-21T03:08:46Z

python/paddle/distributed/auto_parallel/local_map.py

@@ -0,0 +1,280 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.


2024 -> 2025

现在仍然是2024

From00 · 2025-03-21T03:18:17Z

python/paddle/distributed/auto_parallel/local_map.py

+__all__ = ["local_map"]
+
+PlacementType = Sequence[dist.Placement] | None
+InputPlacements = tuple[PlacementType, ...] | None


为什么输入和输出支持的Placements参数类型不同？

已经统一

From00 · 2025-03-21T03:29:40Z

python/paddle/distributed/auto_parallel/local_map.py

+def local_map(
+    func: Callable,
+    out_placements: OutputPlacements,
+    in_placements: InputPlacements | None,


Suggested change

in_placements: InputPlacements | None,

in_placements: Optional[tuple[list[dist.Placement], ...]],

没必要新创建太多类型命名，与框架中其它模块的使用习惯都不相同，反而增加用户理解成本。

这里不能用Optional，pre-commit会报错，但是已经修改了类型命名

From00 · 2025-03-21T03:34:22Z

python/paddle/distributed/auto_parallel/local_map.py

+
+        for idx, arg in enumerate(flat_args):
+            if _is_distributed_tensor(arg):
+                # TODO: the current code doesn't consider the uneven sharding case


这个注释是啥意思？

已经删除，暂时不考虑这个

From00 · 2025-03-21T03:35:11Z

python/paddle/distributed/auto_parallel/local_map.py

+    redistribute_inputs: bool | None,
+):
+    """
+    :meth:`local_map` is an experimental API that allows users to pass dist_tensors


我们没有experimental API这种说法

From00 · 2025-03-21T03:37:10Z

python/paddle/distributed/auto_parallel/local_map.py

+                    if arg.placements != spec:
+                        if redistribute_inputs:
+                            # Redistribute to input placements
+                            arg = arg.redistribute(process_mesh, spec)


我们有redistribute这个接口吗？

已经删除并按照paddle框架改写

From00 · 2025-03-21T03:38:29Z

python/paddle/distributed/auto_parallel/local_map.py

+    in_placements: InputPlacements | None,
+    process_mesh: ProcessMesh | None,
+    *,
+    redistribute_inputs: bool | None,


命名要符合现有框架的习惯，我们没有redistribute这种叫法

From00 · 2025-03-21T03:39:55Z

python/paddle/distributed/auto_parallel/local_map.py

+        else:
+            return out
+
+    def _is_distributed_tensor(tensor) -> bool:


Suggested change

def _is_distributed_tensor(tensor) -> bool:

def is_dist_tensor(tensor) -> bool:

这是一个很基础的方法，应该放在更公共的地方，方便其它模块复用。

From00 · 2025-03-21T03:43:13Z

python/paddle/distributed/auto_parallel/local_map.py

+
+            return pack_sequence_as(out, flat_dist_out)
+        else:
+            return out


如果用户的输入没有dist_tensor，但指定了输出的分布式标记，这个时候直接忽略输出标记，是一种合理的行为吗？

已经增加了这种情况的处理

… pipeline

From00 · 2025-03-25T06:17:29Z

python/paddle/distributed/auto_parallel/local_map.py

+if TYPE_CHECKING:
+    from paddle.distributed import ProcessMesh
+
+__all__ = ["local_map"]


这个不通过auto_parallel.local_map路径对外暴露，不应该加在这个文件的__all__里

zhiqiu · 2025-03-25T06:30:22Z

test/auto_parallel/pir/CMakeLists.txt

@@ -19,6 +19,8 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
  py_test_modules(test_mlp MODULES test_mlp ENVS FLAGS_enable_pir_api=1)
  py_test_modules(test_local_layer MODULES test_local_layer ENVS
                  FLAGS_enable_pir_api=1)
+  py_test_modules(test_local_map MODULES test_local_map ENVS
+                  FLAGS_enable_pir_api=1)


why still need FLAGS_enable_pir_api=1?

之前是为了对标LocalLayer，不过现在好像确实不用标记为pir，单测都会自动走，已修改

From00

LGTM

zhiqiu

LGTM

jeff41404 · 2025-04-08T07:31:04Z

According to the newly added API specification of paddle, it is necessary to write the API Chinese documentation in docs repo for users to refer to the official website. please add link of docs repo PR in description above.

zty-king · 2025-04-08T14:28:13Z

According to the newly added API specification of paddle, it is necessary to write the API Chinese documentation in docs repo for users to refer to the official website. please add link of docs repo PR in description above.

Done
https://github.com/PaddlePaddle/docs/pull/7245/files?short_path=6326c88#diff-6326c883990bd5d2a43546a1c0e4f990a2b4d2dbb02aaa54a6013b42768ca348

jeff41404

LGTM

… pipeline

sunzhongkai588 · 2025-04-10T04:25:00Z

python/paddle/distributed/auto_parallel/local_map.py

+    in_placements: list[list[dist.Placement]] | None,
+    process_mesh: ProcessMesh | None,


默认值的类型注释是不是还得加一下 = None？ @SigureMo 一师傅看看

这得看接口形态，如果没有默认值的话就是不需要加 = None 的啊

这得看接口形态，如果没有默认值的话就是不需要加 = None 的啊

麻烦您看看还有其他什么大问题吗，没有的话，可以先approval吗，格式我会再提一个pr修复

Done

这里我没说过要改，这里具体是否要加 = None 需要看接口形态，这里完全是 @sunzhongkai588 不懂这里写的评论，我只是给他解释这一点

请 review 这里的改动，如果从接口形态上来看确实需要有默认值，且需要为 None，那么可以改这里，否则不需要改这里

sunzhongkai588 · 2025-04-10T04:26:44Z

python/paddle/distributed/auto_parallel/local_map.py

+            Default: None
+
+        reshard_inputs (bool, optional):
+            the bool value indicating whether to reshard the input :dist_tensor` s when


:dist_tensor` 是不是写错了

sunzhongkai588

文档问题之后再提 PR 修复，一师傅别忘了回复一下

SigureMo

中文文档 PR 也没写么？可以下个 PR，反正 @sunzhongkai588 同意了

SigureMo · 2025-04-11T02:31:04Z

python/paddle/distributed/auto_parallel/local_map.py

+
+
+def local_map(
+    func: Callable,


泛型写清楚内部参数类型

在实际使用当中，理论上内部参数接受任意类型变量，这也要写吗

需要，Callable[..., Any]

SigureMo · 2025-04-11T02:32:22Z

python/paddle/distributed/auto_parallel/local_map.py

+    in_placements: list[list[dist.Placement]] | None,
+    process_mesh: ProcessMesh | None,
+    reshard_inputs: bool = False,
+):


写清楚返回值类型

没看到啊……

sorry，理解错了您的意思，现在加了

SigureMo · 2025-04-11T02:33:32Z

python/paddle/distributed/auto_parallel/local_map.py

+            in_placements.
+
+    Example:
+        >>> from __future__ import annotations


这个示例代码格式不对吧，这样能正确渲染吗？就算英文能正确渲染，中文文档也无法使用 COPY-FROM copy 过去

SigureMo · 2025-04-11T02:36:17Z

python/paddle/distributed/auto_parallel/local_map.py

+    Raises:
+        AssertionError: If the number of output placements does not match the number
+            of function outputs.
+
+        AssertionError: If a non-tensor output has a non-None placement specified.
+
+        AssertionError: If process_mesh is None and there are no dist_tensor inputs
+            but out_placements contains non-None values.
+
+        ValueError: If the input dist_tensor placements don't match the required
+            in_placements.


根据文档规范不写 Raises

zty-king · 2025-04-11T02:45:21Z

中文文档写了的，好的，格式问题我统一修改一下

---- 回复的原邮件 ---- | 发件人 | Nyakku ***@***.***> | | 发送日期 | 2025年04月11日 10:41 | | 收件人 | PaddlePaddle/Paddle ***@***.***> | | 抄送人 | zty-king ***@***.***>, Author ***@***.***> | | 主题 | Re: [PaddlePaddle/Paddle] 新增API local_map (PR #71804) | @SigureMo approved this pull request. 中文文档 PR 也没写么？可以下个 PR，反正 @sunzhongkai588 同意了 In python/paddle/distributed/auto_parallel/local_map.py:

+# limitations under the License.

+from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Callable + +import paddle +import paddle.distributed as dist +from paddle.utils import flatten, pack_sequence_as + +if TYPE_CHECKING: + from paddle.distributed import ProcessMesh + + +def local_map( + func: Callable, 泛型写清楚内部参数类型 In python/paddle/distributed/auto_parallel/local_map.py:

+

+import paddle +import paddle.distributed as dist +from paddle.utils import flatten, pack_sequence_as + +if TYPE_CHECKING: + from paddle.distributed import ProcessMesh + + +def local_map( + func: Callable, + out_placements: list[list[dist.Placement]], + in_placements: list[list[dist.Placement]] | None, + process_mesh: ProcessMesh | None, + reshard_inputs: bool = False, +): 写清楚返回值类型 In python/paddle/distributed/auto_parallel/local_map.py:

+ and returns dist_tensors constructed from the return values of ``func``.

+ + Raises: + AssertionError: If the number of output placements does not match the number + of function outputs. + + AssertionError: If a non-tensor output has a non-None placement specified. + + AssertionError: If process_mesh is None and there are no dist_tensor inputs + but out_placements contains non-None values. + + ValueError: If the input dist_tensor placements don't match the required + in_placements. + + Example: + >>> from __future__ import annotations 这个示例代码格式不对吧，这样能正确渲染吗？就算英文能正确渲染，中文文档也无法使用 COPY-FROM copy 过去 In python/paddle/distributed/auto_parallel/local_map.py:

+ Raises:

+ AssertionError: If the number of output placements does not match the number + of function outputs. + + AssertionError: If a non-tensor output has a non-None placement specified. + + AssertionError: If process_mesh is None and there are no dist_tensor inputs + but out_placements contains non-None values. + + ValueError: If the input dist_tensor placements don't match the required + in_placements. 根据文档规范不写 Raises — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.Message ID: ***@***.***>

zty-king · 2025-04-11T03:29:34Z

好的

---- 回复的原邮件 ---- | 发件人 | Nyakku ***@***.***> | | 发送日期 | 2025年04月11日 11:28 | | 收件人 | PaddlePaddle/Paddle ***@***.***> | | 抄送人 | zty-king ***@***.***>, Author ***@***.***> | | 主题 | Re: [PaddlePaddle/Paddle] 新增API local_map (PR #71804) | @SigureMo commented on this pull request. In python/paddle/distributed/auto_parallel/local_map.py:

+# limitations under the License.

+from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Callable + +import paddle +import paddle.distributed as dist +from paddle.utils import flatten, pack_sequence_as + +if TYPE_CHECKING: + from paddle.distributed import ProcessMesh + + +def local_map( + func: Callable, 需要，Callable[..., Any] — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.Message ID: ***@***.***>

… pipeline

paddle-ci-bot · 2025-04-14T02:54:08Z

Sorry to inform you that fdbfcf3's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

zty-king · 2025-04-15T15:58:18Z

中文文档 PR 也没写么？可以下个 PR，反正 @sunzhongkai588 同意了

PaddlePaddle/docs#7245 这个是对应中文文档PR

zty-king · 2025-04-15T16:03:04Z

明白，这里确实是加默认值的

---- 回复的原邮件 ---- | 发件人 | Nyakku ***@***.***> | | 发送日期 | 2025年04月15日 23:58 | | 收件人 | PaddlePaddle/Paddle ***@***.***> | | 抄送人 | zty-king ***@***.***>, Author ***@***.***> | | 主题 | Re: [PaddlePaddle/Paddle] 新增API local_map (PR #71804) | @SigureMo commented on this pull request. In python/paddle/distributed/auto_parallel/local_map.py:

+ in_placements: list[list[dist.Placement]] | None,

+ process_mesh: ProcessMesh | None, Done 这里我没说过要改，这里具体是否要加 = None 需要看接口形态，这里完全是 @sunzhongkai588 不懂这里写的评论，我只是给他解释这一点请 review 这里的改动，如果从接口形态上来看确实需要有默认值，且需要为 None，那么可以改这里，否则不需要改这里 — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.Message ID: ***@***.***>

SigureMo

LGTMeow for type annotations

zty-king · 2025-04-16T13:28:28Z

@jeff41404 按要求修改了一下格式，看起来需要麻烦您再review一下

sunzhongkai588

LGTM

XiaoguangHu01

LGTM

* 新增API local_map * 修正文件格式 * 优化了local_map一些功能 * 新增reshard功能，同时兼容动静态下的local_map调用 * 修改格式规范 * 修正单测的接口命名 * 用local_map替换LocalLayer * 单测使用local_map时的参数设置修改，reshard设置为True * 修正单测 * 修改单测 * 修改格式 * 修改格式 * 修改测试样例格式 * 修改测试样例格式

新增API local_map

5c7f615

zty-king requested review from sneaxiy and raindrops2sea as code owners March 20, 2025 06:18

paddle-bot bot added the contributor External developers label Mar 20, 2025

修正文件格式

e4f6cd6

From00 reviewed Mar 21, 2025

View reviewed changes

zty-king added 4 commits March 23, 2025 17:50

优化了local_map一些功能

035ef2e

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

e0b4ecb

… pipeline

新增reshard功能，同时兼容动静态下的local_map调用

9cf6bf9

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

7423816

… pipeline

From00 reviewed Mar 25, 2025

View reviewed changes

zhiqiu reviewed Mar 25, 2025

View reviewed changes

zty-king added 6 commits March 26, 2025 15:29

修改格式规范

22d43dd

修正单测的接口命名

938bc06

用local_map替换LocalLayer

a206f43

单测使用local_map时的参数设置修改，reshard设置为True

e9a2cb5

修正单测

8d91329

修改单测

fdbfcf3

From00 previously approved these changes Apr 7, 2025

View reviewed changes

zhiqiu previously approved these changes Apr 8, 2025

View reviewed changes

jeff41404 reviewed Apr 9, 2025

View reviewed changes

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

3d21983

… pipeline

sunzhongkai588 reviewed Apr 10, 2025

View reviewed changes

sunzhongkai588 previously approved these changes Apr 10, 2025

View reviewed changes

SigureMo previously approved these changes Apr 11, 2025

View reviewed changes

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

aafbc74

… pipeline

修改格式

798c1a6

zty-king dismissed stale reviews from zhiqiu, From00, SigureMo, and sunzhongkai588 via 798c1a6 April 14, 2025 16:18

zty-king added 3 commits April 14, 2025 16:19

修改格式

187b927

修改测试样例格式

ba757b9

修改测试样例格式

8075976

SigureMo approved these changes Apr 15, 2025

View reviewed changes

From00 approved these changes Apr 23, 2025

View reviewed changes

jeff41404 approved these changes Apr 23, 2025

View reviewed changes

sunzhongkai588 approved these changes Apr 23, 2025

View reviewed changes

XiaoguangHu01 approved these changes Apr 24, 2025

View reviewed changes

From00 merged commit 483c9e1 into PaddlePaddle:develop Apr 24, 2025
36 of 37 checks passed

		@@ -0,0 +1,280 @@
		# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.

	in_placements: InputPlacements \| None,
	in_placements: Optional[tuple[list[dist.Placement], ...]],

	def _is_distributed_tensor(tensor) -> bool:
	def is_dist_tensor(tensor) -> bool:

		in_placements: list[list[dist.Placement]] \| None,
		process_mesh: ProcessMesh \| None,

新增API local_map #71804

新增API local_map #71804

Uh oh!

Conversation

zty-king commented Mar 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Category

PR Types

Description

1、相关背景

2、功能目标

3、意义

4、常见使用场景

5、local_map相比较LocalLayer的优化点

Uh oh!

paddle-bot bot commented Mar 20, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zty-king Mar 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zty-king Mar 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

From00 left a comment

Choose a reason for hiding this comment

Uh oh!

zhiqiu left a comment

Choose a reason for hiding this comment

Uh oh!

jeff41404 commented Apr 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zty-king commented Apr 8, 2025

Uh oh!

zty-king commented Mar 20, 2025 •

edited

Loading

zty-king Mar 23, 2025 •

edited

Loading

zty-king Mar 23, 2025 •

edited

Loading

jeff41404 commented Apr 8, 2025 •

edited

Loading