推理

TorchRec 提供了易于使用的 API，可以将自定义的 TorchRec 模型转换为优化的分布式推理模型，通过 eager 模块替换实现。

这个转换了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）。

Parameters:

• 模型 (torch.nn.Module) – 要进行量化处理的模型

• 量化映射 (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时所需的量化dtype。默认值为int8。

Returns:

量化模型

Return type:

torch.nn.Module

Example:

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模型分片，用于生成最优的推理模型，并且是分布式推理所必需的。

Parameters:

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

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

• 计算设备 (字符串) – 运行模型的设备，默认为“cuda”

• 分片设备 (字符串) – 运行分片的设备，默认为“meta”

• 分块器 (可选[列表[ModuleSharder[torch.nn.Module]]]) – 用于分块量化模型的分块器，默认为QuantEmbeddingBagCollectionSharder, QuantEmbeddingCollectionSharder, QuantFeatureProcessedEmbeddingBagCollectionSharder。

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

• 约束条件 (可选[字典[字符串, 参数约束条件]]) – 用于分片的约束条件，默认为None，这将实现默认约束条件，其中QuantEmbeddingBagCollection按行式进行分片

Returns:

分片模型和分片计划

Return type:

Tuple[torch.nn.Module, ShardingPlan]

Example::

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,

)

量化模型 = 量化推理模型（module） 分块模型，_ = 分块量化模型（量化模型）