WIP: Feature Quantization Scaffolding

gigl/common/data/dataloaders.py

@@ -21,6 +21,7 @@
 class LoadedEntityTensors(NamedTuple):
     ids: torch.Tensor
     features: Optional[torch.Tensor]
+    quantized_features: Optional[torch.Tensor]
     labels: Optional[torch.Tensor]
 
 
@@ -44,6 +45,10 @@ class SerializedTFRecordInfo:
     feature_dim: int
     # Entity ID Key for current entity. If this is a Node Entity, this must be a string. If this is an edge entity, this must be a Tuple[str, str] for the source and destination ids.
     entity_key: Union[str, Tuple[str, str]]
+    # Packed uint8 feature name to load for the current node entity.
+    packed_feature_key: Optional[str] = None
+    # Number of packed uint8 columns for the current node entity.
+    packed_feature_dim: int = 0
     # Name of the label columns for the current entity, defaults to an empty list.
     label_keys: Sequence[str] = field(default_factory=list)
     # The regex pattern to match the TFRecord files at the specified prefix
@@ -367,10 +372,11 @@ def load_as_torch_tensors(
             serialized_tf_record_info (SerializedTFRecordInfo): Information for how TFRecord files are serialized on disk.
             tf_dataset_options (TFDatasetOptions): The options to use when building the dataset.
         Returns:
-            LoadedEntityTensors: The (id_tensor, feature_tensor, label_tensor) for the loaded entities.
+            LoadedEntityTensors: The (id_tensor, feature_tensor, quantized_feature_tensor, label_tensor) for the loaded entities.
         """
         entity_key = serialized_tf_record_info.entity_key
         feature_keys = serialized_tf_record_info.feature_keys
+        packed_feature_key = serialized_tf_record_info.packed_feature_key
         label_keys = serialized_tf_record_info.label_keys
 
         # We make a deep copy of the feature spec dict so that future modifications don't redirect to the input
@@ -392,6 +398,16 @@ def load_as_torch_tensors(
                 feature_spec_dict[entity_key] = tf.io.FixedLenFeature(
                     shape=[], dtype=tf.int64
                 )
+            if (
+                packed_feature_key is not None
+                and packed_feature_key not in feature_spec_dict
+            ):
+                logger.info(
+                    f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`"
+                )
+                feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature(
+                    shape=[], dtype=tf.string
+                )
         else:
             id_concat_axis = 1
             proccess_id_tensor = lambda t: tf.stack(
@@ -435,13 +451,24 @@ def load_as_torch_tensors(
             else:
                 empty_feature = None
 
+            if packed_feature_key is not None:
+                empty_quantized_feature = torch.empty(
+                    (0, serialized_tf_record_info.packed_feature_dim),
+                    dtype=torch.uint8,
+                )
+            else:
+                empty_quantized_feature = None
+
             if label_keys:
                 empty_label = torch.empty(0, len(label_keys))
             else:
                 empty_label = None
 
             return LoadedEntityTensors(
-                ids=empty_entity, features=empty_feature, labels=empty_label
+                ids=empty_entity,
+                features=empty_feature,
+                quantized_features=empty_quantized_feature,
+                labels=empty_label,
             )
 
         dataset = TFRecordDataLoader._build_dataset_for_uris(
@@ -454,6 +481,7 @@ def load_as_torch_tensors(
         num_entities_processed = 0
         id_tensors: list[torch.Tensor] = []
         feature_tensors: list[torch.Tensor] = []
+        quantized_feature_tensors: list[torch.Tensor] = []
         label_tensors: list[torch.Tensor] = []
         for idx, batch in enumerate(dataset):
             id_tensors.append(proccess_id_tensor(batch))
@@ -465,6 +493,15 @@ def load_as_torch_tensors(
                     feature_tensors.append(feature_tensor)
                 if label_tensor is not None:
                     label_tensors.append(label_tensor)
+            if packed_feature_key is not None:
+                quantized_feature_tensor = tf.io.decode_raw(
+                    batch[packed_feature_key], tf.uint8
+                )
+                quantized_feature_tensor = tf.reshape(
+                    quantized_feature_tensor,
+                    [-1, serialized_tf_record_info.packed_feature_dim],
+                )
+                quantized_feature_tensors.append(quantized_feature_tensor)
             num_entities_processed += (
                 id_tensors[-1].shape[0]
                 if entity_type == FeatureTypes.NODE
@@ -483,11 +520,16 @@ def load_as_torch_tensors(
             tf.concat(id_tensors, axis=id_concat_axis)
         )
         output_feature_tensor: Optional[torch.Tensor] = None
+        output_quantized_feature_tensor: Optional[torch.Tensor] = None
         output_label_tensor: Optional[torch.Tensor] = None
         if feature_tensors:
             output_feature_tensor = _tf_tensor_to_torch_tensor(
                 tf.concat(feature_tensors, axis=0)
             )
+        if quantized_feature_tensors:
+            output_quantized_feature_tensor = _tf_tensor_to_torch_tensor(
+                tf.concat(quantized_feature_tensors, axis=0)
+            ).to(torch.uint8)
         if label_tensors:
             output_label_tensor = _tf_tensor_to_torch_tensor(
                 tf.concat(label_tensors, axis=0)
@@ -503,5 +545,8 @@ def load_as_torch_tensors(
             f"Converted {num_entities_processed:,} {entity_type.name} to torch tensors in {end - start:.2f} seconds"
         )
         return LoadedEntityTensors(
-            ids=id_tensor, features=output_feature_tensor, labels=output_label_tensor
+            ids=id_tensor,
+            features=output_feature_tensor,
+            quantized_features=output_quantized_feature_tensor,
+            labels=output_label_tensor,
         )

gigl/common/data/load_torch_tensors.py

@@ -18,6 +18,7 @@
 from gigl.types.graph import (
     DEFAULT_HOMOGENEOUS_EDGE_TYPE,
     DEFAULT_HOMOGENEOUS_NODE_TYPE,
+    FeatureQuantizationMetadata,
     LoadedGraphTensors,
 )
 from gigl.utils.share_memory import share_memory
@@ -26,6 +27,7 @@
 
 _ID_FMT = "{entity}_ids"
 _FEATURE_FMT = "{entity}_features"
+_PACKED_FEATURE_FMT = "{entity}_packed_features"
 _LABEL_FMT = "{entity}_labels"
 _EDGE_WEIGHTS_KEY = "edge_weights"
 _NODE_KEY = "node"
@@ -113,6 +115,10 @@ class SerializedGraphMetadata:
     negative_label_entity_info: Optional[
         Union[SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]]
     ] = None
+    # Optional node quantization metadata.
+    node_quantization_metadata: Optional[
+        Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]]
+    ] = None
 
 
 def _data_loading_process(
@@ -178,6 +184,7 @@ def _data_loading_process(
 
         ids: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
         features: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
+        quantized_features: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
         labels: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
         weights: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
         for (
@@ -192,12 +199,20 @@ def _data_loading_process(
                 raise NotImplementedError(
                     "Label keys are not supported for edge entities"
                 )
+            if (
+                serialized_entity_tf_record_info.packed_feature_key is not None
+                and not serialized_entity_tf_record_info.is_node_entity
+            ):
+                raise NotImplementedError(
+                    "Packed feature keys are not supported for edge entities"
+                )
             loaded_entity = tf_record_dataloader.load_as_torch_tensors(
                 serialized_tf_record_info=serialized_entity_tf_record_info,
                 tf_dataset_options=tf_dataset_options,
             )
             entity_ids = loaded_entity.ids
             entity_features = loaded_entity.features
+            entity_quantized_features = loaded_entity.quantized_features
             entity_labels = loaded_entity.labels
             ids[graph_type] = entity_ids
             logger.info(
@@ -213,6 +228,16 @@ def _data_loading_process(
                     f"Rank {rank} did not detect {entity_type} features for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                 )
 
+            if entity_quantized_features is not None:
+                quantized_features[graph_type] = entity_quantized_features
+                logger.info(
+                    f"Rank {rank} finished loading {entity_type} quantized features of shape {entity_quantized_features.shape} for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
+                )
+            else:
+                logger.info(
+                    f"Rank {rank} did not detect {entity_type} quantized features for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
+                )
+
             if entity_labels is not None:
                 labels[graph_type] = entity_labels
                 logger.info(
@@ -291,6 +316,12 @@ def _data_loading_process(
             share_memory(features)
             # We convert the features back to homogeneous from the default heterogeneous setup if our provided input was homogeneous
 
+        if quantized_features:
+            logger.info(
+                f"Rank {rank} is attempting to share {entity_type} quantized feature memory for tfrecord directories: {all_tf_record_uris}"
+            )
+            share_memory(quantized_features)
+
         if labels:
             logger.info(
                 f"Rank {rank} is attempting to share {entity_type} label memory for tfrecord directories: {all_tf_record_uris}"
@@ -310,6 +341,12 @@ def _data_loading_process(
             output_dict[_FEATURE_FMT.format(entity=entity_type)] = (
                 list(features.values())[0] if is_input_homogeneous else features
             )
+        if quantized_features:
+            output_dict[_PACKED_FEATURE_FMT.format(entity=entity_type)] = (
+                list(quantized_features.values())[0]
+                if is_input_homogeneous
+                else quantized_features
+            )
         if labels:
             output_dict[_LABEL_FMT.format(entity=entity_type)] = (
                 list(labels.values())[0] if is_input_homogeneous else labels
@@ -480,6 +517,9 @@ def load_torch_tensors_from_tf_record(
 
     node_ids = node_output_dict[_ID_FMT.format(entity=_NODE_KEY)]
     node_features = node_output_dict.get(_FEATURE_FMT.format(entity=_NODE_KEY), None)
+    node_quantized_features = node_output_dict.get(
+        _PACKED_FEATURE_FMT.format(entity=_NODE_KEY), None
+    )
     node_labels = node_output_dict.get(_LABEL_FMT.format(entity=_NODE_KEY), None)
 
     edge_index = edge_output_dict[_ID_FMT.format(entity=_EDGE_KEY)]
@@ -507,6 +547,7 @@ def load_torch_tensors_from_tf_record(
     return LoadedGraphTensors(
         node_ids=node_ids,
         node_features=node_features,
+        node_quantized_features=node_quantized_features,
         node_labels=node_labels,
         edge_index=edge_index,
         edge_features=edge_features,

gigl/common/utils/feature_quantization/numpy_ops.py

@@ -0,0 +1,38 @@
+from collections.abc import Mapping
+
+import numpy as np
+
+
+def quantize_ndarray(
+    features: np.ndarray, *, bits: int, stats: Mapping[str, float]
+) -> np.ndarray:
+    """Quantize a 2D float array into packed uint8 codes."""
+    if bits not in (1, 2, 4, 8):
+        raise ValueError(f"bits must be one of 1, 2, 4, or 8, got {bits}.")
+    if features.ndim != 2:
+        raise ValueError(f"Expected a 2D feature array, got shape {features.shape}.")
+    if bits == 1:
+        # 1-bit quantization keeps only sign; values restore from neg/pos means.
+        codes = (features > 0).astype(np.uint8)
+    else:
+        # Linearly map clipped values into integer buckets.
+        levels = (1 << bits) - 1
+        lo, hi = stats["clip_min"], stats["clip_max"]
+        clipped = np.clip(features, lo, hi)
+        scaled = (clipped - lo) / (hi - lo)
+        codes = np.rint(scaled * levels).astype(np.uint8)
+    return pack_codes(codes, bits)
+
+
+def pack_codes(codes: np.ndarray, bits: int) -> np.ndarray:
+    """Pack low-bit feature codes high-bits-first along the final dimension."""
+    per_byte = 8 // bits
+    pad = (-codes.shape[-1]) % per_byte
+    if pad:
+        # Pad only the feature dimension of this 2D [row, feature] array.
+        codes = np.pad(codes, ((0, 0), (0, pad)), constant_values=0)
+    # Group the padded feature dimension into chunks that each form one byte.
+    codes = codes.reshape(codes.shape[0], -1, per_byte).astype(np.uint16)
+    shifts = bits * np.arange(per_byte - 1, -1, -1, dtype=np.uint16)
+    weights = (1 << shifts).astype(np.uint16)
+    return np.sum(codes * weights, axis=-1).astype(np.uint8)

gigl/common/utils/feature_quantization/torch_ops.py

@@ -0,0 +1,42 @@
+import torch
+
+from gigl.types.graph import FeatureQuantizationMetadata
+
+
+def dequantize_torch_tensor(
+    packed_features: torch.Tensor,
+    metadata: FeatureQuantizationMetadata,
+) -> torch.Tensor:
+    """Reconstruct approximate float features from packed uint8 codes."""
+    q = metadata
+
+    if packed_features.size(-1) != q.packed_feature_dim:
+        raise ValueError(
+            f"Expected packed feature dim {q.packed_feature_dim} for "
+            f"{q.dequantized_feature_dim} {q.bits}-bit features, got "
+            f"{packed_features.size(-1)}."
+        )
+
+    codes = _unpack_torch_tensor(
+        packed_features, dim=q.dequantized_feature_dim, bits=q.bits
+    ).float()
+    if q.bits == 1:
+        if q.neg_mean is None or q.pos_mean is None:
+            raise ValueError("1-bit dequantization requires pos_mean/neg_mean")
+        return torch.where(codes.bool(), q.pos_mean, q.neg_mean)
+    else:
+        if q.clip_min is None or q.clip_max is None:
+            raise ValueError(f"{q.bits}-bit dequantization requires clip_min/clip_max")
+        levels = (1 << q.bits) - 1
+        return q.clip_min + (codes / levels) * (q.clip_max - q.clip_min)
+
+
+def _unpack_torch_tensor(
+    packed_features: torch.Tensor, *, dim: int, bits: int
+) -> torch.Tensor:
+    per_byte = 8 // bits
+    mask = (1 << bits) - 1
+    # Extract high-bits-first codes from each packed byte.
+    shifts = bits * torch.arange(per_byte - 1, -1, -1, device=packed_features.device)
+    codes = (packed_features.unsqueeze(-1).to(torch.int16) >> shifts).bitwise_and(mask)
+    return codes.reshape(*packed_features.shape[:-1], -1)[..., :dim].to(torch.uint8)

gigl/distributed/base_dist_loader.py

@@ -242,6 +242,7 @@ def __init__(
         )
         self._node_feature_info = dataset_schema.node_feature_info
         self._edge_feature_info = dataset_schema.edge_feature_info
+        self._node_quantization_metadata = dataset_schema.node_quantization_metadata
 
         self._sampler_options = sampler_options
         self._non_blocking_transfers = non_blocking_transfers