-sdy-add-data-flow-edges
为每个数据流边插入 DataFlowEdgeOp
。
为数据流边的所有所有者值(即 getDataFlowEdgeOwners
针对模块中的每个操作返回的所有值)插入 DataFlowEdgeOp
。
插入的 DataFlowEdgeOp
将采用所有者目标的现有分片(如果存在)。
TODO(b/330339693):移除 getDataFlowEdgeOwners
后更新此文档。
-sdy-apply-sharding-constraints
应用用于决定其输入分片的约束条件。
如果 ShardingConstraintOp
满足以下所有条件,则将其分片复制到其输入中:
- 输入没有现有的分片。
- 输入不是由
DataFlowEdgeOp
生成的,它存放着边缘的所有目标的分片。 ShardingConstraintOp
的分片已完全关闭。- 输入没有任何其他具有其他分片的
ShardingConstraintOp
或ManualComputationOp
用户。
这表示 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
中手动轴的使用
1) 对于任何未指定手动轴的输入/输出分片,请将该手动轴添加到其 duplicate_axes。这是为了确保始终完全指定手动轴。
2) 按网格轴声明顺序对手动轴进行排序。
-sdy-sharding-group-import
对分片组进行了规范化和验证。
在导入时对分片组应用规范化和验证。具体包括:
1) 分片组合并 - 使用组成员资格的传递性属性组合分片组。每当张量 T 同时位于分片组 G1 和分片组 G2 中时,我们都可以推断出 G1 和 G2 中的所有成员都应采用相同的方式进行分片。因此,我们可以将 G1 和 G2 合并为一个组。合并后的规范组 ID 集将为 0,1,...N-1(对于最小组集)。
2) 分片组验证:验证分片组的格式是否正确,以及是否符合实现中的假设。目前,此断言表示,如果分片组包含在 ManualComputationOp
块中定义的 Value
,则该组中的所有其他值都必须位于同一块中。