directly pass update_util as int flag without syncing iter (#2695)

Summary:

as title, this change will eliminate device to host sync for the iter buffer during each iteration, which achieves better performance and avoid pottential bottleneck due to the sync point

this change covers all frontend usage of remap util which covers both ITEP & CEL

Reviewed By: dstaay-fb, sryap

Differential Revision: D68466136

Author: Fei Yu

torchrec/modules/itep_embedding_modules.py

@@ -75,6 +75,7 @@ def forward(
 
         features = self._itep_module(features, self._iter.item())
         pooled_embeddings = self._embedding_bag_collection(features)
+
         self._iter += 1
 
         return pooled_embeddings

torchrec/modules/itep_modules.py

@@ -464,13 +464,13 @@ def forward(
             feature_offsets,
         ) = self.get_remap_info(sparse_features)
 
-        update_utils: bool = (
+        update_util: bool = (
             (cur_iter < 10)
             or (cur_iter < 100 and (cur_iter + 1) % 19 == 0)
             or ((cur_iter + 1) % 39 == 0)
         )
         full_values_list = None
-        if update_utils and sparse_features.variable_stride_per_key():
+        if update_util and sparse_features.variable_stride_per_key():
             if sparse_features.inverse_indices_or_none() is not None:
                 # full util update mode require reconstructing original input indicies from VBE input
                 full_values_list = self.get_full_values_list(sparse_features)
@@ -490,6 +490,7 @@ def forward(
             self.row_util,
             self.buffer_offsets,
             full_values_list=full_values_list,
+            update_util=update_util,
         )
 
         sparse_features._values = remapped_values