fix eviction policy configuration (#3182)

Summary:
Pull Request resolved: #3182

l2 cache size should be used for eviction memory threshold
also, as optimizer offloading is landed, we need to enable it by default.

Reviewed By: kathyxuyy

Differential Revision: D78142858

fbshipit-source-id: 42f6c95e88e9497cdacb951dd8753b9018ccde44

Author: emlin

torchrec/distributed/batched_embedding_kernel.py

@@ -232,6 +232,7 @@ def _populate_zero_collision_tbe_params(
         l2_weight_thresholds = [0.0] * len(config.embedding_tables)
         eviction_strategy = -1
         table_names = [table.name for table in config.embedding_tables]
+        l2_cache_size = tbe_params["l2_cache_size"]
         for i, table in enumerate(config.embedding_tables):
             policy_t = table.virtual_table_eviction_policy
             if policy_t is not None:
@@ -276,6 +277,7 @@ def _populate_zero_collision_tbe_params(
                     )
         eviction_policy = EvictionPolicy(
             eviction_trigger_mode=2,  # 2 means mem_util based eviction
+            eviction_mem_threshold_gb=l2_cache_size,
             eviction_strategy=eviction_strategy,
             counter_thresholds=counter_thresholds,
             ttls_in_mins=ttls_in_mins,
@@ -288,7 +290,7 @@ def _populate_zero_collision_tbe_params(
     tbe_params["kv_zch_params"] = KVZCHParams(
         bucket_offsets=bucket_offsets,
         bucket_sizes=bucket_sizes,
-        enable_optimizer_offloading=False,
+        enable_optimizer_offloading=True,
         eviction_policy=eviction_policy,
     )