Python PyTorch uniform_partition用法及代码示例

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

用法:

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

为每个等级分配一个分片。

例子：

sharding_options = [
    [0,1,2,3],
    [0,1,2,3],
]
# with num_partitions=4

# The final output would be:
[
    partition_0 = [(0,0),(1,0)]
    partition_1 = [(0,1),(1,1)]
    partition_2 = [(0,2),(1,2)]
    partition_3 = [(0,3),(1,3)]
]