fix default uniform init range (#3336)

Summary:
Pull Request resolved: #3336

since zch v.Next is using a very large virtual table size, 2^50, the default uniform init value becomes very small, and when the weight dtype is half, those value essentially becomes 0.
We have observed the weight init value is all 0 from the debug log: https://fburl.com/mlhub/aea9mbzf
 {F1981621246}

Reviewed By: kathyxuyy

Differential Revision: D81296621

fbshipit-source-id: 46a4eb87d6df7a8073efdacffad355b66d683364

Author: emlin

torchrec/distributed/batched_embedding_kernel.py

@@ -14,6 +14,7 @@
 import logging
 import tempfile
 from dataclasses import dataclass
+from math import sqrt
 from typing import (
     Any,
     cast,
@@ -192,6 +193,9 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]:
         )
 
     # populate init min and max
+    if config.is_using_virtual_table:
+        _generate_init_range_for_virtual_tables(ssd_tbe_params, config)
+
     if (
         "ssd_uniform_init_lower" not in ssd_tbe_params
         or "ssd_uniform_init_upper" not in ssd_tbe_params
@@ -245,6 +249,50 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]:
     return ssd_tbe_params
 
 
+def _generate_init_range_for_virtual_tables(
+    tbe_params: Dict[str, Any],
+    config: GroupedEmbeddingConfig,
+) -> None:
+    """
+    Generate uniform init range for zero collision TBE based
+    """
+    # populate init min and max
+    if (
+        "ssd_uniform_init_lower" not in tbe_params
+        or "ssd_uniform_init_upper" not in tbe_params
+    ):
+        # Right now we do not support a per table init max and min. To use
+        # per table init max and min, either we allow it in SSD TBE, or we
+        # create one SSD TBE per table.
+        weights_precision = data_type_to_sparse_type(config.data_type)
+
+        # For Float32: use mathematically correct values, for Half: use safe range
+        max_size = 4_000_000_000  # 4B virtual embeddings
+        default_init_range = (
+            (-sqrt(1 / max_size), sqrt(1 / max_size))
+            if weights_precision.as_dtype() == torch.float32
+            else (-0.001, 0.001)
+        )
+
+        def get_init_value(
+            table_init_val: Optional[float], default_value: float
+        ) -> float:
+            return table_init_val if table_init_val is not None else default_value
+
+        init_mins = [
+            get_init_value(table.weight_init_min, default_init_range[0])
+            for table in config.embedding_tables
+        ]
+        init_maxs = [
+            get_init_value(table.weight_init_max, default_init_range[1])
+            for table in config.embedding_tables
+        ]
+
+        num_tables = len(config.embedding_tables)
+        tbe_params["ssd_uniform_init_lower"] = sum(init_mins) / num_tables
+        tbe_params["ssd_uniform_init_upper"] = sum(init_maxs) / num_tables
+
+
 def _populate_zero_collision_tbe_params(
     tbe_params: Dict[str, Any],
     sharded_local_buckets: List[Tuple[int, int, int]],
@@ -1872,6 +1920,9 @@ def __init__(
         assert (
             len({table.embedding_dim for table in config.embedding_tables}) == 1
         ), "Currently we expect all tables in SSD TBE to have the same embedding dimension."
+        assert (
+            config.is_using_virtual_table
+        ), "Try to create ZeroCollisionKeyValueEmbedding for non virtual tables"
         for table in config.embedding_tables:
             assert table.local_cols % 4 == 0, (
                 f"table {table.name} has local_cols={table.local_cols} "
@@ -2751,6 +2802,9 @@ def __init__(
         assert (
             len({table.embedding_dim for table in config.embedding_tables}) == 1
         ), "Currently we expect all tables in SSD TBE to have the same embedding dimension."
+        assert (
+            config.is_using_virtual_table
+        ), "Try to create ZeroCollisionKeyValueEmbeddingBag for non virtual tables"
 
         for table in config.embedding_tables:
             assert table.local_cols % 4 == 0, (