tf-models-nightly 2.17.0.dev20240707__py2.py3-none-any.whl → 2.17.0.dev20240709__py2.py3-none-any.whl

official/recommendation/ranking/configs/config.py

@@ -112,6 +112,11 @@ class ModelConfig(hyperparams.Config):
       module
     dcn_use_bias: Flag to determine whether to use bias for the dcn interaction
       module
+    use_partial_tpu_embedding: Flag to determine whether to use partial tpu
+      embedding layer or not.
+    max_ids_per_chip_per_sample: Maximum number of ids per chip per sample.
+    max_ids_per_table: Maximum number of ids per table.
+    max_unique_ids_per_table: Maximum number of unique ids per table.
   """
   num_dense_features: int = 13
   vocab_sizes: List[int] = dataclasses.field(default_factory=list)
@@ -128,6 +133,10 @@ class ModelConfig(hyperparams.Config):
   dcn_kernel_initializer: str = 'truncated_normal'
   dcn_bias_initializer: str = 'zeros'
   dcn_use_bias: bool = True
+  use_partial_tpu_embedding: bool = True
+  max_ids_per_chip_per_sample: int | None = None
+  max_ids_per_table: Union[int, List[int]] | None = None
+  max_unique_ids_per_table: Union[int, List[int]] | None = None
 
 
 @dataclasses.dataclass
@@ -424,6 +433,7 @@ def dlrm_dcn_v2_criteo_tb_config() -> Config:
               dcn_use_bias=True,
               concat_dense=False,
               use_multi_hot=True,
+              use_partial_tpu_embedding=False,
               multi_hot_sizes=multi_hot_sizes,
               ),
           loss=Loss(label_smoothing=0.0),

official/recommendation/ranking/task.py

@@ -15,7 +15,7 @@
 """Task for the Ranking model."""
 
 import math
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Tuple
 
 import tensorflow as tf, tf_keras
 import tensorflow_recommenders as tfrs
@@ -35,8 +35,14 @@ def _get_tpu_embedding_feature_config(
     vocab_sizes: List[int],
     embedding_dim: Union[int, List[int]],
     table_name_prefix: str = 'embedding_table',
-    batch_size: Optional[int] = None
-) -> Dict[str, tf.tpu.experimental.embedding.FeatureConfig]:
+    batch_size: Optional[int] = None,
+    max_ids_per_chip_per_sample: Optional[int] = None,
+    max_ids_per_table: Optional[Union[int, List[int]]] = None,
+    max_unique_ids_per_table: Optional[Union[int, List[int]]] = None,
+) -> Tuple[
+    Dict[str, tf.tpu.experimental.embedding.FeatureConfig],
+    Optional[tf.tpu.experimental.embedding.SparseCoreEmbeddingConfig],
+]:
   """Returns TPU embedding feature config.
 
   i'th table config will have vocab size of vocab_sizes[i] and embedding
@@ -47,6 +53,11 @@ def _get_tpu_embedding_feature_config(
     embedding_dim: An integer or a list of embedding table dimensions.
     table_name_prefix: a prefix for embedding tables.
     batch_size: Per-replica batch size.
+    max_ids_per_chip_per_sample: Maximum number of embedding ids per chip per
+      sample.
+    max_ids_per_table: Maximum number of embedding ids per table.
+    max_unique_ids_per_table: Maximum number of unique embedding ids per table.
+
   Returns:
     A dictionary of feature_name, FeatureConfig pairs.
   """
@@ -54,14 +65,49 @@ def _get_tpu_embedding_feature_config(
     if len(vocab_sizes) != len(embedding_dim):
       raise ValueError(
           f'length of vocab_sizes: {len(vocab_sizes)} is not equal to the '
-          f'length of embedding_dim: {len(embedding_dim)}')
+          f'length of embedding_dim: {len(embedding_dim)}'
+      )
   elif isinstance(embedding_dim, int):
     embedding_dim = [embedding_dim] * len(vocab_sizes)
   else:
-    raise ValueError('embedding_dim is not either a list or an int, got '
-                     f'{type(embedding_dim)}')
+    raise ValueError(
+        'embedding_dim is not either a list or an int, got '
+        f'{type(embedding_dim)}'
+    )
+
+  if isinstance(max_ids_per_table, List):
+    if len(vocab_sizes) != len(max_ids_per_table):
+      raise ValueError(
+          f'length of vocab_sizes: {len(vocab_sizes)} is not equal to the '
+          f'length of max_ids_per_table: {len(max_ids_per_table)}'
+      )
+  elif isinstance(max_ids_per_table, int):
+    max_ids_per_table = [max_ids_per_table] * len(vocab_sizes)
+  elif max_ids_per_table is not None:
+    raise ValueError(
+        'max_ids_per_table is not either a list or an int or None, got '
+        f'{type(max_ids_per_table)}'
+    )
+
+  if isinstance(max_unique_ids_per_table, List):
+    if len(vocab_sizes) != len(max_unique_ids_per_table):
+      raise ValueError(
+          f'length of vocab_sizes: {len(vocab_sizes)} is not equal to the '
+          'length of max_unique_ids_per_table: '
+          f'{len(max_unique_ids_per_table)}'
+      )
+  elif isinstance(max_unique_ids_per_table, int):
+    max_unique_ids_per_table = [max_unique_ids_per_table] * len(vocab_sizes)
+  elif max_unique_ids_per_table is not None:
+    raise ValueError(
+        'max_unique_ids_per_table is not either a list or an int or None, '
+        f'got {type(max_unique_ids_per_table)}'
+    )
 
   feature_config = {}
+  sparsecore_config = None
+  max_ids_per_table_dict = {}
+  max_unique_ids_per_table_dict = {}
 
   for i, vocab_size in enumerate(vocab_sizes):
     table_config = tf.tpu.experimental.embedding.TableConfig(
@@ -69,15 +115,35 @@ def _get_tpu_embedding_feature_config(
         dim=embedding_dim[i],
         combiner='mean',
         initializer=tf.initializers.TruncatedNormal(
-            mean=0.0, stddev=1 / math.sqrt(embedding_dim[i])),
-        name=table_name_prefix + '_%02d' % i)
+            mean=0.0, stddev=1 / math.sqrt(embedding_dim[i])
+        ),
+        name=table_name_prefix + '_%02d' % i,
+    )
     feature_config[str(i)] = tf.tpu.experimental.embedding.FeatureConfig(
         name=str(i),
         table=table_config,
         output_shape=[batch_size] if batch_size else None,
     )
+    if max_ids_per_table:
+      max_ids_per_table_dict[str(table_name_prefix + '_%02d' % i)] = (
+          max_ids_per_table[i]
+      )
+    if max_unique_ids_per_table:
+      max_unique_ids_per_table_dict[str(table_name_prefix + '_%02d' % i)] = (
+          max_unique_ids_per_table[i]
+      )
 
-  return feature_config
+  if all((max_ids_per_chip_per_sample, max_ids_per_table,
+          max_unique_ids_per_table)):
+    sparsecore_config = tf.tpu.experimental.embedding.SparseCoreEmbeddingConfig(
+        disable_table_stacking=False,
+        max_ids_per_chip_per_sample=max_ids_per_chip_per_sample,
+        max_ids_per_table=max_ids_per_table_dict,
+        max_unique_ids_per_table=max_unique_ids_per_table_dict,
+        allow_id_dropping=False,
+    )
+
+  return feature_config, sparsecore_config
 
 
 class RankingTask(base_task.Task):
@@ -173,25 +239,33 @@ class RankingTask(base_task.Task):
           decay_start_steps=dense_lr_config.decay_start_steps)
       dense_optimizer.learning_rate = dense_lr_callable
 
-    feature_config = _get_tpu_embedding_feature_config(
-        embedding_dim=self.task_config.model.embedding_dim,
-        vocab_sizes=self.task_config.model.vocab_sizes,
-        batch_size=self.task_config.train_data.global_batch_size
-        // tf.distribute.get_strategy().num_replicas_in_sync,
+    feature_config, sparse_core_embedding_config = (
+        _get_tpu_embedding_feature_config(
+            embedding_dim=self.task_config.model.embedding_dim,
+            vocab_sizes=self.task_config.model.vocab_sizes,
+            batch_size=self.task_config.train_data.global_batch_size
+            // tf.distribute.get_strategy().num_replicas_in_sync,
+            max_ids_per_chip_per_sample=self.task_config.model.max_ids_per_chip_per_sample,
+            max_ids_per_table=self.task_config.model.max_ids_per_table,
+            max_unique_ids_per_table=self.task_config.model.max_unique_ids_per_table,
+        )
     )
 
-    if self.task_config.model.use_multi_hot:
-      embedding_layer = tfrs.layers.embedding.tpu_embedding_layer.TPUEmbedding(
+    # to work around PartialTPUEmbedding issue in v5p and to enable multi hot
+    # features
+    if self.task_config.model.use_partial_tpu_embedding:
+      embedding_layer = tfrs.experimental.layers.embedding.PartialTPUEmbedding(
           feature_config=feature_config,
           optimizer=embedding_optimizer,
           pipeline_execution_with_tensor_core=self.trainer_config.pipeline_sparse_and_dense_execution,
+          size_threshold=self.task_config.model.size_threshold,
       )
     else:
-      embedding_layer = tfrs.experimental.layers.embedding.PartialTPUEmbedding(
+      embedding_layer = tfrs.layers.embedding.tpu_embedding_layer.TPUEmbedding(
           feature_config=feature_config,
           optimizer=embedding_optimizer,
           pipeline_execution_with_tensor_core=self.trainer_config.pipeline_sparse_and_dense_execution,
-          size_threshold=self.task_config.model.size_threshold,
+          sparse_core_embedding_config=sparse_core_embedding_config,
       )
 
     if self.task_config.model.interaction == 'dot':

official/recommendation/ranking/train_test.py

@@ -35,6 +35,7 @@ def _get_params_override(vocab_sizes,
                          concat_dense=True,
                          dcn_num_layers=3,
                          dcn_low_rank_dim=64,
+                         use_partial_tpu_embedding=True,
                          ):
   # Update `data_dir` if `synthetic_data=False`.
   data_dir = ''
@@ -53,6 +54,7 @@ def _get_params_override(vocab_sizes,
               'dcn_num_layers': dcn_num_layers,
               'dcn_low_rank_dim': dcn_low_rank_dim,
               'use_multi_hot': use_multi_hot,
+              'use_partial_tpu_embedding': use_partial_tpu_embedding,
               'multi_hot_sizes': multi_hot_sizes,
           },
           'train_data': {
@@ -90,10 +92,20 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
     super().tearDown()
 
   @parameterized.named_parameters(
-      ('DlrmOneDeviceCTL', 'one_device', 'dot', True, True, 3, False),
-      ('DlrmOneDevice', 'one_device', 'dot', False, True, 3, False),
-      ('DcnOneDeviceCTL', 'one_device', 'cross', True, True, 3, False),
-      ('DcnOneDevice', 'one_device', 'cross', False, True, 3, False),
+      ('DlrmOneDeviceCTL', 'one_device', 'dot', True, True, 3, 64, False, True),
+      ('DlrmOneDevice', 'one_device', 'dot', False, True, 3, 64, False, True),
+      (
+          'DcnOneDeviceCTL',
+          'one_device',
+          'cross',
+          True,
+          True,
+          3,
+          64,
+          False,
+          True,
+      ),
+      ('DcnOneDevice', 'one_device', 'cross', False, True, 3, 64, False, True),
       (
           'DlrmDcnV2OneDeviceCTL',
           'one_device',
@@ -102,7 +114,8 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
       (
           'DlrmDcnV2OneDevice',
@@ -112,12 +125,13 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
-      ('DlrmTPUCTL', 'tpu', 'dot', True, True, 3, 64, False),
-      ('DlrmTPU', 'tpu', 'dot', False, True, 3, 64, False),
-      ('DcnTPUCTL', 'tpu', 'cross', True, True, 3, 64, False),
-      ('DcnTPU', 'tpu', 'cross', False, True, 3, 64, False),
+      ('DlrmTPUCTL', 'tpu', 'dot', True, True, 3, 64, False, True),
+      ('DlrmTPU', 'tpu', 'dot', False, True, 3, 64, False, True),
+      ('DcnTPUCTL', 'tpu', 'cross', True, True, 3, 64, False, True),
+      ('DcnTPU', 'tpu', 'cross', False, True, 3, 64, False, True),
       (
           'DlrmDcnV2TPUCTL',
           'tpu',
@@ -126,7 +140,8 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
       (
           'DlrmDcnV2TPU',
@@ -136,12 +151,13 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
-      ('DlrmMirroredCTL', 'Mirrored', 'dot', True, True, 3, 64, False),
-      ('DlrmMirrored', 'Mirrored', 'dot', False, True, 3, 64, False),
-      ('DcnMirroredCTL', 'Mirrored', 'cross', True, True, 3, 64, False),
-      ('DcnMirrored', 'Mirrored', 'cross', False, True, 3, 64, False),
+      ('DlrmMirroredCTL', 'Mirrored', 'dot', True, True, 3, 64, False, True),
+      ('DlrmMirrored', 'Mirrored', 'dot', False, True, 3, 64, False, True),
+      ('DcnMirroredCTL', 'Mirrored', 'cross', True, True, 3, 64, False, True),
+      ('DcnMirrored', 'Mirrored', 'cross', False, True, 3, 64, False, True),
       (
           'DlrmDcnV2MirroredCTL',
           'Mirrored',
@@ -150,7 +166,8 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
       (
           'DlrmDcnV2Mirrored',
@@ -160,7 +177,8 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
   )
   def testTrainEval(
@@ -172,6 +190,7 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
       dcn_num_layers=3,
       dcn_low_rank_dim=64,
       use_multi_hot=False,
+      use_partial_tpu_embedding=True,
   ):
     # Set up simple trainer with synthetic data.
     # By default the mode must be `train_and_eval`.
@@ -190,23 +209,45 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
         concat_dense=concat_dense,
         dcn_num_layers=dcn_num_layers,
         dcn_low_rank_dim=dcn_low_rank_dim,
+        use_partial_tpu_embedding=use_partial_tpu_embedding,
     )
+
     train.main('unused_args')
     self.assertNotEmpty(
         tf.io.gfile.glob(os.path.join(self._model_dir, 'params.yaml'))
     )
 
   @parameterized.named_parameters(
-      ('DlrmTPUCTL', 'tpu', 'dot', True, True, 3, 64, False),
-      ('DlrmTPU', 'tpu', 'dot', False, True, 3, 64, False),
-      ('DcnTPUCTL', 'tpu', 'cross', True, True, 3, 64, False),
-      ('DcnTPU', 'tpu', 'cross', False, True, 3, 64, False),
-      ('DlrmDcnV2TPUCTL', 'tpu', 'multi_layer_dcn', True, False, 3, 64, True),
-      ('DlrmDcnV2TPU', 'tpu', 'multi_layer_dcn', False, False, 3, 64, True),
-      ('DlrmMirroredCTL', 'Mirrored', 'dot', True, True, 3, 64, False),
-      ('DlrmMirrored', 'Mirrored', 'dot', False, True, 3, 64, False),
-      ('DcnMirroredCTL', 'Mirrored', 'cross', True, True, 3, 64, False),
-      ('DcnMirrored', 'Mirrored', 'cross', False, True, 3, 64, False),
+      ('DlrmTPUCTL', 'tpu', 'dot', True, True, 3, 64, False, True),
+      ('DlrmTPU', 'tpu', 'dot', False, True, 3, 64, False, True),
+      ('DcnTPUCTL', 'tpu', 'cross', True, True, 3, 64, False, True),
+      ('DcnTPU', 'tpu', 'cross', False, True, 3, 64, False, True),
+      (
+          'DlrmDcnV2TPUCTL',
+          'tpu',
+          'multi_layer_dcn',
+          True,
+          False,
+          3,
+          64,
+          True,
+          False,
+      ),
+      (
+          'DlrmDcnV2TPU',
+          'tpu',
+          'multi_layer_dcn',
+          False,
+          False,
+          3,
+          64,
+          True,
+          False,
+      ),
+      ('DlrmMirroredCTL', 'Mirrored', 'dot', True, True, 3, 64, False, True),
+      ('DlrmMirrored', 'Mirrored', 'dot', False, True, 3, 64, False, True),
+      ('DcnMirroredCTL', 'Mirrored', 'cross', True, True, 3, 64, False, True),
+      ('DcnMirrored', 'Mirrored', 'cross', False, True, 3, 64, False, True),
       (
           'DlrmDcnV2MirroredCTL',
           'Mirrored',
@@ -215,7 +256,8 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
       (
           'DlrmDcnV2Mirrored',
@@ -225,7 +267,8 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
           False,
           3,
           64,
-          True
+          True,
+          False,
       ),
   )
   def testTrainThenEval(
@@ -237,6 +280,7 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
       dcn_num_layers=3,
       dcn_low_rank_dim=64,
       use_multi_hot=False,
+      use_partial_tpu_embedding=True,
   ):
     # Set up simple trainer with synthetic data.
     vocab_sizes = [40, 12, 11, 13]
@@ -252,6 +296,7 @@ class TrainTest(parameterized.TestCase, tf.test.TestCase):
         dcn_num_layers=dcn_num_layers,
         dcn_low_rank_dim=dcn_low_rank_dim,
         use_multi_hot=use_multi_hot,
+        use_partial_tpu_embedding=use_partial_tpu_embedding,
     )
 
     default_mode = FLAGS.mode

{tf_models_nightly-2.17.0.dev20240707.dist-info → tf_models_nightly-2.17.0.dev20240709.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.17.0.dev20240707
+Version: 2.17.0.dev20240709
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.17.0.dev20240707.dist-info → tf_models_nightly-2.17.0.dev20240709.dist-info}/RECORD

@@ -877,12 +877,12 @@ official/recommendation/popen_helper.py,sha256=TMWMwsW1DF15YCJ0RG9pE3wsL6njLu8Ed
 official/recommendation/stat_utils.py,sha256=BjWRO2jzmAJyqeRUw0hMHhQqlyDNrdNvQJKAmbDJ4Rc,3076
 official/recommendation/ranking/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/recommendation/ranking/common.py,sha256=fUpBln1auJWwvX5BJIYJYxRQDzXPqCzbS6NYAcn0KaQ,3998
-official/recommendation/ranking/task.py,sha256=ovrPMW7uaLNGjUO7AFXjXpNUHN1MZQyLCq3M5bEXc2k,10850
+official/recommendation/ranking/task.py,sha256=NLoU4ZrCzrxfqg3h-xk_S8kVN1ZedZ-PKjdULmwusQU,13908
 official/recommendation/ranking/task_test.py,sha256=vPN_5oq1tWF3r7GuRTuAppKSF_mP3qoGFzGRoX09ylw,2891
 official/recommendation/ranking/train.py,sha256=_7zC2SVsOOOBN--If0XlWw4gtdaEeTT04PbCQRSzieo,6604
-official/recommendation/ranking/train_test.py,sha256=VlB74V7nPl88TOjLMVqw1w_i-HBTNvvO3X7R8N9y5CA,7839
+official/recommendation/ranking/train_test.py,sha256=n96JN-gI5wj4cbZNJCwQ9XLViKt0ZAg9f9dtaTpAdV8,8756
 official/recommendation/ranking/configs/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
-official/recommendation/ranking/configs/config.py,sha256=C6LHjDsH7cEGj--dHB3l8TBSPfmqmsspVnck4X1ogJ8,14278
+official/recommendation/ranking/configs/config.py,sha256=03KupVFJ9QeslF81Ec7n4VyPn-bVvyZaEd10GQWomVA,14844
 official/recommendation/ranking/configs/config_test.py,sha256=4_W0YUwJ2q0D8HdPdIvpJcAsaySjb_pDMWj26G7T1Ws,1474
 official/recommendation/ranking/data/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/recommendation/ranking/data/data_pipeline.py,sha256=5GSF1naOl6GdqOVin7_Yx_hprmzaq0WAImHqDxysAoI,7329
@@ -1212,9 +1212,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240707.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240707.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240707.dist-info/METADATA,sha256=VdAhQ49dDpv8r-ycZKtDDrbnnqKvxb3o1jVMmma2JUw,1432
-tf_models_nightly-2.17.0.dev20240707.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240707.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240707.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240709.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240709.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240709.dist-info/METADATA,sha256=8vfe5XTduG-zExj5iEl3JbJxCIa5_Qo0J7Fuc0TAD6U,1432
+tf_models_nightly-2.17.0.dev20240709.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240709.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240709.dist-info/RECORD,,