推理

TorchRec 提供了易于使用的 API，通过即时模块交换，将已编写的 TorchRec 模型转换为优化的分布式推理模型。

这会将模型中的 TorchRec 模块（如 EmbeddingBagCollection）转换为量化、分片版本，该版本可以使用 torch.fx 和 TorchScript 进行编译，以便在 C++ 环境中进行推理。

预期用途是在模型上调用 quantize_inference_model，然后调用 shard_quant_model。

torchrec.inference.modules.quantize_inference_model(model: Module, quantization_mapping: Optional[Dict[str, Type[Module]]] = None, per_table_weight_dtype: Optional[Dict[str, dtype]] = None, fp_weight_dtype: dtype = torch.int8, quantization_dtype: dtype = torch.int8, output_dtype: dtype = torch.float32) → Module

量化模型，通过模块交换将 TorchRec 训练模块与其量化对应模块进行替换（例如，EmbeddingBagCollection -> QuantEmbeddingBagCollection）。

参数:

• model (torch.nn.Module) – 要量化的模型

• quantization_mapping (Optional[Dict[str, Type[torch.nn.Module]]]) – 从原始模块类型到量化模块类型的映射。如果未提供，将使用默认映射：（EmbeddingBagCollection -> QuantEmbeddingBagCollection，EmbeddingCollection -> QuantEmbeddingCollection）。

• per_table_weight_dtype (Optional[Dict[str, torch.dtype]]) – 从表名到权重数据类型的映射。如果未提供，将使用默认量化数据类型 (int8)。

• fp_weight_dtype (torch.dtype) – 如果使用 FeatureProcessedEmbeddingBagCollection，则为特征处理器权重所需的量化数据类型。默认为 int8。

返回:

量化后的模型

返回类型:

torch.nn.Module

示例

ebc = EmbeddingBagCollection(tables=eb_configs, device=torch.device("meta"))

module = DLRMPredictModule(
    embedding_bag_collection=ebc,
    dense_in_features=self.model_config.dense_in_features,
    dense_arch_layer_sizes=self.model_config.dense_arch_layer_sizes,
    over_arch_layer_sizes=self.model_config.over_arch_layer_sizes,
    id_list_features_keys=self.model_config.id_list_features_keys,
    dense_device=device,
)

quant_model = quantize_inference_model(module)

torchrec.inference.modules.shard_quant_model(model: Module, world_size: int = 1, compute_device: str = 'cuda', sharding_device: str = 'meta', sharders: Optional[List[ModuleSharder[Module]]] = None, device_memory_size: Optional[int] = None, constraints: Optional[Dict[str, ParameterConstraints]] = None, ddr_cap: Optional[int] = None) → Tuple[Module, ShardingPlan]

对量化的 TorchRec 模型进行分片，用于生成最佳推理模型，并且对于分布式推理是必要的。

参数:

• model (torch.nn.Module) – 要分片的量化模型

• world_size (int) – 用于分片模型的设备数量，默认为 1

• compute_device (str) – 运行模型的设备，默认为 “cuda”

• sharding_device (str) – 运行分片的设备，默认为 “meta”

• sharders (Optional[List[ModuleSharder[torch.nn.Module]]]) – 用于分片量化模型的分片器，默认为 QuantEmbeddingBagCollectionSharder、QuantEmbeddingCollectionSharder、QuantFeatureProcessedEmbeddingBagCollectionSharder。

• device_memory_size (Optional[int]) – cuda 设备的内存限制，默认为 None

• constraints (Optional[Dict[str, ParameterConstraints]]) – 用于分片的约束，默认为 None，然后将实现默认约束，其中 QuantEmbeddingBagCollection 将按 TableWise 分片

返回:

分片模型和分片计划

返回类型:

Tuple[torch.nn.Module, ShardingPlan]

示例：

ebc = EmbeddingBagCollection(tables=eb_configs, device=torch.device(“meta”))

module = DLRMPredictModule(

embedding_bag_collection=ebc, dense_in_features=self.model_config.dense_in_features, dense_arch_layer_sizes=self.model_config.dense_arch_layer_sizes, over_arch_layer_sizes=self.model_config.over_arch_layer_sizes, id_list_features_keys=self.model_config.id_list_features_keys, dense_device=device,

)

quant_model = quantize_inference_model(module) sharded_model, _ = shard_quant_model(quant_model)