Initialization
This API is an open-source API of the TorchRec and is not an external API of the Rec SDK Torch. This section describes the parameter ranges supported by the TorchRec APIs called when the Rec SDK Torch is used.
Function
Creates a table planner to search for the most appropriate table plan.
Prototype
1 2 | class EmbeddingShardingPlanner: def __init__(**kwargs): |
Parameters
Parameter |
Data Type |
Mandatory/Optional |
Description |
|---|---|---|---|
topology |
Topology |
Optional |
For details, see the value range of Topology (TorchRec). This parameter is mandatory when NPU devices are used. |
constraints |
Dict[str, ParameterConstraints] |
Optional |
For details, see the value range of ParameterConstraints (TorchRec). This parameter is mandatory when NPU devices are used. |
batch_size |
int |
Optional |
Value range: [1, 1000000] |
enumerator |
torchrec.distributed.planner.types.Enumerator |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
storage_reservation |
torchrec.distributed.planner.types.StorageReservation |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
proposer |
torchrec.distributed.planner.types.Proposer |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
partitioner |
torchrec.distributed.planner.types.Partitioner |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
performance_model |
torchrec.distributed.planner.types.PerfModel |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
stats |
torchrec.distributed.planner.types.Stats |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
debug |
bool |
Optional |
When the NPU device is used, only the default value True is supported. User-defined values are not supported. |
callbacks |
List[Callable] |
Optional |
When the NPU device is used, only the default value None is supported. User-defined values are not supported. |
Sample
1 2 3 4 5 | from torchrec.distributed.planner import EmbeddingShardingPlanner planner = EmbeddingShardingPlanner( topology=topology, constraints=constraints, ) |
See Also
For details about the API call sequence and example, see Porting and Training.