From bfc96e1992ea10a338b2b240bb17be8976d7d038 Mon Sep 17 00:00:00 2001 From: Raahul Kalyaan Jakka Date: Tue, 18 Nov 2025 10:59:57 -0800 Subject: [PATCH] Adding "enable_optimizer_offloading" and "optimizer_D" to the required function signatures through the flow (#5056) Summary: X-link: https://github.com/meta-pytorch/torchrec/pull/3484 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 --- fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py | 1 + 1 file changed, 1 insertion(+) diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py index f63829845b..b8a5b335e3 100644 --- a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py +++ b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py @@ -186,6 +186,7 @@ def __init__( use_rowwise_bias_correction: bool = False, # For Adam use optimizer_state_dtypes: dict[str, SparseType] = {}, # noqa: B006 pg: Optional[dist.ProcessGroup] = None, + enable_optimizer_offloading: bool = False, ) -> None: super(SSDTableBatchedEmbeddingBags, self).__init__()