Python PyTorch greedy_partition用法及代码示例

本文简要介绍python语言中 torchrec.distributed.planner.partitioners.greedy_partition 的用法。

用法:

torchrec.distributed.planner.partitioners.greedy_partition(num_partitions: int, sharding_options: List[torchrec.distributed.planner.types.ShardingOption], shard_idxes: Optional[List[Tuple[int, int]]] = None, partition_sums: Optional[List[float]] = None, mem_cap: Optional[List[torchrec.distributed.planner.types.Storage]] = None) → List[List[Tuple[int, int]]]

返回：

应该分配给每个分区的 (option_idx, shard_idx) 的索引列表。

返回类型：

列表[列表[元组[int，int]]]

根据与每个 [option_idx, shard_idx] 关联的性能权重，以贪心的方式在 num_partitions 分区之间划分索引。

例子：

sharding_options = [
    [0,1,2,3] with perfs [10,20,30,40]
    [0,1] with perfs [200,300]
]
# with num_partitions=3

# The final output would be:
[
    partition_0 = [(1,1)], with a perf of 300
    partition_1 = [(1,0)], with a perf of 200
    partition_2 = [(0,0),(0,1),(0,2),(0,3)], with a perf of 100 (10+20+30+40)
]