推理

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, sharding_plan: Optional[ShardingPlan] = 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` 将按表进行分片

返回:

分片后的模型和分片计划

返回类型:

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)