Handle Virtual Table sizing in planner

Summary:
Virtual Embedding Tables is a new feature in ZCH v.Next. These do not allocate the embedding table memory on init.
This allows us to have larger embedding tables and paired with Embedding Offloading kernels to DRAM and SSD.
Since the memory is kernel dependent, we check for the specific kernel and override the hbm and dram sizes accordingly.
https://docs.google.com/document/d/1NjtP01PSOHKwyRxAicV7wrdWPVI1NN9A3aPYuBaOpsg/edit?tab=t.0

Differential Revision: D74557713

Author: kausv

torchrec/distributed/planner/planners.py

@@ -422,6 +422,7 @@ def plan(
                     sharders=sharders,
                     debug=self._debug,
                 )
+            logger.info(f"Found sharding plan {sharding_plan}")
             return sharding_plan
         else:
             global_storage_capacity = reduce(

torchrec/distributed/planner/shard_estimators.py

@@ -13,6 +13,7 @@
 
 import torch
 import torchrec.optim as trec_optim
+from libfb.py.pyre import none_throws
 from torch import nn
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel
 from torchrec.distributed.planner.constants import (
@@ -41,6 +42,7 @@
 from torchrec.distributed.types import (
     CacheStatistics,
     CommOp,
+    KeyValueParams,
     ModuleSharder,
     PipelineType,
     ShardingType,
@@ -998,23 +1000,29 @@ def estimate(
                     if hasattr(sharder, "fused_params") and sharder.fused_params
                     else None
                 )
-
-            num_poolings = (
-                cast(List[float], self._constraints[sharding_option.name].num_poolings)
+            constraints: Optional[ParameterConstraints] = (
+                self._constraints.get(sharding_option.name, None)
                 if self._constraints
-                and self._constraints.get(sharding_option.name)
-                and self._constraints[sharding_option.name].num_poolings
+                else None
+            )
+            num_poolings = (
+                constraints.num_poolings
+                if constraints and constraints.num_poolings
                 else [1.0] * sharding_option.num_inputs
             )
             assert len(num_poolings) == sharding_option.num_inputs
             batch_sizes = (
-                cast(List[int], self._constraints[sharding_option.name].batch_sizes)
-                if self._constraints
-                and self._constraints.get(sharding_option.name)
-                and self._constraints[sharding_option.name].batch_sizes
+                constraints.batch_sizes
+                if constraints and constraints.batch_sizes
                 else [sharding_option.batch_size] * sharding_option.num_inputs
             )
 
+            key_value_params: Optional[KeyValueParams] = (
+                constraints.key_value_params
+                if constraints and constraints.key_value_params
+                else None
+            )
+
             # hardcoded as 8 bytes
             # input indices can be of int32, but in TBE they get converted to int64 anyway
             input_data_type_size = BIGINT_DTYPE
@@ -1057,6 +1065,7 @@ def estimate(
                 count_ephemeral_storage_cost=self._run_embedding_at_peak_memory,
                 is_inference=self._is_inference,
                 multipass_prefetch_max_pass=mpp_conf.num_passes if mpp_conf else None,
+                key_value_params=key_value_params,
             )
             for shard, storage in zip(sharding_option.shards, shard_storages):
                 shard.storage = storage
@@ -1125,6 +1134,7 @@ def calculate_shard_storages(
     count_ephemeral_storage_cost: bool = False,
     is_inference: bool = False,
     multipass_prefetch_max_pass: Optional[int] = None,
+    key_value_params: Optional[KeyValueParams] = None,
 ) -> List[Storage]:
     """
     Calculates estimated storage sizes for each sharded tensor, comprised of input,
@@ -1151,6 +1161,7 @@ def calculate_shard_storages(
         output_data_type_size (int): number of bytes of output data type.
         pipeline_type: PipelineType: pipeline type if for training.
         is_inference: bool, whether the model is for inference.
+        key_value_params (Optional[KeyValueParams]): fused params for SSD/DRAM KV cache.
 
     Returns:
         List[Storage]: storage object for each device in topology.
@@ -1184,13 +1195,6 @@ def calculate_shard_storages(
         # TODO(wangj): for ssd/dram kv, most likely we use absolute L1 cache size instead of caching ratio, as denominator is huge
         hbm_storage = round(ddr_storage * caching_ratio)
         table_cached = True
-    if compute_kernel in {
-        EmbeddingComputeKernel.KEY_VALUE.value,
-        EmbeddingComputeKernel.SSD_VIRTUAL_TABLE.value,
-        EmbeddingComputeKernel.DRAM_VIRTUAL_TABLE.value,
-    }:
-        # TODO(wangj): update this to the L2 cache usage and add SSD usage
-        ddr_storage = 0
 
     optimizer_class = getattr(tensor, "_optimizer_classes", [None])[0]
 
@@ -1212,6 +1216,36 @@ def calculate_shard_storages(
         is_inference=is_inference,
     )
 
+    if compute_kernel in {
+        EmbeddingComputeKernel.KEY_VALUE.value,
+        EmbeddingComputeKernel.SSD_VIRTUAL_TABLE.value,
+        EmbeddingComputeKernel.DRAM_VIRTUAL_TABLE.value,
+    }:
+        assert (
+            key_value_params is not None
+        ), "key_value_params cannot be None in ParameterConstraints of planner for embedding compute kernel: {}".format(
+            compute_kernel
+        )
+        assert (
+            key_value_params.max_l1_cache_size is not None
+        ), "key_value_params.max_l1_cache_size cannot be None in ParameterConstraints of planner for embedding compute kernel: {}".format(
+            compute_kernel
+        )
+        assert (
+            key_value_params.l2_cache_size is not None
+        ), "key_value_params.l2_cache_size cannot be None in ParameterConstraints of planner for embedding compute kernel: {}".format(
+            compute_kernel
+        )
+        # TODO(wangj): is this expected?
+        hbm_specific_sizes = [
+            none_throws(key_value_params.max_l1_cache_size) * 1024 * 1024
+            for _ in hbm_specific_sizes
+        ]
+        ddr_specific_sizes = [
+            none_throws(key_value_params.l2_cache_size) * 1024 * 1024 * 1024
+            for _ in ddr_specific_sizes
+        ]
+
     hbm_sizes: List[int] = [
         (
             hbm_specific_size

torchrec/distributed/planner/tests/test_planners.py

@@ -12,11 +12,14 @@
 
 import torch
 from torch import nn
+from torchrec import EmbeddingConfig
+from torchrec.distributed.embedding import EmbeddingCollectionSharder
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel
 from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
 from torchrec.distributed.planner import ParameterConstraints
 from torchrec.distributed.planner.planners import EmbeddingShardingPlanner
 from torchrec.distributed.planner.proposers import EmbeddingOffloadScaleupProposer
+from torchrec.distributed.planner.stats import EmbeddingStats
 from torchrec.distributed.planner.types import (
     PlannerError,
     PlannerErrorType,
@@ -31,6 +34,7 @@
     CacheParams,
     DataType,
     EmbeddingModuleShardingPlan,
+    KeyValueParams,
     ModuleSharder,
     ShardingPlan,
     ShardingType,
@@ -359,3 +363,149 @@ def test_auto_sharder_solution(self) -> None:
         self.assertSetEqual(
             {EmbeddingComputeKernel.FUSED_UVM_CACHING.value}, compute_kernels
         )
+
+    def test_planner_with_virtual_table(self) -> None:
+        table_count = 4
+        tables = [
+            EmbeddingConfig(
+                num_embeddings=1_125_899_902_955_520,
+                embedding_dim=64,
+                name="table_" + str(i),
+                feature_names=["feature_" + str(i)],
+                use_virtual_table=True,
+                total_num_buckets=3_991_680,
+            )
+            for i in range(table_count // 2)
+        ] + [
+            EmbeddingConfig(
+                num_embeddings=100_000,
+                embedding_dim=64,
+                name="table_" + str(i),
+                feature_names=["feature_" + str(i)],
+            )
+            for i in range(table_count // 2, table_count)
+        ]
+        print(tables)
+        model = TestSparseNN(tables=tables, sparse_device=torch.device("meta"))
+
+        constraints = {
+            **{
+                f"table_{i}": ParameterConstraints(
+                    sharding_types=["row_wise"],
+                    compute_kernels=["dram_virtual_table"],
+                )
+                for i in range(table_count // 2)
+            },
+            **{
+                f"table_{i}": ParameterConstraints(
+                    enforce_hbm=False,
+                )
+                for i in range(table_count // 2, table_count)
+            },
+        }
+        planner = EmbeddingShardingPlanner(
+            topology=self.topology,
+            proposer=EmbeddingOffloadScaleupProposer(),
+            constraints=constraints,
+        )
+
+        self.assertRaisesRegex(
+            AssertionError,
+            "key_value_params cannot be None in ParameterConstraints of planner for embedding compute kernel: dram_virtual_table",
+            planner.plan,
+            module=model,
+            sharders=[EmbeddingCollectionSharder()],
+        )
+
+        constraints = {
+            **{
+                f"table_{i}": ParameterConstraints(
+                    sharding_types=["row_wise"],
+                    compute_kernels=["dram_virtual_table"],
+                    key_value_params=KeyValueParams(
+                        l2_cache_size=64, max_l1_cache_size=128
+                    ),
+                )
+                for i in range(table_count // 2)
+            },
+            **{
+                f"table_{i}": ParameterConstraints(
+                    cache_params=CacheParams(algorithm=CacheAlgorithm.LRU),
+                )
+                for i in range(table_count // 2, table_count)
+            },
+        }
+
+        topology = Topology(
+            world_size=2,
+            hbm_cap=1024 * 1024 * 1024 * 2,
+            ddr_cap=1024 * 1024 * 1024 * 256,
+            compute_device="cuda",
+        )
+
+        planner = EmbeddingShardingPlanner(
+            topology=topology,
+            proposer=EmbeddingOffloadScaleupProposer(),
+            constraints=constraints,
+        )
+        sharding_plan = planner.plan(
+            module=model, sharders=[EmbeddingCollectionSharder()]  # pyre-ignore
+        )
+
+        expected_ranks = [[0, 1], [0, 1], [0, 1], [0, 1]]
+        ranks = [
+            cast(List[int], param_shard.ranks)
+            for param_shard in cast(
+                EmbeddingModuleShardingPlan, sharding_plan.plan["sparse.ec"]
+            ).values()
+        ]
+        compute_kernels = {
+            param_shard.compute_kernel
+            for param_shard in cast(
+                EmbeddingModuleShardingPlan, sharding_plan.plan["sparse.ec"]
+            ).values()
+        }
+        self.assertEqual(sorted(expected_ranks), sorted(ranks))
+        self.assertSetEqual(
+            {
+                EmbeddingComputeKernel.DRAM_VIRTUAL_TABLE.value,
+                EmbeddingComputeKernel.FUSED_UVM_CACHING.value,
+            },
+            compute_kernels,
+        )
+
+        for table_index in range(4):
+            shards = sharding_plan.plan["sparse.ec"][f"table_{table_index}"].sharding_spec.shards  # pyre-ignore
+            self.assertEqual(len(shards), 2)
+            self.assertEqual(shards[0].shard_offsets, [0, 0])
+            self.assertEqual(
+                shards[0].shard_sizes,
+                [562949951477760 if table_index < 2 else 50_000, 64],
+            )
+            self.assertEqual(
+                shards[1].shard_offsets,
+                [562949951477760 if table_index < 2 else 50_000, 0],
+            )
+            self.assertEqual(
+                shards[1].shard_sizes,
+                [562949951477760 if table_index < 2 else 50_000, 64],
+            )
+        stats: List[str] = cast(EmbeddingStats, planner._stats[0])._stats_table
+        self.assertTrue(
+            any(
+                "dram_virtual_table: HBM: 0.501 GB, DDR: 256.0 GB" in line
+                for line in stats
+            )
+        )
+        self.assertTrue(
+            any(
+                "fused_uvm_caching: HBM: 0.011 GB, DDR: 0.048 GB" in line
+                for line in stats
+            )
+        )
+        self.assertTrue(
+            any("Max HBM: 0.256 GB on ranks [0, 1]" in line for line in stats)
+        )
+        self.assertTrue(
+            any("Min HBM: 0.256 GB on ranks [0, 1]" in line for line in stats)
+        )

torchrec/distributed/test_utils/test_model.py

@@ -15,6 +15,7 @@
 import torch
 import torch.nn as nn
 from tensordict import TensorDict
+from torchrec import EmbeddingCollection
 from torchrec.distributed.embedding import EmbeddingCollectionSharder
 from torchrec.distributed.embedding_tower_sharding import (
     EmbeddingTowerCollectionSharder,
@@ -1207,6 +1208,47 @@ def _post_sparsenn_forward(
         )
 
 
+class TestECSparseArch(nn.Module):
+    """
+    Basic nn.Module for testing
+
+    Args:
+        tables
+        device
+
+    Call Args:
+        features
+
+    Returns:
+        KeyedTensor
+    """
+
+    def __init__(
+        self,
+        tables: List[EmbeddingConfig],
+        # weighted_tables: List[EmbeddingBagConfig],
+        device: Optional[torch.device] = None,
+        # max_feature_lengths: Optional[Dict[str, int]] = None,
+    ) -> None:
+        super().__init__()
+        if device is None:
+            device = torch.device("cpu")
+        self.ec: EmbeddingCollection = EmbeddingCollection(
+            tables=tables,
+            device=device,
+        )
+
+    def forward(
+        self,
+        features: KeyedJaggedTensor,
+        weighted_features: Optional[KeyedJaggedTensor] = None,
+        batch_size: Optional[int] = None,
+    ) -> KeyedTensor:
+        ec = self.ec(features)
+        result = _post_sparsenn_forward(ec, None, None, batch_size)
+        return result
+
+
 class TestSparseArch(nn.Module):
     """
     Basic nn.Module for testing
@@ -1349,7 +1391,7 @@ class TestSparseNN(TestSparseNNBase, CopyableMixin):
 
     def __init__(
         self,
-        tables: List[EmbeddingBagConfig],
+        tables: Union[List[EmbeddingBagConfig], List[EmbeddingConfig]],
         num_float_features: int = 10,
         weighted_tables: Optional[List[EmbeddingBagConfig]] = None,
         embedding_groups: Optional[Dict[str, List[str]]] = None,
@@ -1373,14 +1415,19 @@ def __init__(
         self.dense = TestDenseArch(num_float_features, dense_device)
         if zch:
             self.sparse: nn.Module = TestSparseArchZCH(
-                tables,
+                tables,  # pyre-ignore
                 weighted_tables,
                 torch.device("meta"),
                 return_remapped=True,
             )
+        elif isinstance(tables[0], EmbeddingConfig):
+            self.sparse = TestECSparseArch(
+                tables,  # pyre-ignore [6]
+                sparse_device,
+            )
         else:
             self.sparse = TestSparseArch(
-                tables,
+                tables,  # pyre-ignore
                 weighted_tables,
                 sparse_device,
                 max_feature_lengths,