推理

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) → 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)