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)
]