规划器¶

TorchRec 规划器负责确定分布式训练和推理的最具性能和平衡的切片计划。

生成切片计划的主要 API 是 EmbeddingShardingPlanner.plan

class torchrec.distributed.types.ShardingPlan(plan: Dict[str, ModuleShardingPlan])¶

切片计划的表示。这使用更大包装模型的 FQN（即使用 DistributedModelParallel 包装的模型）EmbeddingModuleShardingPlan 应在需要 TorchRec 可组合性时使用。

plan¶

以模块路径为键的字典，以参数名称为键的参数切片规范字典。

类型:: Dict[str, EmbeddingModuleShardingPlan]

get_plan_for_module(module_path: str) → Optional[ModuleShardingPlan]¶

参数:: module_path (str) –
返回值:: 以参数名称为键的参数切片规范字典。如果给定 module_path 不存在切片规范，则为 None。
返回类型:: Optional[ModuleShardingPlan]

class torchrec.distributed.planner.planners.EmbeddingShardingPlanner(topology: Optional[Topology] = None, batch_size: Optional[int] = None, enumerator: Optional[Enumerator] = None, storage_reservation: Optional[StorageReservation] = None, proposer: Optional[Union[Proposer, List[Proposer]]] = None, partitioner: Optional[Partitioner] = None, performance_model: Optional[PerfModel] = None, stats: Optional[Union[Stats, List[Stats]]] = None, constraints: Optional[Dict[str, ParameterConstraints]] = None, debug: bool = True, callbacks: Optional[List[Callable[[List[ShardingOption]], List[ShardingOption]]]] = None)¶

根据提供的切片器、拓扑结构和约束条件，为具有可切片参数的给定模块提供优化的切片计划。

参数:

topology (Optional[Topology]) – 当前进程组的拓扑结构。
batch_size (Optional[int]) – 模型的批处理大小。
enumerator (Optional[Enumerator]) – 要使用的枚举器
storage_reservation (Optional[StorageReservation]) – 要使用的存储预留
proposer (Optional[Union[Proposer, List[Proposer]]]) – 要使用的建议者
partitioner (Optional[Partitioner]) – 要使用的分区器
performance_model (Optional[PerfModel]) – 要使用的性能模型
stats (Optional[Union[Stats, List[Stats]]]) – 要使用的统计信息
constraints (Optional[Dict[str, ParameterConstraints]]) – 每个表的切片约束。
debug (bool) – 是否打印调试信息。

示例

ebc = EmbeddingBagCollection(tables=eb_configs, device=torch.device("meta"))
planner = EmbeddingShardingPlanner()
plan = planner.plan(
    module=ebc,
    sharders=[EmbeddingBagCollectionSharder()],
)

collective_plan(module: Module, sharders: Optional[List[ModuleSharder[Module]]] = None, pg: Optional[ProcessGroup] = None) → ShardingPlan¶

在 rank 0 上调用 self.plan(…) 并广播

参数:

module (nn.Module) – 要分片的模块。
sharders (Optional[List[ModuleSharder[nn.Module]]]) – 用于分片的 sharders
pg (Optional[dist.ProcessGroup]) – 用于集体操作的进程组

返回值:

模块的分片计划。

返回类型:

ShardingPlan

plan(module: Module, sharders: List[ModuleSharder[Module]]) → ShardingPlan¶

根据提供的切片器、拓扑结构和约束条件，为具有可切片参数的给定模块提供优化的切片计划。

参数:

module (nn.Module) – 要分片的模块。
sharders (List[ModuleSharder[nn.Module]]) – 用于分片的 sharders。

返回值:

模块的分片计划。

返回类型:

ShardingPlan

class torchrec.distributed.planner.enumerators.EmbeddingEnumerator(topology: Topology, batch_size: int, constraints: Optional[Dict[str, ParameterConstraints]] = None, estimator: Optional[Union[ShardEstimator, List[ShardEstimator]]] = None, use_exact_enumerate_order: Optional[bool] = False)¶

针对给定的 nn.Module 生成嵌入分片选项，并考虑用户提供的约束。

参数:

topology (Topology) – 设备拓扑。
batch_size (int) – 批量大小。
constraints (Optional[Dict[str, ParameterConstraints]]) – 参数名称到提供的 ParameterConstraints 的字典。
estimator (Optional[Union[ShardEstimator, List[ShardEstimator]]]) – 分片性能估算器。
use_exact_enumerate_order (bool) – 是否以确切的 name_children 枚举顺序枚举可分片参数

enumerate(module: Module, sharders: List[ModuleSharder[Module]]) → List[ShardingOption]¶

根据模块和 sharders 生成相关分片选项。

参数:

module (nn.Module) – 要分片的模块。
sharders (List[ModuleSharder[nn.Module]]) – 为模块提供的 sharders。

返回值:

具有填充值的有效分片选项。

返回类型:

List[ShardingOption]

populate_estimates(sharding_options: List[ShardingOption]) → None¶: 请参阅类说明。

class torchrec.distributed.planner.partitioners.GreedyPerfPartitioner(sort_by: SortBy = SortBy.STORAGE, balance_modules: bool = False)¶

贪婪分区器。

参数:

sort_by (SortBy) – 按存储或性能降序排序分片选项（即，较大的表将被优先放置）。
balance_modules (bool) – 是否首先按模块排序，其中较小的模块将被优先排序。实际上，这将以平衡的方式将每个模块中的表放置在其中。

partition(proposal: List[ShardingOption], storage_constraint: Topology) → List[ShardingOption]¶

基于每个分片选项的 partition_by 属性，将分片选项放置在拓扑结构上。拓扑、存储和性能将在放置完成时更新。

参数:

proposal (List[ShardingOption]) – 已填充分片选项列表。
storage_constraint (Topology) – 设备拓扑。

返回值:

选定计划的分片选项列表。

返回类型:

List[ShardingOption]

示例

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.total = (1,2)
topology.devices[1].perf.total = (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.total = (1,2) + (3,4)
topology.devices[1].perf.total = (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).

class torchrec.distributed.planner.storage_reservations.HeuristicalStorageReservation(percentage: float, parameter_multiplier: float = 6.0, dense_tensor_estimate: Optional[int] = None)¶

使用启发式计算为要分片的模型保留存储空间。存储保留包括密集张量存储、KJT 存储以及总存储空间的额外百分比。

参数:

percentage (float) – 要保留的额外存储百分比，充当超出启发式存储计算的误差裕度。
parameter_multiplier (float) – 总参数存储的启发式乘数。
dense_tensor_estimate (Optional[int]) – 密集张量的存储估计，如果未提供，则使用默认的启发式估计。

class torchrec.distributed.planner.proposers.GreedyProposer(use_depth: bool = True, threshold: Optional[int] = None)¶

以贪婪方式提出分片计划。

按性能对每个可分片参数的分片选项进行排序。在每次迭代中，找到当前存储使用量最大的参数，并尝试其下一个分片选项。

参数:

use_depth (bool) – 启用后，fqn 的 sharding_options 将按 max(shard.perf.total) 进行排序，否则 sharding_options 将按 sum(shard.perf.total) 进行排序。
threshold (Optional[int]) – 早停阈值。如果指定，则当提案的连续性能等级比最佳性能等级差时，提案器停止提出提案。

feedback(partitionable: bool, plan: Optional[List[ShardingOption]] = None, perf_rating: Optional[float] = None, storage_constraint: Optional[Topology] = None) → None¶

向建议者提供反馈。

参数:

partitionable (bool) – 计划是否可分区。
plan (Optional[List[ShardingOption]]) – 要提供反馈的计划。
perf_rating (Optional[float]) – 计划的性能评分。
storage_constraint (Optional[Topology]) – 计划的存储约束。

load(search_space: List[ShardingOption], enumerator: Optional[Enumerator] = None) → None¶

将搜索空间加载到建议者中。

参数:

search_space (List[ShardingOption]) – 要加载的搜索空间。
enumerator (Enumerator) – 用于生成搜索空间的枚举器。

propose() → Optional[List[ShardingOption]]¶

建议一个分片计划。

返回值:: 建议的计划。
返回类型:: Optional[List[ShardingOption]]

class torchrec.distributed.planner.shard_estimators.EmbeddingPerfEstimator(topology: Topology, constraints: Optional[Dict[str, ParameterConstraints]] = None, is_inference: bool = False)¶

嵌入式壁钟时间性能评估器。此评估器估计给定分片选项的壁钟时间。

参数:

topology (Topology) – 设备拓扑。
constraints (Optional[Dict[str, ParameterConstraints]]) – 参数约束。
is_inference (bool) – 评估器是否用于推理。

estimate(sharding_options: List[ShardingOption], sharder_map: Optional[Dict[str, ModuleSharder[Module]]] = None) → None¶

估计给定分片选项的壁钟时间。

参数:

sharding_options (List[ShardingOption]) – 分片选项列表。
sharder_map (Optional[Dict[str, ModuleSharder[nn.Module]]]) – 分片器映射。

classmethod perf_func_emb_wall_time(shard_sizes: List[List[int]], compute_kernel: str, compute_device: str, sharding_type: str, batch_sizes: List[int], world_size: int, local_world_size: int, input_lengths: List[float], input_data_type_size: float, table_data_type_size: float, output_data_type_size: float, fwd_a2a_comm_data_type_size: float, bwd_a2a_comm_data_type_size: float, fwd_sr_comm_data_type_size: float, bwd_sr_comm_data_type_size: float, num_poolings: List[float], hbm_mem_bw: float, ddr_mem_bw: float, intra_host_bw: float, inter_host_bw: float, bwd_compute_multiplier: float, weighted_feature_bwd_compute_multiplier: float, is_pooled: bool, is_weighted: bool = False, caching_ratio: Optional[float] = None, is_inference: bool = False, prefetch_pipeline: bool = False, expected_cache_fetches: float = 0, uneven_sharding_perf_multiplier: float = 1.0) → List[Perf]¶

尝试将性能建模为相对壁钟时间的函数。

参数:

shard_sizes (List[List[int]]) – 每个分片的（local_rows, local_cols）列表。
compute_kernel (str) – 计算内核。
compute_device (str) – 计算设备。
sharding_type (str) – tw, rw, cw, twrw, dp。
batch_sizes (List[int]) – 每个输入特征的批次大小。
world_size (int) – 所有主机的设备数量。
local_world_size (int) – 每个主机的设备数量。
input_lengths (List[float]) – 输入查询特征的平均查找次数列表。
input_data_type_size (float) – 分布式数据并行输入的数据类型大小。
table_data_type_size (float) – 表格的数据类型大小。
output_data_type_size (float) – 输出嵌入的数据类型大小。
fwd_comm_data_type_size (float) – 正向通信期间分布式数据并行输入的数据类型大小。
bwd_comm_data_type_size (float) – 反向通信期间分布式数据并行输入的数据类型大小。
num_poolings (List[float]) – 每个样本的池化次数，通常为 1.0。
hbm_mem_bw (float) – 设备 HBM 的带宽。
ddr_mem_bw (float) – 系统 DDR 内存的带宽。
intra_host_bw (float) – 单个主机（如多个线程）内的带宽。
inter_host_bw (float) – 两个主机（如多台机器）之间的带宽。
is_pooled (bool) – 如果嵌入输出是池化的（例如 EmbeddingBag），则为 True；如果是非池化的/顺序的（例如 Embedding），则为 False。
is_weighted (bool = False) – 如果模块是 EBC 且加权，通常表示 id 得分列表特征。
is_inference (bool = False) – 如果计划用于推理。
caching_ratio (Optional[float] = None) – 缓存比率，用于确定设备带宽。
prefetch_pipeline (bool = False) – 是否启用预取管道。
expected_cache_fetches (float) – 全局批次中预期的缓存提取次数
uneven_sharding_perf_multiplier (float = 1.0) – 用于说明不均匀分片性能的乘数

返回值:

每个分片的性能列表。

返回类型:

List[float]

class torchrec.distributed.planner.shard_estimators.EmbeddingStorageEstimator(topology: Topology, constraints: Optional[Dict[str, ParameterConstraints]] = None, pipeline_type: PipelineType = PipelineType.NONE, run_embedding_at_peak_memory: bool = False, is_inference: bool = False)¶

嵌入存储使用量估算器

参数:

topology (Topology) – 设备拓扑。
constraints (Optional[Dict[str, ParameterConstraints]]) – 参数约束。
pipeline_type (PipelineType) – 管道类型（如果有）。将在内存估算期间确定输入复制因子。
run_embedding_at_peak_memory (bool) –
如果嵌入正向/反向将在 HBM 使用量达到峰值时执行。当设置为 TRUE 时，嵌入正向/反向期间的任何临时内存分配，只要输出_dist 之前的输出大小，将被计入 HBM 存储成本。否则，它们不会被计入，因为它们将被真实的内存峰值“隐藏”。

仅在为向后兼容性设置 pipeline_type 时生效（不会影响使用旧的与管道无关的公式的模型）

默认值为 false，因为这通常在 RecSys 中为 false，因为内存峰值发生在密集正向结束/密集反向开始时。
is_inference (bool) – 如果模型是推理模型。默认为 False。

estimate(sharding_options: List[ShardingOption], sharder_map: Optional[Dict[str, ModuleSharder[Module]]] = None) → None¶

估计每个分片选项的存储成本。

参数:

sharding_options (List[ShardingOption]) – 分片选项列表。
sharder_map (Optional[Dict[str, ModuleSharder[nn.Module]]]) – 模块类型到分片器的映射。

规划器¶

文档

Tutorials

Resources