ParameterConstraints (TorchRec)

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

Query range of a table sharding plan.

Prototype

class ParameterConstraints:
    def __init__(**kwargs):

Parameters

Parameter	Data Type	Mandatory/Optional	Description
sharding_type	List[str]	Optional	Table sharding type. This parameter is mandatory when NPU devices are used. The options are as follows: "row_wise": Tables are sharded by row number. "data_parallel": One table copy is reserved for each rank. NOTE: Different table sharding types cannot be used together.
compute_kernels	List[str]	Optional	Kernel type. This parameter is mandatory when NPU devices are used. The options are as follows: "fused": Query in the combined table mode. This mode is used only when sharding_type is set to "row_wise". "dense": Query in the sharding table mode. This mode is used only when sharding_type is set to "data_parallel".
min_partition	List[int]	Optional	When an NPU is used, the value can only be the default value None and cannot be customized.
pooling_factors	List[float]	Optional	When an NPU is used, the value can only be the default value POOLING_FACTOR and cannot be customized.
num_poolings	List[float]	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
batch_sizes	List[int]	Optional	When an NPU is used, the value can only be the default value None and cannot be customized.
is_weighted	bool	Optional	When the NPU device is used, only the default value False is supported. User-defined values are not supported.
cache_params	torchrec.distributed.types.CacheParams	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
enforce_hbm	bool	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
stochastic_rounding	bool	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
bounds_check_mode	enum.IntEnum	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
feature_names	List[str]	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
output_dtype	Enum	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
device_group	str	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.
key_value_params	torchrec.distributed.types.KeyValueParams	Optional	When the NPU device is used, only the default value None is supported. User-defined values are not supported.

Sample

from torchrec.distributed.planner import ParameterConstraints
constraints = {
   "table0": ParameterConstraints(sharding_types=["row_wise"], compute_kernels=["fused"])
}

Parent topic: Table Splitting APIs