Preserve unquantized embedding tables through sharding

Summary:
PNI builds TorchRec per-table weight dtype maps from `select_quantization_dtype()`. The selector uses `None` to mean that a table must remain unquantized, but omitting that table from TorchRec's map means to use the default INT8 dtype. That mismatch produced an INT8 TBE specification for a table whose packed weight remained floating point, so ICE attempted to load a Float weight into a Byte table.

Record selector `None` explicitly as `torch.float` in the existing ADS EBC, IG, and retrieval PNI quantization callsites. Explicit INT8, INT4, and INT2 selections are unchanged, and float-to-half policies continue to return and preserve `torch.half`.

Mixed floating and quantized tables also expose a TorchRec state initialization gap when split scale/bias state is enabled. Floating TBEs correctly return neither qscale nor qbias. Preserve their weight state while omitting quantization-parameter state, and reject an asymmetric qscale/qbias pair as invalid instead of producing an incomplete state dict.

Together these changes preserve each table's selected representation across PNI authoring, TorchRec sharding, model packing, and ICE loading without special-casing a model or transform.

Differential Revision: D113501049

Author: shawnxu26

torchrec/distributed/quant_state.py

@@ -176,6 +176,16 @@ def _initialize_torch_state(  # noqa: C901
                 # end of weight shards section
 
                 # weight_qscale & weight_qbias section:
+                if (tbe_split_qscale is None) != (tbe_split_qbias is None):
+                    raise RuntimeError(
+                        "Quantized table must provide both qscale and qbias: "
+                        f"table={table.name}, data_type={table.data_type}, "
+                        f"qscale_is_none={tbe_split_qscale is None}, "
+                        f"qbias_is_none={tbe_split_qbias is None}"
+                    )
+                if tbe_split_qscale is None:
+                    continue
+
                 # For RW - ShardedTensorBase
                 # For CW - List[Tensor] that logically corresponds to the same unsharded Tensor, but present on each sharded rank
                 for (

torchrec/distributed/tests/test_infer_shardings.py

@@ -147,6 +147,42 @@ def setUp(self) -> None:
         super().setUp()
         set_propogate_device(True)
 
+    def test_tw_mixed_float_and_int8_split_scale_bias_cpu(self) -> None:
+        device = torch.device("cpu")
+        mi = create_test_model(
+            num_embeddings=16,
+            emb_dim=8,
+            world_size=1,
+            batch_size=2,
+            dense_device=device,
+            sparse_device=device,
+            num_features=2,
+            num_weighted_features=0,
+        )
+        mi.quant_model = quantize(
+            module=mi.model,
+            inplace=False,
+            quant_state_dict_split_scale_bias=True,
+            per_table_weight_dtypes={
+                "table_0": torch.float,
+                "table_1": torch.quint8,
+            },
+        )
+
+        sharded_model = shard_qebc(
+            mi,
+            sharding_type=ShardingType.TABLE_WISE,
+            device=device,
+        )
+        state_keys = sharded_model.state_dict().keys()
+
+        self.assertTrue(any(key.endswith("table_0.weight") for key in state_keys))
+        self.assertFalse(any("table_0.weight_q" in key for key in state_keys))
+        self.assertTrue(
+            any(key.endswith("table_1.weight_qscale") for key in state_keys)
+        )
+        self.assertTrue(any(key.endswith("table_1.weight_qbias") for key in state_keys))
+
     @unittest.skipIf(
         torch.cuda.device_count() <= 1,
         "Not enough GPUs available",