From f3b5fb22ecbbb6ce4355378a11025acbc1064413 Mon Sep 17 00:00:00 2001 From: Raahul Kalyaan Jakka Date: Mon, 27 Oct 2025 13:12:08 -0700 Subject: [PATCH] Changing Backend Tensor initialization (#3484) Summary: X-link: https://github.com/pytorch/FBGEMM/pull/5056 X-link: https://github.com/facebookresearch/FBGEMM/pull/2066 **Context:** Currently, we are enabling SSD optimizer offloading for the ssd tbe kernel **In this diff:** We retrieve the newly added parameters from the tbe config and pass it down to the tbe Differential Revision: D85353134 --- torchrec/distributed/batched_embedding_kernel.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/torchrec/distributed/batched_embedding_kernel.py b/torchrec/distributed/batched_embedding_kernel.py index 968d7aa15..2919f4684 100644 --- a/torchrec/distributed/batched_embedding_kernel.py +++ b/torchrec/distributed/batched_embedding_kernel.py @@ -246,6 +246,12 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]: ssd_tbe_params["kvzch_eviction_trigger_mode"] = fused_params.get( "kvzch_eviction_trigger_mode" ) + if "enable_optimizer_offloading" in fused_params: + ssd_tbe_params["enable_optimizer_offloading"] = fused_params.get( + "enable_optimizer_offloading" + ) + else: + ssd_tbe_params["enable_optimizer_offloading"] = False ssd_tbe_params["table_names"] = [table.name for table in config.embedding_tables]