From b4f3ddada5b5ac335313dd6b996358ff48ba9a1d Mon Sep 17 00:00:00 2001 From: Raahul Kalyaan Jakka Date: Fri, 19 Dec 2025 13:00:27 -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 bff406aab0..52ab385914 100644 --- a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py +++ b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py @@ -187,6 +187,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__()