Python PyTorch GreedyPerfPartitioner.partition用法及代碼示例

本文簡要介紹python語言中 torchrec.distributed.planner.partitioners.GreedyPerfPartitioner.partition 的用法。

用法:

partition(proposal: List[torchrec.distributed.planner.types.ShardingOption], storage_constraint: torchrec.distributed.planner.types.Topology) → List[torchrec.distributed.planner.types.ShardingOption]

參數：

• proposal(List[ShardingOption]) -填充的分片選項列表。

• storage_constraint(Topology) -設備拓撲。

返回：

所選計劃的分片選項列表。

返回類型：

列表[ShardingOption]

根據每個分片選項的partition_by 屬性在拓撲上放置分片選項。拓撲存儲和性能在放置結束時更新。

例子：

sharding_options = [
        ShardingOption(partition_by="uniform",
                shards=[
                    Shards(storage=1, perf=1),
                    Shards(storage=1, perf=1),
                ]),
        ShardingOption(partition_by="uniform",
                shards=[
                    Shards(storage=2, perf=2),
                    Shards(storage=2, perf=2),
                ]),
        ShardingOption(partition_by="device",
                shards=[
                    Shards(storage=3, perf=3),
                    Shards(storage=3, perf=3),
                ])
        ShardingOption(partition_by="device",
                shards=[
                    Shards(storage=4, perf=4),
                    Shards(storage=4, perf=4),
                ]),
    ]
topology = Topology(world_size=2)

# First [sharding_options[0] and sharding_options[1]] will be placed on the
# topology with the uniform strategy, resulting in

topology.devices[0].perf = (1,2)
topology.devices[1].perf = (1,2)

# Finally sharding_options[2] and sharding_options[3]] will be placed on the
# topology with the device strategy (see docstring of `partition_by_device` for
# more details).

topology.devices[0].perf = (1,2) + (3,4)
topology.devices[1].perf = (1,2) + (3,4)

# The topology updates are done after the end of all the placements (the other
# in the example is just for clarity).