add virtual table eviction policy (#3172)

Summary:
Pull Request resolved: #3172

X-link: facebookresearch/FBGEMM#1498

X-link: pytorch/FBGEMM#4433

Add eviction policy to embedding config and also enable config in mvai model family

Reviewed By: duduyi2013, yixin94

Differential Revision: D75660955

fbshipit-source-id: e514f56a88b46f5000f8d54478531f7d4e739f21

Author: emlin

torchrec/distributed/batched_embedding_kernel.py

@@ -31,6 +31,7 @@
 import torch.distributed as dist
 from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
     BackendType,
+    EvictionPolicy,
     KVZCHParams,
 )
 from fbgemm_gpu.split_table_batched_embeddings_ops_inference import (
@@ -78,8 +79,14 @@
 )
 from torchrec.distributed.utils import append_prefix, none_throws
 from torchrec.modules.embedding_configs import (
+    CountBasedEvictionPolicy,
+    CountTimestampMixedEvictionPolicy,
     data_type_to_sparse_type,
+    FeatureL2NormBasedEvictionPolicy,
+    NoEvictionPolicy,
     pooling_type_to_pooling_mode,
+    TimestampBasedEvictionPolicy,
+    VirtualTableEvictionPolicy,
 )
 from torchrec.optim.fused import (
     EmptyFusedOptimizer,
@@ -201,6 +208,7 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]:
 def _populate_zero_collision_tbe_params(
     tbe_params: Dict[str, Any],
     sharded_local_buckets: List[Tuple[int, int, int]],
+    config: GroupedEmbeddingConfig,
 ) -> None:
     """
     Construct Zero Collision TBE params from config and fused params dict.
@@ -211,10 +219,77 @@ def _populate_zero_collision_tbe_params(
     ]
     bucket_sizes: List[int] = [size for _, _, size in sharded_local_buckets]
 
+    enabled = False
+    for table in config.embedding_tables:
+        if table.virtual_table_eviction_policy is not None and not isinstance(
+            table.virtual_table_eviction_policy, NoEvictionPolicy
+        ):
+            enabled = True
+    if enabled:
+        counter_thresholds = [0] * len(config.embedding_tables)
+        ttls_in_mins = [0] * len(config.embedding_tables)
+        counter_decay_rates = [0.0] * len(config.embedding_tables)
+        l2_weight_thresholds = [0.0] * len(config.embedding_tables)
+        eviction_strategy = -1
+        table_names = [table.name for table in config.embedding_tables]
+        for i, table in enumerate(config.embedding_tables):
+            policy_t = table.virtual_table_eviction_policy
+            if policy_t is not None:
+                if isinstance(policy_t, CountBasedEvictionPolicy):
+                    counter_thresholds[i] = policy_t.eviction_threshold
+                    counter_decay_rates[i] = policy_t.decay_rate
+                    if eviction_strategy == -1 or eviction_strategy == 1:
+                        eviction_strategy = 1
+                    else:
+                        raise ValueError(
+                            f"Do not support multiple eviction strategy in one tbe {eviction_strategy} and 1 for tables {table_names}"
+                        )
+                elif isinstance(policy_t, TimestampBasedEvictionPolicy):
+                    ttls_in_mins[i] = policy_t.eviction_ttl_mins
+                    if eviction_strategy == -1 or eviction_strategy == 0:
+                        eviction_strategy = 0
+                    else:
+                        raise ValueError(
+                            f"Do not support multiple eviction strategy in one tbe {eviction_strategy} and 0 for tables {table_names}"
+                        )
+                elif isinstance(policy_t, FeatureL2NormBasedEvictionPolicy):
+                    l2_weight_thresholds[i] = policy_t.eviction_threshold
+                    if eviction_strategy == -1 or eviction_strategy == 3:
+                        eviction_strategy = 3
+                    else:
+                        raise ValueError(
+                            f"Do not support multiple eviction strategy in one tbe {eviction_strategy} and 3 for tables {table_names}"
+                        )
+                elif isinstance(policy_t, CountTimestampMixedEvictionPolicy):
+                    counter_thresholds[i] = policy_t.eviction_threshold
+                    counter_decay_rates[i] = policy_t.decay_rate
+                    ttls_in_mins[i] = policy_t.eviction_ttl_mins
+                    if eviction_strategy == -1 or eviction_strategy == 2:
+                        eviction_strategy = 2
+                    else:
+                        raise ValueError(
+                            f"Do not support multiple eviction strategy in one tbe {eviction_strategy} and 2 for tables {table_names}"
+                        )
+                else:
+                    raise ValueError(
+                        f"Unsupported eviction policy {policy_t} for table {table.name}"
+                    )
+        eviction_policy = EvictionPolicy(
+            eviction_trigger_mode=2,  # 2 means mem_util based eviction
+            eviction_strategy=eviction_strategy,
+            counter_thresholds=counter_thresholds,
+            ttls_in_mins=ttls_in_mins,
+            counter_decay_rates=counter_decay_rates,
+            l2_weight_thresholds=l2_weight_thresholds,
+        )
+    else:
+        eviction_policy = None
+
     tbe_params["kv_zch_params"] = KVZCHParams(
         bucket_offsets=bucket_offsets,
         bucket_sizes=bucket_sizes,
         enable_optimizer_offloading=False,
+        eviction_policy=eviction_policy,
     )
 
 
@@ -1318,7 +1393,7 @@ def __init__(
                 self._config.embedding_tables, self._pg
             )
         )
-        _populate_zero_collision_tbe_params(ssd_tbe_params, self._bucket_spec)
+        _populate_zero_collision_tbe_params(ssd_tbe_params, self._bucket_spec, config)
         compute_kernel = config.embedding_tables[0].compute_kernel
         embedding_location = compute_kernel_to_embedding_location(compute_kernel)
 
@@ -2124,7 +2199,7 @@ def __init__(
                 self._config.embedding_tables, self._pg
             )
         )
-        _populate_zero_collision_tbe_params(ssd_tbe_params, self._bucket_spec)
+        _populate_zero_collision_tbe_params(ssd_tbe_params, self._bucket_spec, config)
         compute_kernel = config.embedding_tables[0].compute_kernel
         embedding_location = compute_kernel_to_embedding_location(compute_kernel)
 

torchrec/distributed/embedding.py

@@ -248,6 +248,7 @@ def create_sharding_infos_by_sharding_device_group(
                         weight_init_min=config.weight_init_min,
                         total_num_buckets=config.total_num_buckets,
                         use_virtual_table=config.use_virtual_table,
+                        virtual_table_eviction_policy=config.virtual_table_eviction_policy,
                     ),
                     param_sharding=parameter_sharding,
                     param=param,
@@ -613,6 +614,7 @@ def create_grouped_sharding_infos(
                             weight_init_min=config.weight_init_min,
                             total_num_buckets=config.total_num_buckets,
                             use_virtual_table=config.use_virtual_table,
+                            virtual_table_eviction_policy=config.virtual_table_eviction_policy,
                         ),
                         param_sharding=parameter_sharding,
                         param=param,

torchrec/distributed/embeddingbag.py

@@ -295,6 +295,7 @@ def create_sharding_infos_by_sharding_device_group(
                     ),
                     total_num_buckets=config.total_num_buckets,
                     use_virtual_table=config.use_virtual_table,
+                    virtual_table_eviction_policy=config.virtual_table_eviction_policy,
                 ),
                 param_sharding=parameter_sharding,
                 param=param,
@@ -693,6 +694,7 @@ def create_grouped_sharding_infos(
                     ),
                     total_num_buckets=config.total_num_buckets,
                     use_virtual_table=config.use_virtual_table,
+                    virtual_table_eviction_policy=config.virtual_table_eviction_policy,
                 ),
                 param_sharding=parameter_sharding,
                 param=param,

torchrec/distributed/quant_state.py

@@ -385,6 +385,9 @@ class WeightSpec:
     shard_offsets: List[int]  # shard offsets
     shard_sizes: List[int]  # shard sizes
     sharding_type: Optional[str]  # e.g. ShardingType.ROW_WISE.value=="row_wise"
+    virtual_table_dim_offsets: Optional[List[int]] = (
+        None  # for virtual table, weight dim offsets for quantization. e.g. [8, 264] for 256 dim tables, the first 8 elements are the metaheader
+    )
 
 
 def get_bucket_offsets_per_virtual_table(
@@ -504,6 +507,18 @@ def sharded_tbes_weights_spec(
                 tables = config.embedding_tables
                 for table_idx, table in enumerate(tables):
                     table_name: str = table.name
+                    table_dim_offsets: Optional[List[int]] = (
+                        None
+                        if not table.use_virtual_table
+                        else [0, table.embedding_dim]
+                    )
+                    if table.virtual_table_eviction_policy:
+                        table_dim_offsets = [
+                            table.virtual_table_eviction_policy.get_meta_header_len(),
+                            # pyre-ignore [16]
+                            table.virtual_table_eviction_policy.get_meta_header_len()
+                            + table.embedding_dim,
+                        ]
                     # pyre-ignore
                     table_metadata: ShardMetadata = table.local_metadata
                     local_rows = table.local_rows
@@ -577,6 +592,7 @@ def sharded_tbes_weights_spec(
                         shard_offsets=shard_offsets,
                         shard_sizes=shard_sizes,
                         sharding_type=sharding_type,
+                        virtual_table_dim_offsets=table_dim_offsets,
                     )
 
                     # We also need to populate weight_id tensor for vritual

torchrec/distributed/sharding/rw_sharding.py

@@ -219,6 +219,7 @@ def _shard(
                         num_embeddings_post_pruning=info.embedding_config.num_embeddings_post_pruning,
                         total_num_buckets=info.embedding_config.total_num_buckets,
                         use_virtual_table=info.embedding_config.use_virtual_table,
+                        virtual_table_eviction_policy=info.embedding_config.virtual_table_eviction_policy,
                     )
                 )
         return tables_per_rank

torchrec/distributed/test_utils/infer_utils.py

@@ -77,6 +77,7 @@
     dtype_to_data_type,
     EmbeddingBagConfig,
     QuantConfig,
+    VirtualTableEvictionPolicy,
 )
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import PositionWeightedModuleCollection
@@ -642,6 +643,7 @@ def create_test_model(
     constraints: Optional[Dict[str, ParameterConstraints]] = None,
     weight_dtype: torch.dtype = torch.qint8,
     pruning_dict: Optional[Dict[str, int]] = None,
+    virtual_table_eviction_policy: Optional[VirtualTableEvictionPolicy] = None,
 ) -> TestModelInfo:
     topology: Topology = Topology(
         world_size=world_size, compute_device=sparse_device.type
@@ -675,6 +677,8 @@ def create_test_model(
             embedding_dim=emb_dim,
             name="table_" + str(i),
             feature_names=["feature_" + str(i)],
+            use_virtual_table=True if virtual_table_eviction_policy else False,
+            virtual_table_eviction_policy=virtual_table_eviction_policy,
         )
         for i in range(mi.num_features)
     ]
@@ -685,6 +689,8 @@ def create_test_model(
             embedding_dim=emb_dim,
             name="weighted_table_" + str(i),
             feature_names=["weighted_feature_" + str(i)],
+            use_virtual_table=True if virtual_table_eviction_policy else False,
+            virtual_table_eviction_policy=virtual_table_eviction_policy,
         )
         for i in range(mi.num_weighted_features)
     ]

torchrec/distributed/tests/test_infer_shardings.py

@@ -69,6 +69,10 @@
 from torchrec.distributed.test_utils.test_model import ModelInput
 from torchrec.distributed.types import ShardingEnv, ShardingPlan
 from torchrec.fx import symbolic_trace
+from torchrec.modules.embedding_configs import (
+    dtype_to_data_type,
+    TimestampBasedEvictionPolicy,
+)
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import (
     FeatureProcessorsCollection,
@@ -357,6 +361,90 @@ def test_rw(self, weight_dtype: torch.dtype, device_type: str) -> None:
             ShardingType.ROW_WISE.value,
         )
 
+    @unittest.skipIf(
+        torch.cuda.device_count() <= 1,
+        "Not enough GPUs available",
+    )
+    # pyre-ignore
+    @given(
+        weight_dtype=st.sampled_from([torch.qint8, torch.quint4x2]),
+        device_type=st.sampled_from(["cuda", "cpu"]),
+    )
+    @settings(max_examples=4, deadline=None)
+    def test_rw_with_virtual_table_eviction(
+        self, weight_dtype: torch.dtype, device_type: str
+    ) -> None:
+        num_embeddings = 256
+        emb_dim = 16
+        world_size = 2
+        batch_size = 4
+        local_device = torch.device(f"{device_type}:0")
+        eviction_policy = TimestampBasedEvictionPolicy()
+        eviction_policy.init_metaheader_config(dtype_to_data_type(torch.float16))
+        mi = create_test_model(
+            num_embeddings,
+            emb_dim,
+            world_size,
+            batch_size,
+            dense_device=local_device,
+            sparse_device=local_device,
+            quant_state_dict_split_scale_bias=True,
+            weight_dtype=weight_dtype,
+            virtual_table_eviction_policy=eviction_policy,
+        )
+
+        non_sharded_model = mi.quant_model
+        num_emb_half = num_embeddings // 2
+        expected_shards = [
+            [
+                ((0, 0, num_emb_half, emb_dim), placement(device_type, 0, world_size)),
+                (
+                    (num_emb_half, 0, num_emb_half, emb_dim),
+                    placement(device_type, 1, world_size),
+                ),
+            ]
+        ]
+        sharded_model = shard_qebc(
+            mi,
+            sharding_type=ShardingType.ROW_WISE,
+            device=local_device,
+            expected_shards=expected_shards,
+        )
+        inputs = [
+            model_input_to_forward_args(inp.to(local_device))
+            for inp in prep_inputs(mi, world_size, batch_size, long_indices=False)
+        ]
+
+        sharded_model.load_state_dict(non_sharded_model.state_dict())
+
+        sharded_output = sharded_model(*inputs[0])
+        non_sharded_output = non_sharded_model(*inputs[0])
+        assert_close(sharded_output, non_sharded_output)
+
+        weights_spec: Dict[str, WeightSpec] = sharded_tbes_weights_spec(sharded_model)
+        assert_weight_spec(
+            weights_spec,
+            expected_shards,
+            "_module.sparse.ebc",
+            "embedding_bags",
+            ["table_0"],
+            ShardingType.ROW_WISE.value,
+        )
+        print(weights_spec)
+        assert (
+            weights_spec[
+                "_module.sparse.ebc.tbes.0.0.table_0.weight"
+            ].virtual_table_dim_offsets
+            is not None
+        )
+        assert (
+            # pyre-ignore [16]
+            weights_spec[
+                "_module.sparse.ebc.tbes.0.0.table_0.weight"
+            ].virtual_table_dim_offsets[0]
+            == 8
+        )
+
     @unittest.skipIf(
         torch.cuda.device_count() <= 1,
         "Not enough GPUs available",