-sdy-add-data-flow-edges
为每个数据流边插入 DataFlowEdgeOp
。
为数据流边的所有所有者值(即 getDataFlowEdgeOwners
针对模块中的每个操作返回的所有值)插入 DataFlowEdgeOp
。
插入的 DataFlowEdgeOp
将采用所有者目标的现有分片(如果存在)。
-sdy-apply-sharding-constraints
应用用于决定其输入分片的约束条件。
如果 ShardingConstraintOp
满足以下所有条件,则将其分片复制到其输入:
- 输入没有现有的分片。
ShardingConstraintOp
的分片已完全关闭。- 输入中没有其他类型为
ShardingConstraintOp
或ManualComputationOp
且采用不同分片的用户。
这些条件表明 ShardingConstraintOp
决定了其输入的分片方式。
如果输入由 DataFlowEdgeOp
生成,则我们会将所有 input
用法替换为 ShardingConstraintOp
,而不是设置操作的分片,以避免限制边缘的所有目标的分片。
请注意,无论是否进行此传递,ShardingConstraintOp
的分片都会在传播期间传播到其输入或用户,但由于维度的 closed 属性不会传播,因此请务必复制分片,以便在上述情况下完全遵循约束条件。
此外,如果张量被满足以下所有条件的 ShardingConstraintOp
链使用:
- 该张量不是由
ShardingConstraintOp
生成的,并且没有任何其他类型为ShardingConstraintOp
或ManualComputationOp
的其他用户; - 链中的所有
ShardingConstraintOp
(除了最后一个)都不会重复使用; - 链中的最后一个
ShardingConstraintOp
没有任何类型为ShardingConstraintOp
或ManualComputationOp
的用户(否则,它不是链中的最后一个);
然后,此传递会将链的输入的所有其他用法(在链中的最后一个 ShardingConstraintOp
之后 [并在同一代码块内] 定义)替换为链的结果,因为它应决定这些用法的分片。
-sdy-constant-splitter
拆分常量子计算,以便每个子计算都只使用一次。
拆分常量子计算,使其具有单个用户。
这可确保在常量子计算的不同用法之间不会传播分片,因为这被视为虚假依赖项(常量的用法不应仅因使用相同的常量而以相同的方式进行分片)。实际上,每种用法可以具有不同的分片,这些分片可以独立传播到常量子计算的各自副本。
常量子计算可以是:
- 常量或 iota 运算(无运算数)
- 广播、切片或纯元素级运算,其操作数都由常量子计算(递归)定义,以及定义其操作数的整个子计算。
请注意,在常量子计算中,某个值可以在该子计算中多次使用。
-sdy-lift-inlined-meshes
将分片中的内嵌 MeshAttr
提升为符号 MeshOp
。
将 TensorShardingAttr
中的任何内嵌 MeshAttr
替换为网格符号名称,引用模块中的现有或新 MeshOp
,以便没有两个 MeshOp
具有相同的 MeshAttr
(现有 MeshOp
也会被删除重复项)。
每个新 MeshOp
的名称将是:
maximal_mesh_{device-id}
(对于最大网格,即空轴列表和单个设备 ID),或- [
mesh
、mesh_0
、mesh_1
、…] 中第一个可用的名称。
-sdy-manual-axes-cleanup
清理了 ManualComputationOp
中手动轴的使用
- 对于未指定手动轴的任何入/出分片,请将该手动轴添加到其 replicated_axes。这是为了确保始终完全指定手动轴。
- 按网格轴声明顺序对手动轴进行排序。
-sdy-sharding-group-import
对分片组进行了规范化和验证。
在导入时对分片群组应用规范化和验证。具体包括:
分片群组统一
使用组成员资格的传递性属性组合分片组。每当张量 T 同时位于分片组 G1 和分片组 G2 中时,我们都可以推断出 G1 和 G2 中的所有成员都应采用相同的方式进行分片。因此,我们可以将 G1 和 G2 合并为一个组。合并后,一组最小组的规范组 ID 将为 0,1,...N-1。
分片群组验证
验证分片组的格式是否正确,以及是否符合实现中的假设。目前,此断言表示,如果分片组包含在
ManualComputationOp
块中定义的Value
,则该组中的所有其他值都必须位于同一块中。