From cde04462cca18292da975bb03bd9050c450c6fb3 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 29 Jun 2026 21:29:19 +0000 Subject: [PATCH 01/44] Initial draft --- gigl/common/data/dataloaders.py | 68 ++++++- gigl/common/data/load_torch_tensors.py | 34 ++++ gigl/distributed/base_dist_loader.py | 2 + gigl/distributed/base_sampler.py | 52 +++++ gigl/distributed/dataset_factory.py | 6 + gigl/distributed/dist_ablp_neighborloader.py | 10 + gigl/distributed/dist_dataset.py | 154 +++++++++++++++ gigl/distributed/dist_partitioner.py | 180 +++++++++++++++--- gigl/distributed/dist_range_partitioner.py | 29 ++- .../distributed/distributed_neighborloader.py | 10 + gigl/distributed/utils/neighborloader.py | 165 +++++++++++++++- .../serialized_graph_metadata_translator.py | 72 ++++++- .../pb_wrappers/preprocessed_metadata.py | 72 ++++++- gigl/types/graph.py | 24 +++ .../research/gbml/preprocessed_metadata.proto | 36 ++++ .../gbml/preprocessed_metadata_pb2.py | 57 ++++-- .../gbml/preprocessed_metadata_pb2.pyi | 86 ++++++++- 17 files changed, 1011 insertions(+), 46 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 0ac83f0e7..df9d1b436 100644 --- a/gigl/common/data/dataloaders.py +++ b/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 names to load for the current node entity. + quantized_feature_keys: Sequence[str] = field(default_factory=list) + # Number of packed uint8 columns for the current node entity. + quantized_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 @@ -187,6 +192,25 @@ def _concatenate_features_by_names( return _get_labels_from_features(combined_feature_tensor, label_dim) +def _concatenate_quantized_features_by_names( + feature_key_to_tf_tensor: dict[str, tf.Tensor], + feature_keys: Sequence[str], +) -> Optional[tf.Tensor]: + """Concatenate packed quantized feature tensors as uint8 columns.""" + if not feature_keys: + return None + + features: list[tf.Tensor] = [] + for feature_key in feature_keys: + tensor = feature_key_to_tf_tensor[feature_key] + if tensor.dtype != tf.uint8: + tensor = tf.cast(tensor, tf.uint8) + if len(tensor.shape) == 1: + tensor = tf.expand_dims(tensor, axis=-1) + features.append(tensor) + return tf.concat(features, axis=1) + + def _tf_tensor_to_torch_tensor(tf_tensor: tf.Tensor) -> torch.Tensor: """ Converts a TensorFlow tensor to a PyTorch tensor using DLPack to ensure zero-copy conversion. @@ -371,6 +395,7 @@ def load_as_torch_tensors( """ entity_key = serialized_tf_record_info.entity_key feature_keys = serialized_tf_record_info.feature_keys + quantized_feature_keys = serialized_tf_record_info.quantized_feature_keys 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 +417,19 @@ def load_as_torch_tensors( feature_spec_dict[entity_key] = tf.io.FixedLenFeature( shape=[], dtype=tf.int64 ) + for feature_key in quantized_feature_keys: + if feature_key not in feature_spec_dict: + feature_shape = ( + [serialized_tf_record_info.quantized_feature_dim] + if len(quantized_feature_keys) == 1 + else [] + ) + logger.info( + f"Injecting quantized feature key {feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape={feature_shape}, dtype=tf.int64)`" + ) + feature_spec_dict[feature_key] = tf.io.FixedLenFeature( + shape=feature_shape, dtype=tf.int64 + ) else: id_concat_axis = 1 proccess_id_tensor = lambda t: tf.stack( @@ -435,13 +473,24 @@ def load_as_torch_tensors( else: empty_feature = None + if quantized_feature_keys: + empty_quantized_feature = torch.empty( + (0, serialized_tf_record_info.quantized_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 +503,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 +515,12 @@ def load_as_torch_tensors( feature_tensors.append(feature_tensor) if label_tensor is not None: label_tensors.append(label_tensor) + if quantized_feature_keys: + quantized_feature_tensor = _concatenate_quantized_features_by_names( + batch, quantized_feature_keys + ) + if quantized_feature_tensor is not None: + quantized_feature_tensors.append(quantized_feature_tensor) num_entities_processed += ( id_tensors[-1].shape[0] if entity_type == FeatureTypes.NODE @@ -483,11 +539,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 +564,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, ) diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index 2ba43ce10..f852da49b 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -19,6 +19,7 @@ DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, LoadedGraphTensors, + NodeQuantizationMetadata, ) from gigl.utils.share_memory import share_memory @@ -29,6 +30,7 @@ _LABEL_FMT = "{entity}_labels" _EDGE_WEIGHTS_KEY = "edge_weights" _NODE_KEY = "node" +_NODE_QUANTIZED_KEY = "node_quantized" def _extract_weight_col( @@ -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[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ] = 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 ( @@ -198,6 +205,7 @@ def _data_loading_process( ) 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 +221,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 +309,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 +334,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 and entity_type == _NODE_KEY: + output_dict[_FEATURE_FMT.format(entity=_NODE_QUANTIZED_KEY)] = ( + 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 +510,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( + _FEATURE_FMT.format(entity=_NODE_QUANTIZED_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 +540,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, diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index 02d7c69e1..b9d2b059b 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -241,6 +241,8 @@ def __init__( dataset_schema.is_homogeneous_with_labeled_edge_type ) self._node_feature_info = dataset_schema.node_feature_info + self._node_quantized_feature_info = dataset_schema.node_quantized_feature_info + self._node_quantization_metadata = dataset_schema.node_quantization_metadata self._edge_feature_info = dataset_schema.edge_feature_info self._sampler_options = sampler_options diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 986ba5d58..1e3a97d77 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -20,6 +20,7 @@ POSITIVE_LABEL_METADATA_KEY, ABLPNodeSamplerInput, ) +from gigl.distributed.utils.neighborloader import NODE_QUANTIZED_FEATURES_METADATA_KEY from gigl.utils.data_splitters import PADDING_NODE @@ -91,6 +92,25 @@ class BaseDistNeighborSampler(GLTDistNeighborSampler): sampling strategy (e.g., k-hop neighbor sampling, PPR-based sampling). """ + def __init__(self, *args, **kwargs): + data = kwargs.get("data") + super().__init__(*args, **kwargs) + self.dist_node_quantized_feature = None + if ( + self.collect_features + and data is not None + and getattr(data, "node_quantized_features", None) is not None + ): + self.dist_node_quantized_feature = DistFeature( + data.num_partitions, + data.partition_idx, + data.node_quantized_features, + data.node_quantized_feature_pb, + local_only=False, + rpc_router=self.rpc_router, + device=self.device, + ) + def _prepare_sample_loop_inputs( self, inputs: NodeSamplerInput, @@ -322,6 +342,32 @@ async def _collate_fn( for ntype, fut in nfeat_fut_dict.items(): nfeats = await wrap_torch_future(fut) result_map[f"{as_str(ntype)}.nfeats"] = nfeats + if self.dist_node_quantized_feature is not None: + if self.use_all2all: + sorted_ntype = sorted( + self.dist_node_quantized_feature.feature_pb.keys() + ) + quantized_nfeat_dict = ( + self.dist_node_quantized_feature.get_all2all( + output, sorted_ntype + ) + ) + for ntype, quantized_nfeats in quantized_nfeat_dict.items(): + result_map[ + f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{as_str(ntype)}" + ] = quantized_nfeats + else: + quantized_nfeat_fut_dict = {} + for ntype, nodes in output.node.items(): + nodes = nodes.to(torch.long) + quantized_nfeat_fut_dict[ntype] = ( + self.dist_node_quantized_feature.async_get(nodes, ntype) + ) + for ntype, fut in quantized_nfeat_fut_dict.items(): + quantized_nfeats = await wrap_torch_future(fut) + result_map[ + f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{as_str(ntype)}" + ] = quantized_nfeats if self.dist_edge_feature is not None and self.with_edge: efeat_fut_dict = {} for etype in self.edge_types: @@ -382,6 +428,12 @@ async def _collate_fn( fut = self.dist_node_feature.async_get(output.node) nfeats = await wrap_torch_future(fut) result_map["nfeats"] = nfeats + if self.dist_node_quantized_feature is not None: + fut = self.dist_node_quantized_feature.async_get(output.node) + quantized_nfeats = await wrap_torch_future(fut) + result_map[ + f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}" + ] = quantized_nfeats if self.dist_edge_feature is not None: eids = result_map["eids"] fut = self.dist_edge_feature.async_get(eids) diff --git a/gigl/distributed/dataset_factory.py b/gigl/distributed/dataset_factory.py index cbd7006f8..dc9e154ff 100644 --- a/gigl/distributed/dataset_factory.py +++ b/gigl/distributed/dataset_factory.py @@ -180,6 +180,10 @@ def _load_and_build_partitioned_dataset( partitioner.register_node_features( node_features=loaded_graph_tensors.node_features ) + if loaded_graph_tensors.node_quantized_features is not None: + partitioner.register_node_quantized_features( + node_quantized_features=loaded_graph_tensors.node_quantized_features + ) if loaded_graph_tensors.node_labels is not None: partitioner.register_node_labels(node_labels=loaded_graph_tensors.node_labels) if loaded_graph_tensors.edge_weights is not None: @@ -205,6 +209,7 @@ def _load_and_build_partitioned_dataset( del ( loaded_graph_tensors.node_ids, loaded_graph_tensors.node_features, + loaded_graph_tensors.node_quantized_features, loaded_graph_tensors.edge_index, loaded_graph_tensors.edge_features, loaded_graph_tensors.edge_weights, @@ -222,6 +227,7 @@ def _load_and_build_partitioned_dataset( dataset.build( partition_output=partition_output, splitter=splitter, + node_quantization_metadata=serialized_graph_metadata.node_quantization_metadata, ) return dataset diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 50f42f5a9..757eefcd7 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -30,6 +30,7 @@ from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, + append_dequantized_node_features, extract_edge_type_metadata, extract_metadata, labeled_to_homogeneous, @@ -576,6 +577,8 @@ def _setup_for_colocated( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=dataset.node_feature_info, + node_quantized_feature_info=dataset.node_quantized_feature_info, + node_quantization_metadata=dataset.node_quantization_metadata, edge_feature_info=dataset.edge_feature_info, edge_dir=dataset.edge_dir, ), @@ -740,6 +743,8 @@ def _setup_for_graph_store( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=node_feature_info, + node_quantized_feature_info=None, + node_quantization_metadata=None, edge_feature_info=edge_feature_info, edge_dir=dataset.fetch_edge_dir(), ), @@ -876,6 +881,11 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: data = self._set_labels(data, positive_labels, negative_labels) data, metadata = self._apply_ppr_outputs(data, metadata) + data, metadata = append_dequantized_node_features( + data=data, + metadata=metadata, + node_quantization_metadata=self._node_quantization_metadata, + ) # Attach any remaining metadata (e.g. custom user-defined keys) directly onto the # data object so downstream code can access them via attribute lookup. diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index 181c2c7d9..a2b958271 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -24,6 +24,7 @@ FeatureInfo, FeaturePartitionData, GraphPartitionData, + NodeQuantizationMetadata, PartitionOutput, ) from gigl.utils.data_splitters import ( @@ -54,6 +55,9 @@ def __init__( node_feature_partition: Optional[ Union[Feature, dict[NodeType, Feature]] ] = None, + node_quantized_feature_partition: Optional[ + Union[Feature, dict[NodeType, Feature]] + ] = None, edge_feature_partition: Optional[ Union[Feature, dict[EdgeType, Feature]] ] = None, @@ -77,6 +81,12 @@ def __init__( node_feature_info: Optional[ Union[FeatureInfo, dict[NodeType, FeatureInfo]] ] = None, + node_quantized_feature_info: Optional[ + Union[FeatureInfo, dict[NodeType, FeatureInfo]] + ] = None, + node_quantization_metadata: Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ] = None, edge_feature_info: Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] ] = None, @@ -97,6 +107,7 @@ def __init__( The below arguments are only expected to be provided when re-serializing an instance of the DistDataset class after build() has been called graph_partition (Optional[Union[Graph, dict[EdgeType, Graph]]]): Partitioned Graph Data node_feature_partition (Optional[Union[Feature, dict[NodeType, Feature]]]): Partitioned Node Feature Data + node_quantized_feature_partition (Optional[Union[Feature, dict[NodeType, Feature]]]): Partitioned packed uint8 node feature data edge_feature_partition (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Partitioned Edge Feature Data node_labels (Optional[Union[Feature, dict[NodeType, Feature]]]): The labels of each node on the current machine node_partition_book (Optional[Union[PartitionBook, dict[NodeType, PartitionBook]]]): Node Partition Book @@ -109,6 +120,8 @@ def __init__( num_test: (Optional[Union[int, dict[NodeType, int]]]): Number of test nodes on the current machine. Will be a dict if heterogeneous. node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension of node features and its data type, will be a dict if heterogeneous. Note this will be None in the homogeneous case if the data has no node features, or will only contain node types with node features in the heterogeneous case. + node_quantized_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension and dtype for packed uint8 node features. + node_quantization_metadata: Optional[Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]]]: Metadata required for dequantizing packed node features. edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Dimension of edge features and its data type, will be a dict if heterogeneous. Note this will be None in the homogeneous case if the data has no edge features, or will only contain edge types with edge features in the heterogeneous case. degree_tensor: Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Pre-computed degree tensor. Lazily computed on first access via the degree_tensor property. @@ -142,6 +155,7 @@ def __init__( self._node_ids: Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]] = ( node_ids ) + self._node_quantized_features = node_quantized_feature_partition self._num_train = num_train self._num_val = num_val @@ -149,6 +163,8 @@ def __init__( # These fields are added so we can extract the node and edge feature dimensions and data type in the dataloader without having to lazily initialize the features. self._node_feature_info = node_feature_info + self._node_quantized_feature_info = node_quantized_feature_info + self._node_quantization_metadata = node_quantization_metadata self._edge_feature_info = edge_feature_info self._degree_tensor: Optional[ @@ -209,6 +225,18 @@ def node_features( ): self._node_features = new_node_features + @property + def node_quantized_features( + self, + ) -> Optional[Union[Feature, dict[NodeType, Feature]]]: + return self._node_quantized_features + + @property + def node_quantized_feature_pb( + self, + ) -> Optional[Union[PartitionBook, dict[NodeType, PartitionBook]]]: + return self.node_pb if self._node_quantized_features is not None else None + @property def edge_features(self) -> Optional[Union[Feature, dict[EdgeType, Feature]]]: """ @@ -303,6 +331,23 @@ def node_feature_info( """ return self._node_feature_info + @property + def node_quantized_feature_info( + self, + ) -> Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: + """ + Contains dimension and dtype for packed uint8 node features. + """ + return self._node_quantized_feature_info + + @property + def node_quantization_metadata( + self, + ) -> Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ]: + return self._node_quantization_metadata + @property def edge_feature_info( self, @@ -723,6 +768,57 @@ def _initialize_node_features( ) logger.info("Initialized node features for homogeneous graph to dataset") + def _initialize_node_quantized_features( + self, + node_partition_book: Union[PartitionBook, dict[NodeType, PartitionBook]], + partitioned_node_quantized_features: Optional[ + Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]] + ], + ) -> None: + """ + Initializes packed uint8 node features in a separate feature store. + """ + + node_quantized_features, node_quantized_feature_id_to_index = ( + _prepare_feature_data( + partition_book=node_partition_book, + partitioned_data=partitioned_node_quantized_features, + ) + ) + + if ( + node_quantized_features is None + or node_quantized_feature_id_to_index is None + ): + logger.info("Found no node quantized features to initialize") + return + + self._node_quantized_features = _build_feature_store( + feature_data=node_quantized_features, + id2idx=node_quantized_feature_id_to_index, + dtype=torch.uint8, + ) + + if isinstance(node_quantized_features, Mapping): + self._node_quantized_feature_info = {} + for node_type, features_per_node_type in node_quantized_features.items(): + assert not isinstance(node_type, EdgeType) + self._node_quantized_feature_info[node_type] = FeatureInfo( + dim=features_per_node_type.size(1), + dtype=features_per_node_type.dtype, + ) + logger.info( + f"Initialized node quantized features for heterogeneous graph to dataset with node types: {node_quantized_features.keys()}" + ) + else: + self._node_quantized_feature_info = FeatureInfo( + dim=node_quantized_features.size(1), + dtype=node_quantized_features.dtype, + ) + logger.info( + "Initialized node quantized features for homogeneous graph to dataset" + ) + def _initialize_node_labels( self, node_partition_book: Union[PartitionBook, dict[NodeType, PartitionBook]], @@ -810,6 +906,9 @@ def build( self, partition_output: PartitionOutput, splitter: Optional[Union[NodeSplitter, NodeAnchorLinkSplitter]] = None, + node_quantization_metadata: Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ] = None, ) -> None: """ Provided some partition graph information, this method stores these tensors inside of the class for @@ -824,6 +923,7 @@ def build( * a tuple of train, val, and test node ids, if heterogeneous * a dict[NodeType, tuple[train, val, test]] of node ids, if homogeneous Optional as not all datasets need to be split on, e.g. if we're doing inference. + node_quantization_metadata: Optional metadata for append-only packed node features. """ logger.info( f"Rank {self._rank} starting building dataset class from partitioned graph ..." @@ -897,6 +997,14 @@ def build( partition_output.partitioned_node_features = None gc.collect() + self._initialize_node_quantized_features( + node_partition_book=partition_output.node_partition_book, + partitioned_node_quantized_features=partition_output.partitioned_node_quantized_features, + ) + self._node_quantization_metadata = node_quantization_metadata + partition_output.partitioned_node_quantized_features = None + gc.collect() + self._initialize_node_labels( node_partition_book=partition_output.node_partition_book, partitioned_node_labels=partition_output.partitioned_node_labels, @@ -935,6 +1043,7 @@ def share_ipc( Literal["in", "out"], Optional[Union[Graph, dict[EdgeType, Graph]]], Optional[Union[Feature, dict[NodeType, Feature]]], + Optional[Union[Feature, dict[NodeType, Feature]]], Optional[Union[Feature, dict[EdgeType, Feature]]], Optional[Union[Feature, dict[NodeType, Feature]]], Optional[Union[PartitionBook, dict[NodeType, PartitionBook]]], @@ -946,6 +1055,10 @@ def share_ipc( Optional[Union[int, dict[NodeType, int]]], Optional[Union[int, dict[NodeType, int]]], Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]], + Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]], + Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ], Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]], Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]], Optional[int], @@ -959,6 +1072,7 @@ def share_ipc( Literal["in", "out"]: Graph Edge Direction Optional[Union[Graph, dict[EdgeType, Graph]]]: Partitioned Graph Data Optional[Union[Feature, dict[NodeType, Feature]]]: Partitioned Node Feature Data + Optional[Union[Feature, dict[NodeType, Feature]]]: Partitioned packed uint8 node feature data Optional[Union[Feature, dict[EdgeType, Feature]]]: Partitioned Edge Feature Data Optional[Union[Feature, dict[NodeType, Feature]]]: Node labels on the current machine. Will be a dict if heterogeneous. Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Node Partition Book Tensor @@ -970,6 +1084,8 @@ def share_ipc( Optional[Union[int, dict[NodeType, int]]]: Number of validation nodes on the current machine. Will be a dict if heterogeneous. Optional[Union[int, dict[NodeType, int]]]: Number of test nodes on the current machine. Will be a dict if heterogeneous. Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Node feature dim and its data type, will be a dict if heterogeneous + Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Packed uint8 node feature dim and dtype + Optional[Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]]]: Node quantization metadata Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Edge feature dim and its data type, will be a dict if heterogeneous Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Degree tensors Optional[int]: Optional per-anchor label cap for ABLP label fetching @@ -990,6 +1106,7 @@ def share_ipc( self._edge_dir, self._graph, self._node_features, + self._node_quantized_features, self._edge_features, self._node_labels, self._node_partition_book, @@ -1001,6 +1118,8 @@ def share_ipc( self._num_val, # Additional field unique to DistDataset class self._num_test, # Additional field unique to DistDataset class self._node_feature_info, # Additional field unique to DistDataset class + self._node_quantized_feature_info, # Additional field unique to DistDataset class + self._node_quantization_metadata, # Additional field unique to DistDataset class self._edge_feature_info, # Additional field unique to DistDataset class self._degree_tensor, # Additional field unique to DistDataset class self._max_labels_per_anchor_node, # Additional field unique to DistDataset class @@ -1213,6 +1332,32 @@ def _prepare_feature_data( return None, None +def _build_feature_store( + feature_data: Union[torch.Tensor, dict[_EntityType, torch.Tensor]], + id2idx: Union[TensorDataType, dict[_EntityType, TensorDataType]], + dtype: torch.dtype, +) -> Union[Feature, dict[_EntityType, Feature]]: + """Build a GLT Feature store without registering it as standard node features.""" + if isinstance(feature_data, Mapping): + assert isinstance(id2idx, Mapping) + return { + entity_key: Feature( + feature_tensor=features, + id2index=id2idx[entity_key], + with_gpu=False, + dtype=dtype, + ) + for entity_key, features in feature_data.items() + } + assert not isinstance(id2idx, Mapping) + return Feature( + feature_tensor=feature_data, + id2index=id2idx, + with_gpu=False, + dtype=dtype, + ) + + ## Pickling Registration # The serialization function (share_ipc) first pushes all member variable tensors # to the shared memory, and then packages all references to the tensors in one ipc @@ -1234,6 +1379,9 @@ def _rebuild_distributed_dataset( Optional[ Union[Feature, dict[NodeType, Feature]] ], # Partitioned Node Feature Data + Optional[ + Union[Feature, dict[NodeType, Feature]] + ], # Partitioned packed uint8 node feature data Optional[ Union[Feature, dict[EdgeType, Feature]] ], # Partitioned Edge Feature Data @@ -1257,6 +1405,12 @@ def _rebuild_distributed_dataset( Optional[ Union[FeatureInfo, dict[NodeType, FeatureInfo]] ], # Node feature dim and its data type + Optional[ + Union[FeatureInfo, dict[NodeType, FeatureInfo]] + ], # Packed uint8 node feature dim and dtype + Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ], # Node quantization metadata Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] ], # Edge feature dim and its data type diff --git a/gigl/distributed/dist_partitioner.py b/gigl/distributed/dist_partitioner.py index 2198557ba..3728740d3 100644 --- a/gigl/distributed/dist_partitioner.py +++ b/gigl/distributed/dist_partitioner.py @@ -150,6 +150,9 @@ def __init__( node_features: Optional[ Union[torch.Tensor, dict[NodeType, torch.Tensor]] ] = None, + node_quantized_features: Optional[ + Union[torch.Tensor, dict[NodeType, torch.Tensor]] + ] = None, edge_index: Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]] = None, edge_features: Optional[ Union[torch.Tensor, dict[EdgeType, torch.Tensor]] @@ -171,6 +174,7 @@ def __init__( should_assign_edges_by_src_node (bool): Whether edges should be assigned to the machine of the source nodes during partitioning node_ids (Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]): Optionally registered node ids from input. Tensors should be of shape [num_nodes_on_current_rank] node_features (Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]): Optionally registered node feats from input. Tensors should be of shope [num_nodes_on_current_rank, node_feat_dim] + node_quantized_features (Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]): Optionally registered packed uint8 node features from input. Tensors should be of shape [num_nodes_on_current_rank, packed_node_feat_dim] edge_index (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Optionally registered edge indexes from input. Tensors should be of shape [2, num_edges_on_current_rank] edge_features (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Optionally registered edge features from input. Tensors should be of shape [num_edges_on_current_rank, edge_feat_dim] positive_labels (Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]]): Optionally registered positive labels from input. Tensors should be of shape [2, num_pos_labels_on_current_rank] @@ -195,6 +199,8 @@ def __init__( self._max_node_ids: Optional[dict[NodeType, int]] = None self._node_feat: Optional[dict[NodeType, torch.Tensor]] = None self._node_feat_dim: Optional[dict[NodeType, int]] = None + self._node_quantized_feat: Optional[dict[NodeType, torch.Tensor]] = None + self._node_quantized_feat_dim: Optional[dict[NodeType, int]] = None self._node_labels: Optional[dict[NodeType, torch.Tensor]] = None self._node_labels_dim: Optional[dict[NodeType, int]] = None @@ -218,6 +224,11 @@ def __init__( if node_features is not None: self.register_node_features(node_features=node_features) + if node_quantized_features is not None: + self.register_node_quantized_features( + node_quantized_features=node_quantized_features + ) + if edge_features is not None: self.register_edge_features(edge_features=edge_features) @@ -541,6 +552,40 @@ def register_node_features( for node_type in input_node_features: self._node_feat_dim[node_type] = input_node_features[node_type].shape[1] + def register_node_quantized_features( + self, + node_quantized_features: Union[torch.Tensor, dict[NodeType, torch.Tensor]], + ) -> None: + """Registers packed uint8 node features to the partitioner.""" + + self._assert_and_get_rpc_setup() + + if self._node_quantized_feat is not None: + raise ValueError( + "Node quantized features have already been registered. Cannot re-register node quantized feature data." + ) + + logger.info("Registering Node Quantized Features ...") + + input_node_quantized_features = ( + self._convert_node_entity_to_heterogeneous_format( + input_node_entity=node_quantized_features + ) + ) + + assert input_node_quantized_features, ( + "Node quantized features is an empty dictionary. Please provide node quantized features to register." + ) + + self._node_quantized_feat = convert_to_tensor( + input_node_quantized_features, dtype=torch.uint8 + ) + self._node_quantized_feat_dim = {} + for node_type in input_node_quantized_features: + self._node_quantized_feat_dim[node_type] = ( + input_node_quantized_features[node_type].shape[1] + ) + def register_node_labels( self, node_labels: Union[torch.Tensor, dict[NodeType, torch.Tensor]] ) -> None: @@ -928,7 +973,11 @@ def _partition_node_features_and_labels( self, node_partition_book: dict[NodeType, PartitionBook], node_type: NodeType, - ) -> tuple[Optional[FeaturePartitionData], Optional[FeaturePartitionData]]: + ) -> tuple[ + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + ]: """ Partitions node features and labels according to the node partition book. @@ -941,6 +990,7 @@ def _partition_node_features_and_labels( Returns: Optional[FeaturePartitionData]: Partitioned data of node features for current node type. + Optional[FeaturePartitionData]: Partitioned data of packed uint8 node features for current node type. Optional[FeaturePartitionData]: Partitioned data of node labels for current node type. """ @@ -980,6 +1030,18 @@ def _partition_node_features_and_labels( if self._node_feat_dim is not None and node_type in self._node_feat_dim else None ) + node_quantized_features = ( + self._node_quantized_feat[node_type] + if self._node_quantized_feat is not None + and node_type in self._node_quantized_feat + else None + ) + node_quantized_feat_dim = ( + self._node_quantized_feat_dim[node_type] + if self._node_quantized_feat_dim is not None + and node_type in self._node_quantized_feat_dim + else None + ) node_labels = ( self._node_labels[node_type] @@ -992,22 +1054,29 @@ def _partition_node_features_and_labels( else None ) - if node_features is not None and node_labels is not None: - input_data: Tuple[torch.Tensor, ...] = ( - node_features, - node_labels, - node_ids, - ) - elif node_features is not None: - input_data = (node_features, node_ids) - elif node_labels is not None: - input_data = (node_labels, node_ids) - else: + input_parts: list[torch.Tensor] = [] + node_feature_ind: Optional[int] = None + node_quantized_feature_ind: Optional[int] = None + node_label_ind: Optional[int] = None + + if node_features is not None: + node_feature_ind = len(input_parts) + input_parts.append(node_features) + if node_quantized_features is not None: + node_quantized_feature_ind = len(input_parts) + input_parts.append(node_quantized_features) + if node_labels is not None: + node_label_ind = len(input_parts) + input_parts.append(node_labels) + if not input_parts: raise ValueError( - f"Found no node features or node labels to partition for node type {node_type}" + f"Found no node features, quantized node features, or node labels to partition for node type {node_type}" ) + input_parts.append(node_ids) + input_data: Tuple[torch.Tensor, ...] = tuple(input_parts) has_node_features = node_features is not None + has_node_quantized_features = node_quantized_features is not None has_node_labels = node_labels is not None def _node_feature_partition_fn(node_feature_ids, _): @@ -1028,12 +1097,21 @@ def _node_feature_partition_fn(node_feature_ids, _): self._remove_key_from_member_dict("_max_node_ids", node_type) self._remove_key_from_member_dict("_node_feat", node_type) self._remove_key_from_member_dict("_node_feat_dim", node_type) + self._remove_key_from_member_dict("_node_quantized_feat", node_type) + self._remove_key_from_member_dict("_node_quantized_feat_dim", node_type) self._remove_key_from_member_dict("_node_labels", node_type) self._remove_key_from_member_dict("_node_labels_dim", node_type) # Since the unpartitioned node ids, features and labels are large, we would like to delete them when # they are no longer needed to free memory. - del node_ids, num_nodes, max_node_ids, node_features, node_labels + del ( + node_ids, + num_nodes, + max_node_ids, + node_features, + node_quantized_features, + node_labels, + ) gc.collect() @@ -1049,18 +1127,31 @@ def _node_feature_partition_fn(node_feature_ids, _): # Partitioned node features are stored at the 0th index ineach tuple of the partitioned results. if has_node_features: + assert node_feature_ind is not None node_feature_partition_data = FeaturePartitionData( - feats=torch.cat([r[0] for r in partitioned_results]), + feats=torch.cat( + [r[node_feature_ind] for r in partitioned_results] + ), ids=partitioned_ids, ) else: node_feature_partition_data = None - # Partitioned node labels are stored at the 1st index in each tuple of the partitioned results if we have - # both node features and node labels. Otherwise, it is stored at the 0th index + if has_node_quantized_features: + assert node_quantized_feature_ind is not None + node_quantized_feature_partition_data = FeaturePartitionData( + feats=torch.cat( + [r[node_quantized_feature_ind] for r in partitioned_results] + ), + ids=partitioned_ids, + ) + + else: + node_quantized_feature_partition_data = None + if has_node_labels: - node_label_ind = 1 if has_node_features else 0 + assert node_label_ind is not None node_label_partition_data = FeaturePartitionData( feats=torch.cat([r[node_label_ind] for r in partitioned_results]), ids=partitioned_ids, @@ -1077,6 +1168,16 @@ def _node_feature_partition_fn(node_feature_ids, _): else: node_feature_partition_data = None + if node_quantized_feat_dim is not None: + node_quantized_feature_partition_data = FeaturePartitionData( + feats=torch.empty( + (0, node_quantized_feat_dim), dtype=torch.uint8 + ), + ids=torch.empty(0), + ) + else: + node_quantized_feature_partition_data = None + if node_labels_dim is not None: node_label_partition_data = FeaturePartitionData( feats=torch.empty((0, node_labels_dim)), @@ -1093,7 +1194,11 @@ def _node_feature_partition_fn(node_feature_ids, _): f"Node feature and label partitioning for node type {node_type} finished, took {time.time() - start_time:.3f}s" ) - return node_feature_partition_data, node_label_partition_data + return ( + node_feature_partition_data, + node_quantized_feature_partition_data, + node_label_partition_data, + ) def _partition_edge_index_and_edge_features( self, @@ -1459,6 +1564,7 @@ def partition_node_features_and_labels( ) -> tuple[ Optional[Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]]], Optional[Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]]], + Optional[Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]]], ]: """ Partitions node features and labels of a graph. If heterogeneous, partitions features and labels for all node type. @@ -1476,6 +1582,7 @@ def partition_node_features_and_labels( node_partition_book (Union[PartitionBook, dict[NodeType, PartitionBook]]): The Computed Node Partition Book Returns: Optional[Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]]]: Partitioned data of node features. + Optional[Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]]]: Partitioned data of packed uint8 node features. Optional[Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]]]: Partitioned data of node labels. """ assert self._num_nodes is not None and self._node_ids is not None, ( @@ -1509,22 +1616,32 @@ def partition_node_features_and_labels( ) for node_type in self._node_feat.keys(): node_feature_types.add(node_type) - elif self._node_labels is not None: + if self._node_quantized_feat is not None: + self._assert_data_type_consistency( + input_entity=self._node_quantized_feat, + is_node_entity=True, + is_subset=True, + ) + for node_type in self._node_quantized_feat.keys(): + node_feature_types.add(node_type) + if self._node_labels is not None: self._assert_data_type_consistency( input_entity=self._node_labels, is_node_entity=True, is_subset=True ) for node_type in self._node_labels.keys(): node_feature_types.add(node_type) - else: + if not node_feature_types: raise ValueError( - "Node features or labels must be registered prior to partitioning." + "Node features, quantized node features, or labels must be registered prior to partitioning." ) partitioned_node_features: dict[NodeType, FeaturePartitionData] = {} + partitioned_node_quantized_features: dict[NodeType, FeaturePartitionData] = {} partitioned_node_labels: dict[NodeType, FeaturePartitionData] = {} for node_type in sorted(node_feature_types): ( partitioned_node_features_for_node_type, + partitioned_node_quantized_features_for_node_type, partitioned_node_labels_for_node_type, ) = self._partition_node_features_and_labels( node_partition_book=transformed_node_partition_book, node_type=node_type @@ -1533,6 +1650,10 @@ def partition_node_features_and_labels( partitioned_node_features[node_type] = ( partitioned_node_features_for_node_type ) + if partitioned_node_quantized_features_for_node_type is not None: + partitioned_node_quantized_features[node_type] = ( + partitioned_node_quantized_features_for_node_type + ) if partitioned_node_labels_for_node_type is not None: partitioned_node_labels[node_type] = ( partitioned_node_labels_for_node_type @@ -1546,6 +1667,9 @@ def partition_node_features_and_labels( partitioned_node_features[DEFAULT_HOMOGENEOUS_NODE_TYPE] if partitioned_node_features else None, + partitioned_node_quantized_features[DEFAULT_HOMOGENEOUS_NODE_TYPE] + if partitioned_node_quantized_features + else None, partitioned_node_labels[DEFAULT_HOMOGENEOUS_NODE_TYPE] if partitioned_node_labels else None, @@ -1553,6 +1677,9 @@ def partition_node_features_and_labels( else: return ( partitioned_node_features if partitioned_node_features else None, + partitioned_node_quantized_features + if partitioned_node_quantized_features + else None, partitioned_node_labels if partitioned_node_labels else None, ) @@ -1771,15 +1898,21 @@ def partition( node_partition_book=node_partition_book ) - if self._node_feat is not None or self._node_labels is not None: + if ( + self._node_feat is not None + or self._node_quantized_feat is not None + or self._node_labels is not None + ): ( partitioned_node_features, + partitioned_node_quantized_features, partitioned_node_labels, ) = self.partition_node_features_and_labels( node_partition_book=node_partition_book ) else: partitioned_node_features = None + partitioned_node_quantized_features = None partitioned_node_labels = None if self._positive_label_edge_index is not None: @@ -1805,6 +1938,7 @@ def partition( edge_partition_book=edge_partition_book, partitioned_edge_index=partitioned_edge_index, partitioned_node_features=partitioned_node_features, + partitioned_node_quantized_features=partitioned_node_quantized_features, partitioned_edge_features=partitioned_edge_features, partitioned_positive_labels=partitioned_positive_edge_index, partitioned_negative_labels=partitioned_negative_edge_index, diff --git a/gigl/distributed/dist_range_partitioner.py b/gigl/distributed/dist_range_partitioner.py index faa1cfb57..b7b0754f7 100644 --- a/gigl/distributed/dist_range_partitioner.py +++ b/gigl/distributed/dist_range_partitioner.py @@ -128,7 +128,11 @@ def _partition_node_features_and_labels( self, node_partition_book: dict[NodeType, PartitionBook], node_type: NodeType, - ) -> tuple[Optional[FeaturePartitionData], Optional[FeaturePartitionData]]: + ) -> tuple[ + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + ]: """ Partitions node features according to the node partition book. We rely on the functionality from the parent tensor-based partitioner here, and add logic to sort the node features by node indices which is specific to range-based partitioning. This is done so that the range-based @@ -143,6 +147,7 @@ def _partition_node_features_and_labels( """ ( feature_partition_data, + quantized_feature_partition_data, labels_partition_data, ) = super()._partition_node_features_and_labels( node_partition_book=node_partition_book, node_type=node_type @@ -167,6 +172,22 @@ def _partition_node_features_and_labels( else: partitioned_node_feature_data = None + if quantized_feature_partition_data is not None: + ids = quantized_feature_partition_data.ids + assert ids is not None + sorted_node_ids_indices = torch.argsort(ids) + partitioned_node_quantized_features = ( + quantized_feature_partition_data.feats[sorted_node_ids_indices] + ) + partitioned_node_quantized_feature_data = FeaturePartitionData( + feats=partitioned_node_quantized_features, ids=None + ) + + del sorted_node_ids_indices + gc.collect() + else: + partitioned_node_quantized_feature_data = None + if labels_partition_data is not None: ids = labels_partition_data.ids assert ids is not None @@ -183,7 +204,11 @@ def _partition_node_features_and_labels( else: partitioned_node_label_data = None - return partitioned_node_feature_data, partitioned_node_label_data + return ( + partitioned_node_feature_data, + partitioned_node_quantized_feature_data, + partitioned_node_label_data, + ) def _partition_edge_index_and_edge_features( self, diff --git a/gigl/distributed/distributed_neighborloader.py b/gigl/distributed/distributed_neighborloader.py index 0ef76ac9e..e36f291c9 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -27,6 +27,7 @@ from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, + append_dequantized_node_features, extract_metadata, labeled_to_homogeneous, set_missing_features, @@ -408,6 +409,8 @@ def _setup_for_graph_store( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=node_feature_info, + node_quantized_feature_info=None, + node_quantization_metadata=None, edge_feature_info=edge_feature_info, edge_dir=dataset.fetch_edge_dir(), ), @@ -525,6 +528,8 @@ def _setup_for_colocated( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=dataset.node_feature_info, + node_quantized_feature_info=dataset.node_quantized_feature_info, + node_quantization_metadata=dataset.node_quantization_metadata, edge_feature_info=dataset.edge_feature_info, edge_dir=dataset.edge_dir, ), @@ -550,6 +555,11 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: data = labeled_to_homogeneous(DEFAULT_HOMOGENEOUS_EDGE_TYPE, data) data, metadata = self._apply_ppr_outputs(data, metadata) + data, metadata = append_dequantized_node_features( + data=data, + metadata=metadata, + node_quantization_metadata=self._node_quantization_metadata, + ) # Attach any remaining metadata (e.g. custom user-defined keys) directly onto the # data object so downstream code can access them via attribute lookup. diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 876582208..e0d38d702 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -13,11 +13,17 @@ from torch_geometric.typing import EdgeType, NodeType from gigl.common.logger import Logger -from gigl.types.graph import FeatureInfo, is_label_edge_type +from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_NODE_TYPE, + FeatureInfo, + NodeQuantizationMetadata, + is_label_edge_type, +) logger = Logger() _GraphType = TypeVar("_GraphType", Data, HeteroData) +NODE_QUANTIZED_FEATURES_METADATA_KEY = "node_quantized_features" class SamplingClusterSetup(Enum): @@ -42,6 +48,14 @@ class DatasetSchema: edge_types: Optional[list[EdgeType]] # Node feature info. node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]] + # Packed uint8 node feature info. + node_quantized_feature_info: Optional[ + Union[FeatureInfo, dict[NodeType, FeatureInfo]] + ] + # Quantization metadata for append-only packed node features. + node_quantization_metadata: Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ] # Edge feature info. edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]] # Edge direction. @@ -324,6 +338,155 @@ def set_missing_features( return data +def _unpack_quantized_features( + packed_features: torch.Tensor, + quantization_metadata: NodeQuantizationMetadata, +) -> torch.Tensor: + bits = quantization_metadata.bits + dequantized_feature_dim = quantization_metadata.dequantized_feature_dim + if bits not in {1, 2, 4, 8}: + raise ValueError(f"Expected bits to be one of 1, 2, 4, or 8, got {bits}.") + + if bits == 8: + if packed_features.size(1) < dequantized_feature_dim: + raise ValueError( + f"Packed feature dim {packed_features.size(1)} is smaller than dequantized dim {dequantized_feature_dim}." + ) + return packed_features[:, :dequantized_feature_dim].to(torch.float32) + + values_per_byte = 8 // bits + feature_indices = torch.arange( + dequantized_feature_dim, device=packed_features.device + ) + byte_indices = torch.div( + feature_indices, values_per_byte, rounding_mode="floor" + ) + bit_offsets = ((feature_indices % values_per_byte) * bits).to(torch.int16) + selected_bytes = packed_features[:, byte_indices].to(torch.int16) + return torch.bitwise_and( + torch.bitwise_right_shift(selected_bytes, bit_offsets), + (1 << bits) - 1, + ).to(torch.float32) + + +def _dequantize_node_features( + packed_features: torch.Tensor, + quantization_metadata: NodeQuantizationMetadata, +) -> torch.Tensor: + codes = _unpack_quantized_features(packed_features, quantization_metadata) + if quantization_metadata.bits == 1: + if ( + quantization_metadata.bucket_0_value is None + or quantization_metadata.bucket_1_value is None + ): + raise ValueError("Centroid quantization requires both bucket values.") + return torch.where( + codes.bool(), + torch.tensor( + quantization_metadata.bucket_1_value, + dtype=torch.float32, + device=packed_features.device, + ), + torch.tensor( + quantization_metadata.bucket_0_value, + dtype=torch.float32, + device=packed_features.device, + ), + ) + + if quantization_metadata.clip_min is None or quantization_metadata.clip_max is None: + raise ValueError("Linear quantization requires both clip bounds.") + levels = (1 << quantization_metadata.bits) - 1 + clip_min = torch.tensor( + quantization_metadata.clip_min, + dtype=torch.float32, + device=packed_features.device, + ) + clip_max = torch.tensor( + quantization_metadata.clip_max, + dtype=torch.float32, + device=packed_features.device, + ) + return clip_min + (codes / levels) * (clip_max - clip_min) + + +def append_dequantized_node_features( + data: _GraphType, + metadata: dict[str, torch.Tensor], + node_quantization_metadata: Optional[ + Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + ], +) -> tuple[_GraphType, dict[str, torch.Tensor]]: + """Append dequantized packed node features to PyG node feature tensors.""" + if node_quantization_metadata is None: + return data, metadata + + if isinstance(data, Data): + quantization_metadata = node_quantization_metadata + metadata_key = NODE_QUANTIZED_FEATURES_METADATA_KEY + if isinstance(node_quantization_metadata, dict): + quantization_metadata = node_quantization_metadata.get( + DEFAULT_HOMOGENEOUS_NODE_TYPE + ) + metadata_key = ( + f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}." + f"{DEFAULT_HOMOGENEOUS_NODE_TYPE}" + ) + if quantization_metadata is None: + return data, metadata + packed_features = metadata.pop(metadata_key, None) + if packed_features is None: + packed_features = metadata.pop(NODE_QUANTIZED_FEATURES_METADATA_KEY, None) + if packed_features is None: + return data, metadata + dequantized_features = _dequantize_node_features( + packed_features=packed_features, + quantization_metadata=quantization_metadata, + ) + data_x = data.x + if data_x is None: + data.x = dequantized_features + else: + if data_x.size(0) != dequantized_features.size(0): + raise ValueError( + "Cannot append dequantized features with " + f"{dequantized_features.size(0)} rows to existing x with " + f"{data_x.size(0)} rows." + ) + data.x = torch.cat([data_x.to(torch.float32), dequantized_features], dim=1) + return data, metadata + + if isinstance(node_quantization_metadata, NodeQuantizationMetadata): + raise ValueError( + "Expected per-node-type quantization metadata for heterogeneous data." + ) + + for node_type, quantization_metadata in node_quantization_metadata.items(): + metadata_key = f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{node_type}" + packed_features = metadata.pop(metadata_key, None) + if packed_features is None: + continue + dequantized_features = _dequantize_node_features( + packed_features=packed_features, + quantization_metadata=quantization_metadata, + ) + node_data = data[node_type] + node_x = getattr(node_data, "x", None) + if node_x is None: + node_data.x = dequantized_features + else: + if node_x.size(0) != dequantized_features.size(0): + raise ValueError( + "Cannot append dequantized features with " + f"{dequantized_features.size(0)} rows to existing x with " + f"{node_x.size(0)} rows." + ) + node_data.x = torch.cat( + [node_x.to(torch.float32), dequantized_features], dim=1 + ) + return data, metadata + + def extract_metadata( msg: SampleMessage, device: torch.device ) -> tuple[dict[str, torch.Tensor], SampleMessage]: diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index c774a57e6..5381787c6 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Optional, Tuple, Union from gigl.common import UriFactory from gigl.common.data.dataloaders import SerializedTFRecordInfo @@ -9,7 +9,7 @@ PreprocessedMetadataPbWrapper, ) from gigl.src.data_preprocessor.lib.types import FeatureSpecDict -from gigl.types.graph import to_homogeneous +from gigl.types.graph import NodeQuantizationMetadata, to_homogeneous from snapchat.research.gbml.preprocessed_metadata_pb2 import PreprocessedMetadata @@ -33,19 +33,75 @@ def _build_serialized_tfrecord_entity_info( Returns: SerializedTFRecordInfo: Stored metadata for current entity """ + quantized_feature_keys: list[str] = [] + quantized_feature_dim = 0 + if isinstance(preprocessed_metadata, PreprocessedMetadata.NodeMetadataOutput): + if preprocessed_metadata.HasField("quantized_feature_metadata"): + quantized_metadata = preprocessed_metadata.quantized_feature_metadata + quantized_feature_keys = list(quantized_metadata.packed_feature_keys) + quantized_feature_dim = quantized_metadata.packed_feature_dim + + stored_keys = set(preprocessed_metadata.feature_keys) + stored_keys.update(preprocessed_metadata.label_keys) + stored_keys.update(quantized_feature_keys) + filtered_feature_spec_dict = { + feature_key: feature_spec + for feature_key, feature_spec in feature_spec_dict.items() + if feature_key in stored_keys + } + return SerializedTFRecordInfo( tfrecord_uri_prefix=UriFactory.create_uri( preprocessed_metadata.tfrecord_uri_prefix ), feature_keys=list(preprocessed_metadata.feature_keys), - feature_spec=feature_spec_dict, + feature_spec=filtered_feature_spec_dict, feature_dim=preprocessed_metadata.feature_dim, entity_key=entity_key, + quantized_feature_keys=quantized_feature_keys, + quantized_feature_dim=quantized_feature_dim, label_keys=list(preprocessed_metadata.label_keys), tfrecord_uri_pattern=tfrecord_uri_pattern, ) +def _build_node_quantization_metadata( + node_metadata: PreprocessedMetadata.NodeMetadataOutput, +) -> Optional[NodeQuantizationMetadata]: + if not node_metadata.HasField("quantized_feature_metadata"): + return None + + quantized_metadata = node_metadata.quantized_feature_metadata + state = quantized_metadata.WhichOneof("state") + if quantized_metadata.bits == 1: + if state != "centroid": + raise ValueError("Expected centroid quantization state for 1-bit features.") + return NodeQuantizationMetadata( + bits=quantized_metadata.bits, + packed_feature_dim=quantized_metadata.packed_feature_dim, + dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, + dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), + bucket_0_value=quantized_metadata.centroid.bucket_0_value, + bucket_1_value=quantized_metadata.centroid.bucket_1_value, + ) + if quantized_metadata.bits not in {2, 4, 8}: + raise ValueError( + f"Expected quantized feature bits to be one of 1, 2, 4, or 8, got {quantized_metadata.bits}." + ) + if state != "linear": + raise ValueError( + f"Expected linear quantization state for {quantized_metadata.bits}-bit features." + ) + return NodeQuantizationMetadata( + bits=quantized_metadata.bits, + packed_feature_dim=quantized_metadata.packed_feature_dim, + dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, + dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), + clip_min=quantized_metadata.linear.clip_min, + clip_max=quantized_metadata.linear.clip_max, + ) + + def convert_pb_to_serialized_graph_metadata( preprocessed_metadata_pb_wrapper: PreprocessedMetadataPbWrapper, graph_metadata_pb_wrapper: GraphMetadataPbWrapper, @@ -65,6 +121,7 @@ def convert_pb_to_serialized_graph_metadata( edge_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} positive_label_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} negative_label_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} + node_quantization_metadata: dict[NodeType, NodeQuantizationMetadata] = {} preprocessed_metadata_pb = preprocessed_metadata_pb_wrapper.preprocessed_metadata_pb @@ -92,6 +149,9 @@ def convert_pb_to_serialized_graph_metadata( entity_key=node_key, tfrecord_uri_pattern=tfrecord_uri_pattern, ) + quantization_metadata = _build_node_quantization_metadata(node_metadata) + if quantization_metadata is not None: + node_quantization_metadata[node_type] = quantization_metadata for edge_type in graph_metadata_pb_wrapper.edge_types: condensed_edge_type = ( @@ -159,6 +219,9 @@ def convert_pb_to_serialized_graph_metadata( negative_label_entity_info=to_homogeneous(negative_label_entity_info) if len(negative_label_entity_info) > 0 else None, + node_quantization_metadata=to_homogeneous(node_quantization_metadata) + if len(node_quantization_metadata) > 0 + else None, ) else: return SerializedGraphMetadata( @@ -170,4 +233,7 @@ def convert_pb_to_serialized_graph_metadata( negative_label_entity_info=negative_label_entity_info if len(negative_label_entity_info) > 0 else None, + node_quantization_metadata=node_quantization_metadata + if len(node_quantization_metadata) > 0 + else None, ) diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index b49deb9f1..6fe17b965 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -1,7 +1,9 @@ from dataclasses import dataclass, field from typing import Callable, cast +import tensorflow as tf from tensorflow_metadata.proto.v0.schema_pb2 import Schema +from tensorflow_metadata.proto.v0.schema_pb2 import Feature as SchemaFeature import gigl.common.utils.local_fs as LocalFsUtils from gigl.common import GcsUri, LocalUri, Uri, UriFactory @@ -65,9 +67,25 @@ def __post_init__(self): for condensed_node_type, node_metadata in dict( self.preprocessed_metadata_pb.condensed_node_type_to_preprocessed_metadata ).items(): + dequantized_feature_keys: list[str] = [] + if node_metadata.HasField("quantized_feature_metadata"): + quantized_metadata = node_metadata.quantized_feature_metadata + dequantized_feature_keys = list( + quantized_metadata.dequantized_feature_keys + ) + if ( + quantized_metadata.dequantized_feature_dim + != len(dequantized_feature_keys) + ): + raise ValueError( + f"Expected quantized feature metadata for condensed node type {condensed_node_type} " + f"to have one dequantized_feature_key per output feature dimension, got " + f"{len(dequantized_feature_keys)} keys and dimension " + f"{quantized_metadata.dequantized_feature_dim}." + ) condensed_node_type_to_feature_dim_map[ CondensedNodeType(condensed_node_type) - ] = node_metadata.feature_dim + ] = node_metadata.feature_dim + len(dequantized_feature_keys) # Note that sorting the node feature/label keys breaks training with DDP. The root cause for why this is happening # is still under investigation. TODO (mkolodner-sc): Once the reason for why sorting the feature/label keys @@ -82,6 +100,12 @@ def __post_init__(self): ), feature_keys=node_feature_keys + label_keys, ) + if dequantized_feature_keys: + node_feature_schema = self.__append_dequantized_feature_schema( + feature_schema=node_feature_schema, + dequantized_feature_keys=dequantized_feature_keys, + label_keys=label_keys, + ) condensed_node_type_to_feature_schema_map[ CondensedNodeType(condensed_node_type) ] = node_feature_schema @@ -276,6 +300,52 @@ def __build_feature_schema( feature_vocab=feature_vocab, ) + def __append_dequantized_feature_schema( + self, + feature_schema: FeatureSchema, + dequantized_feature_keys: list[str], + label_keys: list[str], + ) -> FeatureSchema: + """Append synthetic fp32 schema entries for dequantized node features.""" + label_key_set = set(label_keys) + schema: FeatureSchemaDict = {} + feature_spec: FeatureSpecDict = {} + + for feature_key in dequantized_feature_keys: + if feature_key in feature_schema.feature_spec: + raise ValueError( + f"Dequantized feature key {feature_key} already exists in raw feature schema." + ) + + def add_dequantized_features() -> None: + for dequantized_feature_key in dequantized_feature_keys: + schema[dequantized_feature_key] = SchemaFeature( + name=dequantized_feature_key + ) + feature_spec[dequantized_feature_key] = tf.io.FixedLenFeature( + shape=[], dtype=tf.float32 + ) + + inserted_dequantized_features = False + for feature_key, raw_feature_spec in feature_schema.feature_spec.items(): + if feature_key in label_key_set and not inserted_dequantized_features: + add_dequantized_features() + inserted_dequantized_features = True + if feature_key in feature_spec: + raise ValueError(f"Duplicate feature key {feature_key} in schema.") + schema[feature_key] = feature_schema.schema[feature_key] + feature_spec[feature_key] = raw_feature_spec + + if not inserted_dequantized_features: + add_dequantized_features() + + return FeatureSchema( + schema=schema, + feature_spec=feature_spec, + feature_index=feature_spec_to_feature_index_map(feature_spec), + feature_vocab=feature_schema.feature_vocab, + ) + def __get_feature_to_vocab_list_map( self, transform_fn_assets_uri: Uri, diff --git a/gigl/types/graph.py b/gigl/types/graph.py index e7323693e..9324dbfcb 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -98,6 +98,12 @@ class PartitionOutput: Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]] ] + # Quantized node features on current rank. These are packed uint8 features + # aligned by node id and dequantized/appended in the sampler collate path. + partitioned_node_quantized_features: Optional[ + Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]] + ] = None + @dataclass(frozen=True) class FeatureInfo: @@ -107,6 +113,20 @@ class FeatureInfo: dtype: torch.dtype +@dataclass(frozen=True) +class NodeQuantizationMetadata: + """Metadata needed to unpack/dequantize append-only node features.""" + + bits: int + packed_feature_dim: int + dequantized_feature_dim: int + dequantized_feature_keys: tuple[str, ...] + clip_min: Optional[float] = None + clip_max: Optional[float] = None + bucket_0_value: Optional[float] = None + bucket_1_value: Optional[float] = None + + def _get_label_edges( labeled_edge_index: torch.Tensor, edge_dir: Literal["in", "out"], @@ -151,6 +171,10 @@ class LoadedGraphTensors: negative_label: Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]] # Unpartitioned Edge Weights (per-edge sampling weights, one scalar per edge) edge_weights: Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]] = None + # Unpartitioned packed uint8 node features. + node_quantized_features: Optional[ + Union[torch.Tensor, dict[NodeType, torch.Tensor]] + ] = None def treat_labels_as_edges(self, edge_dir: Literal["in", "out"]) -> None: """ diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index b1b7c9e15..5691025c6 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -3,6 +3,40 @@ syntax = "proto3"; package snapchat.research.gbml; message PreprocessedMetadata{ + message LinearQuantizationState{ + // Value produced by quantized code 0. Inputs are clipped to this value + // before linear quantization. + float clip_min = 1; + // Value produced by the largest quantized code. Inputs are clipped to this + // value before linear quantization. + float clip_max = 2; + } + + message CentroidQuantizationState{ + // Value produced by packed bit/code 0. + float bucket_0_value = 1; + // Value produced by packed bit/code 1. + float bucket_1_value = 2; + } + + message NodeQuantizedFeatureMetadata{ + // Packed uint8 fields in output TFRecords. + repeated string packed_feature_keys = 1; + // Model-visible feature names produced after dequantizing and appending. + repeated string dequantized_feature_keys = 2; + // Number of uint8 columns in the packed sidecar tensor. + uint32 packed_feature_dim = 3; + // Number of fp32 columns produced after unpacking/dequantizing. + uint32 dequantized_feature_dim = 4; + // 1 means centroid; 2, 4, and 8 mean linear buckets. + uint32 bits = 5; + + oneof state { + LinearQuantizationState linear = 6; + CentroidQuantizationState centroid = 7; + } + } + // Houses metadata about node TFTransform output from DataPreprocessor. message NodeMetadataOutput{ // The field in output TFRecords which references the node identifier. @@ -23,6 +57,8 @@ message PreprocessedMetadata{ optional uint32 feature_dim = 8; // Contains categorical feature vocabularies string transform_fn_assets_uri = 9; + // Optional append-only quantized node feature sidecar metadata. + NodeQuantizedFeatureMetadata quantized_feature_metadata = 10; } // Houses metadata of edge features output from DataPreprocessor diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.py b/snapchat/research/gbml/preprocessed_metadata_pb2.py index 8f2f818f3..23dd6fb05 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.py +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.py @@ -14,11 +14,14 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xed\x0b\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a\x9c\x02\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\tB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xd0\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1aK\n\x19\x43\x65ntroidQuantizationState\x12\x16\n\x0e\x62ucket_0_value\x18\x01 \x01(\x02\x12\x16\n\x0e\x62ucket_1_value\x18\x02 \x01(\x02\x1a\xe5\x02\n\x1cNodeQuantizedFeatureMetadata\x12\x1b\n\x13packed_feature_keys\x18\x01 \x03(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8b\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12m\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32I.snapchat.research.gbml.PreprocessedMetadata.NodeQuantizedFeatureMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') _PREPROCESSEDMETADATA = DESCRIPTOR.message_types_by_name['PreprocessedMetadata'] +_PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE = _PREPROCESSEDMETADATA.nested_types_by_name['LinearQuantizationState'] +_PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE = _PREPROCESSEDMETADATA.nested_types_by_name['CentroidQuantizationState'] +_PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA = _PREPROCESSEDMETADATA.nested_types_by_name['NodeQuantizedFeatureMetadata'] _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT = _PREPROCESSEDMETADATA.nested_types_by_name['NodeMetadataOutput'] _PREPROCESSEDMETADATA_EDGEMETADATAINFO = _PREPROCESSEDMETADATA.nested_types_by_name['EdgeMetadataInfo'] _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT = _PREPROCESSEDMETADATA.nested_types_by_name['EdgeMetadataOutput'] @@ -26,6 +29,27 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY = _PREPROCESSEDMETADATA.nested_types_by_name['CondensedEdgeTypeToPreprocessedMetadataEntry'] PreprocessedMetadata = _reflection.GeneratedProtocolMessageType('PreprocessedMetadata', (_message.Message,), { + 'LinearQuantizationState' : _reflection.GeneratedProtocolMessageType('LinearQuantizationState', (_message.Message,), { + 'DESCRIPTOR' : _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE, + '__module__' : 'snapchat.research.gbml.preprocessed_metadata_pb2' + # @@protoc_insertion_point(class_scope:snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationState) + }) + , + + 'CentroidQuantizationState' : _reflection.GeneratedProtocolMessageType('CentroidQuantizationState', (_message.Message,), { + 'DESCRIPTOR' : _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE, + '__module__' : 'snapchat.research.gbml.preprocessed_metadata_pb2' + # @@protoc_insertion_point(class_scope:snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationState) + }) + , + + 'NodeQuantizedFeatureMetadata' : _reflection.GeneratedProtocolMessageType('NodeQuantizedFeatureMetadata', (_message.Message,), { + 'DESCRIPTOR' : _PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA, + '__module__' : 'snapchat.research.gbml.preprocessed_metadata_pb2' + # @@protoc_insertion_point(class_scope:snapchat.research.gbml.PreprocessedMetadata.NodeQuantizedFeatureMetadata) + }) + , + 'NodeMetadataOutput' : _reflection.GeneratedProtocolMessageType('NodeMetadataOutput', (_message.Message,), { 'DESCRIPTOR' : _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT, '__module__' : 'snapchat.research.gbml.preprocessed_metadata_pb2' @@ -65,6 +89,9 @@ # @@protoc_insertion_point(class_scope:snapchat.research.gbml.PreprocessedMetadata) }) _sym_db.RegisterMessage(PreprocessedMetadata) +_sym_db.RegisterMessage(PreprocessedMetadata.LinearQuantizationState) +_sym_db.RegisterMessage(PreprocessedMetadata.CentroidQuantizationState) +_sym_db.RegisterMessage(PreprocessedMetadata.NodeQuantizedFeatureMetadata) _sym_db.RegisterMessage(PreprocessedMetadata.NodeMetadataOutput) _sym_db.RegisterMessage(PreprocessedMetadata.EdgeMetadataInfo) _sym_db.RegisterMessage(PreprocessedMetadata.EdgeMetadataOutput) @@ -79,15 +106,21 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._options = None _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_options = b'8\001' _PREPROCESSEDMETADATA._serialized_start=79 - _PREPROCESSEDMETADATA._serialized_end=1596 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=396 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=680 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=683 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=906 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=909 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1304 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1307 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=1450 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1453 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=1596 + _PREPROCESSEDMETADATA._serialized_end=2207 + _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_start=395 + _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_end=456 + _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_start=458 + _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=533 + _PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA._serialized_start=536 + _PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA._serialized_end=893 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=896 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1291 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1294 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1517 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1520 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1915 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1918 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2061 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2064 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2207 # @@protoc_insertion_point(module_scope) diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index 2271a7fec..a59fecd0a 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -20,6 +20,83 @@ DESCRIPTOR: google.protobuf.descriptor.FileDescriptor class PreprocessedMetadata(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor + class LinearQuantizationState(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CLIP_MIN_FIELD_NUMBER: builtins.int + CLIP_MAX_FIELD_NUMBER: builtins.int + clip_min: builtins.float + """Value produced by quantized code 0. Inputs are clipped to this value + before linear quantization.""" + clip_max: builtins.float + """Value produced by the largest quantized code. Inputs are clipped to this + value before linear quantization.""" + def __init__( + self, + *, + clip_min: builtins.float = ..., + clip_max: builtins.float = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["clip_max", b"clip_max", "clip_min", b"clip_min"]) -> None: ... + + class CentroidQuantizationState(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + BUCKET_0_VALUE_FIELD_NUMBER: builtins.int + BUCKET_1_VALUE_FIELD_NUMBER: builtins.int + bucket_0_value: builtins.float + """Value produced by packed bit/code 0.""" + bucket_1_value: builtins.float + """Value produced by packed bit/code 1.""" + def __init__( + self, + *, + bucket_0_value: builtins.float = ..., + bucket_1_value: builtins.float = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["bucket_0_value", b"bucket_0_value", "bucket_1_value", b"bucket_1_value"]) -> None: ... + + class NodeQuantizedFeatureMetadata(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PACKED_FEATURE_KEYS_FIELD_NUMBER: builtins.int + DEQUANTIZED_FEATURE_KEYS_FIELD_NUMBER: builtins.int + PACKED_FEATURE_DIM_FIELD_NUMBER: builtins.int + DEQUANTIZED_FEATURE_DIM_FIELD_NUMBER: builtins.int + BITS_FIELD_NUMBER: builtins.int + LINEAR_FIELD_NUMBER: builtins.int + CENTROID_FIELD_NUMBER: builtins.int + @property + def packed_feature_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Packed uint8 fields in output TFRecords.""" + @property + def dequantized_feature_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Model-visible feature names produced after dequantizing and appending.""" + packed_feature_dim: builtins.int + """Number of uint8 columns in the packed sidecar tensor.""" + dequantized_feature_dim: builtins.int + """Number of fp32 columns produced after unpacking/dequantizing.""" + bits: builtins.int + """1 means centroid; 2, 4, and 8 mean linear buckets.""" + @property + def linear(self) -> global___PreprocessedMetadata.LinearQuantizationState: ... + @property + def centroid(self) -> global___PreprocessedMetadata.CentroidQuantizationState: ... + def __init__( + self, + *, + packed_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., + dequantized_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., + packed_feature_dim: builtins.int = ..., + dequantized_feature_dim: builtins.int = ..., + bits: builtins.int = ..., + linear: global___PreprocessedMetadata.LinearQuantizationState | None = ..., + centroid: global___PreprocessedMetadata.CentroidQuantizationState | None = ..., + ) -> None: ... + def HasField(self, field_name: typing_extensions.Literal["centroid", b"centroid", "linear", b"linear"]) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_dim", b"packed_feature_dim", "packed_feature_keys", b"packed_feature_keys"]) -> None: ... + def WhichOneof(self, oneof_group: typing_extensions.Literal["state", b"state"]) -> typing_extensions.Literal["linear", "centroid"] | None: ... + class NodeMetadataOutput(google.protobuf.message.Message): """Houses metadata about node TFTransform output from DataPreprocessor.""" @@ -34,6 +111,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): ENUMERATED_NODE_DATA_BQ_TABLE_FIELD_NUMBER: builtins.int FEATURE_DIM_FIELD_NUMBER: builtins.int TRANSFORM_FN_ASSETS_URI_FIELD_NUMBER: builtins.int + QUANTIZED_FEATURE_METADATA_FIELD_NUMBER: builtins.int node_id_key: builtins.str """The field in output TFRecords which references the node identifier.""" @property @@ -54,6 +132,9 @@ class PreprocessedMetadata(google.protobuf.message.Message): """Feature dimension after preprocessing""" transform_fn_assets_uri: builtins.str """Contains categorical feature vocabularies""" + @property + def quantized_feature_metadata(self) -> global___PreprocessedMetadata.NodeQuantizedFeatureMetadata: + """Optional append-only quantized node feature sidecar metadata.""" def __init__( self, *, @@ -66,9 +147,10 @@ class PreprocessedMetadata(google.protobuf.message.Message): enumerated_node_data_bq_table: builtins.str = ..., feature_dim: builtins.int | None = ..., transform_fn_assets_uri: builtins.str = ..., + quantized_feature_metadata: global___PreprocessedMetadata.NodeQuantizedFeatureMetadata | None = ..., ) -> None: ... - def HasField(self, field_name: typing_extensions.Literal["_feature_dim", b"_feature_dim", "feature_dim", b"feature_dim"]) -> builtins.bool: ... - def ClearField(self, field_name: typing_extensions.Literal["_feature_dim", b"_feature_dim", "enumerated_node_data_bq_table", b"enumerated_node_data_bq_table", "enumerated_node_ids_bq_table", b"enumerated_node_ids_bq_table", "feature_dim", b"feature_dim", "feature_keys", b"feature_keys", "label_keys", b"label_keys", "node_id_key", b"node_id_key", "schema_uri", b"schema_uri", "tfrecord_uri_prefix", b"tfrecord_uri_prefix", "transform_fn_assets_uri", b"transform_fn_assets_uri"]) -> None: ... + def HasField(self, field_name: typing_extensions.Literal["_feature_dim", b"_feature_dim", "feature_dim", b"feature_dim", "quantized_feature_metadata", b"quantized_feature_metadata"]) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["_feature_dim", b"_feature_dim", "enumerated_node_data_bq_table", b"enumerated_node_data_bq_table", "enumerated_node_ids_bq_table", b"enumerated_node_ids_bq_table", "feature_dim", b"feature_dim", "feature_keys", b"feature_keys", "label_keys", b"label_keys", "node_id_key", b"node_id_key", "quantized_feature_metadata", b"quantized_feature_metadata", "schema_uri", b"schema_uri", "tfrecord_uri_prefix", b"tfrecord_uri_prefix", "transform_fn_assets_uri", b"transform_fn_assets_uri"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["_feature_dim", b"_feature_dim"]) -> typing_extensions.Literal["feature_dim"] | None: ... class EdgeMetadataInfo(google.protobuf.message.Message): From ee9cbc697dc2f550a96b9d43739878217703d698 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 15:57:04 +0000 Subject: [PATCH 02/44] Streamline schema parsing to accomodate synthetic feature keys/schema --- .../pb_wrappers/preprocessed_metadata.py | 87 +++++++------------ 1 file changed, 30 insertions(+), 57 deletions(-) diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index 6fe17b965..0886dda87 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -98,14 +98,9 @@ def __post_init__(self): transform_fn_assets_uri=UriFactory.create_uri( node_metadata.transform_fn_assets_uri ), - feature_keys=node_feature_keys + label_keys, + feature_keys=node_feature_keys + dequantized_feature_keys + label_keys, + synthetic_feature_keys=set(dequantized_feature_keys), ) - if dequantized_feature_keys: - node_feature_schema = self.__append_dequantized_feature_schema( - feature_schema=node_feature_schema, - dequantized_feature_keys=dequantized_feature_keys, - label_keys=label_keys, - ) condensed_node_type_to_feature_schema_map[ CondensedNodeType(condensed_node_type) ] = node_feature_schema @@ -228,25 +223,36 @@ def __post_init__(self): ) def __get_feature_spec_for_feature_keys( - self, feature_spec: FeatureSpecDict, feature_keys: list[str] + self, + feature_spec: FeatureSpecDict, + feature_keys: list[str], + synthetic_feature_keys: set[str], ) -> FeatureSpecDict: """ Return feature spec for the given feature keys """ - return {feature: feature_spec[feature] for feature in feature_keys} + return { + feature: tf.io.FixedLenFeature(shape=[], dtype=tf.float32) + if feature in synthetic_feature_keys + else feature_spec[feature] + for feature in feature_keys + } def __get_schema_for_feature_keys( self, feature_schema: Schema, feature_spec: FeatureSpecDict, feature_keys: list[str], + synthetic_feature_keys: set[str], ) -> FeatureSchemaDict: """ Return feature schema for the given feature keys """ all_features_in_feature_spec = list(feature_spec.keys()) return { - feature: feature_schema.feature[all_features_in_feature_spec.index(feature)] + feature: SchemaFeature(name=feature) + if feature in synthetic_feature_keys + else feature_schema.feature[all_features_in_feature_spec.index(feature)] for feature in feature_keys } @@ -255,6 +261,7 @@ def __build_feature_schema( schema_uri: Uri, transform_fn_assets_uri: Uri, feature_keys: list[str], + synthetic_feature_keys: set[str] | None = None, ) -> FeatureSchema: """ Return FeatureSchema NamedTuple for the given feature keys @@ -264,15 +271,26 @@ def __build_feature_schema( feature_spec will be based on the order of feature keys in preprocessed metadata which is also how SGS processes the features currently into float vectors """ + synthetic_feature_keys = synthetic_feature_keys or set() raw_feature_schema, raw_feature_spec = ( load_tf_schema_uri_str_to_feature_spec(uri=schema_uri) if schema_uri.uri else (None, {}) ) + for synthetic_feature_key in synthetic_feature_keys: + if synthetic_feature_key in raw_feature_spec: + raise ValueError( + f"Synthetic feature key {synthetic_feature_key} already exists " + "in raw feature schema." + ) + # Dequantized features are materialized later from packed uint8 tensors, + # so they need logical fp32 schema entries even though TFT did not write them. feature_spec = ( self.__get_feature_spec_for_feature_keys( - feature_spec=raw_feature_spec, feature_keys=feature_keys + feature_spec=raw_feature_spec, + feature_keys=feature_keys, + synthetic_feature_keys=synthetic_feature_keys, ) if feature_keys and raw_feature_schema else {} @@ -282,6 +300,7 @@ def __build_feature_schema( feature_schema=raw_feature_schema, feature_spec=raw_feature_spec, feature_keys=feature_keys, + synthetic_feature_keys=synthetic_feature_keys, ) if feature_keys and raw_feature_schema else {} @@ -300,52 +319,6 @@ def __build_feature_schema( feature_vocab=feature_vocab, ) - def __append_dequantized_feature_schema( - self, - feature_schema: FeatureSchema, - dequantized_feature_keys: list[str], - label_keys: list[str], - ) -> FeatureSchema: - """Append synthetic fp32 schema entries for dequantized node features.""" - label_key_set = set(label_keys) - schema: FeatureSchemaDict = {} - feature_spec: FeatureSpecDict = {} - - for feature_key in dequantized_feature_keys: - if feature_key in feature_schema.feature_spec: - raise ValueError( - f"Dequantized feature key {feature_key} already exists in raw feature schema." - ) - - def add_dequantized_features() -> None: - for dequantized_feature_key in dequantized_feature_keys: - schema[dequantized_feature_key] = SchemaFeature( - name=dequantized_feature_key - ) - feature_spec[dequantized_feature_key] = tf.io.FixedLenFeature( - shape=[], dtype=tf.float32 - ) - - inserted_dequantized_features = False - for feature_key, raw_feature_spec in feature_schema.feature_spec.items(): - if feature_key in label_key_set and not inserted_dequantized_features: - add_dequantized_features() - inserted_dequantized_features = True - if feature_key in feature_spec: - raise ValueError(f"Duplicate feature key {feature_key} in schema.") - schema[feature_key] = feature_schema.schema[feature_key] - feature_spec[feature_key] = raw_feature_spec - - if not inserted_dequantized_features: - add_dequantized_features() - - return FeatureSchema( - schema=schema, - feature_spec=feature_spec, - feature_index=feature_spec_to_feature_index_map(feature_spec), - feature_vocab=feature_schema.feature_vocab, - ) - def __get_feature_to_vocab_list_map( self, transform_fn_assets_uri: Uri, From 3d51b273b98cb7252eac961e158f0a9df888dda3 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 16:24:25 +0000 Subject: [PATCH 03/44] Rename FeatureQuantizationMetadata proto --- gigl/common/data/load_torch_tensors.py | 6 ++- gigl/distributed/dist_dataset.py | 30 +++++++++++---- gigl/distributed/utils/neighborloader.py | 16 +++++--- .../serialized_graph_metadata_translator.py | 29 +++++++------- .../pb_wrappers/preprocessed_metadata.py | 8 ++-- gigl/types/graph.py | 4 +- .../research/gbml/preprocessed_metadata.proto | 4 +- .../gbml/preprocessed_metadata_pb2.py | 38 +++++++++---------- .../gbml/preprocessed_metadata_pb2.pyi | 16 ++++---- 9 files changed, 86 insertions(+), 65 deletions(-) diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index f852da49b..5348f96ab 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -19,7 +19,7 @@ DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, LoadedGraphTensors, - NodeQuantizationMetadata, + FeatureQuantizationMetadata, ) from gigl.utils.share_memory import share_memory @@ -117,7 +117,9 @@ class SerializedGraphMetadata: ] = None # Optional node quantization metadata. node_quantization_metadata: Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] + ] ] = None diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index a2b958271..a951fddc4 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -24,7 +24,7 @@ FeatureInfo, FeaturePartitionData, GraphPartitionData, - NodeQuantizationMetadata, + FeatureQuantizationMetadata, PartitionOutput, ) from gigl.utils.data_splitters import ( @@ -85,7 +85,10 @@ def __init__( Union[FeatureInfo, dict[NodeType, FeatureInfo]] ] = None, node_quantization_metadata: Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] ] = None, edge_feature_info: Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] @@ -121,7 +124,7 @@ def __init__( node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension of node features and its data type, will be a dict if heterogeneous. Note this will be None in the homogeneous case if the data has no node features, or will only contain node types with node features in the heterogeneous case. node_quantized_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension and dtype for packed uint8 node features. - node_quantization_metadata: Optional[Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]]]: Metadata required for dequantizing packed node features. + node_quantization_metadata: Metadata for packed node features. edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Dimension of edge features and its data type, will be a dict if heterogeneous. Note this will be None in the homogeneous case if the data has no edge features, or will only contain edge types with edge features in the heterogeneous case. degree_tensor: Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Pre-computed degree tensor. Lazily computed on first access via the degree_tensor property. @@ -344,7 +347,9 @@ def node_quantized_feature_info( def node_quantization_metadata( self, ) -> Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] + ] ]: return self._node_quantization_metadata @@ -907,7 +912,10 @@ def build( partition_output: PartitionOutput, splitter: Optional[Union[NodeSplitter, NodeAnchorLinkSplitter]] = None, node_quantization_metadata: Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] ] = None, ) -> None: """ @@ -1057,7 +1065,10 @@ def share_ipc( Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]], Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]], Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] ], Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]], Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]], @@ -1085,7 +1096,7 @@ def share_ipc( Optional[Union[int, dict[NodeType, int]]]: Number of test nodes on the current machine. Will be a dict if heterogeneous. Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Node feature dim and its data type, will be a dict if heterogeneous Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Packed uint8 node feature dim and dtype - Optional[Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]]]: Node quantization metadata + Optional node quantization metadata. Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Edge feature dim and its data type, will be a dict if heterogeneous Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Degree tensors Optional[int]: Optional per-anchor label cap for ABLP label fetching @@ -1409,7 +1420,10 @@ def _rebuild_distributed_dataset( Union[FeatureInfo, dict[NodeType, FeatureInfo]] ], # Packed uint8 node feature dim and dtype Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] ], # Node quantization metadata Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index e0d38d702..9d8d04709 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -16,7 +16,7 @@ from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_NODE_TYPE, FeatureInfo, - NodeQuantizationMetadata, + FeatureQuantizationMetadata, is_label_edge_type, ) @@ -54,7 +54,9 @@ class DatasetSchema: ] # Quantization metadata for append-only packed node features. node_quantization_metadata: Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] + ] ] # Edge feature info. edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]] @@ -340,7 +342,7 @@ def set_missing_features( def _unpack_quantized_features( packed_features: torch.Tensor, - quantization_metadata: NodeQuantizationMetadata, + quantization_metadata: FeatureQuantizationMetadata, ) -> torch.Tensor: bits = quantization_metadata.bits dequantized_feature_dim = quantization_metadata.dequantized_feature_dim @@ -371,7 +373,7 @@ def _unpack_quantized_features( def _dequantize_node_features( packed_features: torch.Tensor, - quantization_metadata: NodeQuantizationMetadata, + quantization_metadata: FeatureQuantizationMetadata, ) -> torch.Tensor: codes = _unpack_quantized_features(packed_features, quantization_metadata) if quantization_metadata.bits == 1: @@ -414,7 +416,9 @@ def append_dequantized_node_features( data: _GraphType, metadata: dict[str, torch.Tensor], node_quantization_metadata: Optional[ - Union[NodeQuantizationMetadata, dict[NodeType, NodeQuantizationMetadata]] + Union[ + FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] + ] ], ) -> tuple[_GraphType, dict[str, torch.Tensor]]: """Append dequantized packed node features to PyG node feature tensors.""" @@ -456,7 +460,7 @@ def append_dequantized_node_features( data.x = torch.cat([data_x.to(torch.float32), dequantized_features], dim=1) return data, metadata - if isinstance(node_quantization_metadata, NodeQuantizationMetadata): + if isinstance(node_quantization_metadata, FeatureQuantizationMetadata): raise ValueError( "Expected per-node-type quantization metadata for heterogeneous data." ) diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index 5381787c6..1e7554317 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Tuple, Union from gigl.common import UriFactory from gigl.common.data.dataloaders import SerializedTFRecordInfo @@ -9,7 +9,7 @@ PreprocessedMetadataPbWrapper, ) from gigl.src.data_preprocessor.lib.types import FeatureSpecDict -from gigl.types.graph import NodeQuantizationMetadata, to_homogeneous +from gigl.types.graph import FeatureQuantizationMetadata, to_homogeneous from snapchat.research.gbml.preprocessed_metadata_pb2 import PreprocessedMetadata @@ -65,18 +65,14 @@ def _build_serialized_tfrecord_entity_info( ) -def _build_node_quantization_metadata( - node_metadata: PreprocessedMetadata.NodeMetadataOutput, -) -> Optional[NodeQuantizationMetadata]: - if not node_metadata.HasField("quantized_feature_metadata"): - return None - - quantized_metadata = node_metadata.quantized_feature_metadata +def _build_feature_quantization_metadata( + quantized_metadata: PreprocessedMetadata.FeatureQuantizationMetadata, +) -> FeatureQuantizationMetadata: state = quantized_metadata.WhichOneof("state") if quantized_metadata.bits == 1: if state != "centroid": raise ValueError("Expected centroid quantization state for 1-bit features.") - return NodeQuantizationMetadata( + return FeatureQuantizationMetadata( bits=quantized_metadata.bits, packed_feature_dim=quantized_metadata.packed_feature_dim, dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, @@ -92,7 +88,7 @@ def _build_node_quantization_metadata( raise ValueError( f"Expected linear quantization state for {quantized_metadata.bits}-bit features." ) - return NodeQuantizationMetadata( + return FeatureQuantizationMetadata( bits=quantized_metadata.bits, packed_feature_dim=quantized_metadata.packed_feature_dim, dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, @@ -121,7 +117,7 @@ def convert_pb_to_serialized_graph_metadata( edge_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} positive_label_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} negative_label_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} - node_quantization_metadata: dict[NodeType, NodeQuantizationMetadata] = {} + node_quantization_metadata: dict[NodeType, FeatureQuantizationMetadata] = {} preprocessed_metadata_pb = preprocessed_metadata_pb_wrapper.preprocessed_metadata_pb @@ -149,9 +145,12 @@ def convert_pb_to_serialized_graph_metadata( entity_key=node_key, tfrecord_uri_pattern=tfrecord_uri_pattern, ) - quantization_metadata = _build_node_quantization_metadata(node_metadata) - if quantization_metadata is not None: - node_quantization_metadata[node_type] = quantization_metadata + if node_metadata.HasField("quantized_feature_metadata"): + node_quantization_metadata[node_type] = ( + _build_feature_quantization_metadata( + node_metadata.quantized_feature_metadata + ) + ) for edge_type in graph_metadata_pb_wrapper.edge_types: condensed_edge_type = ( diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index 0886dda87..4448ce97d 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -2,8 +2,7 @@ from typing import Callable, cast import tensorflow as tf -from tensorflow_metadata.proto.v0.schema_pb2 import Schema -from tensorflow_metadata.proto.v0.schema_pb2 import Feature as SchemaFeature +from tensorflow_metadata.proto.v0.schema_pb2 import Feature, Schema import gigl.common.utils.local_fs as LocalFsUtils from gigl.common import GcsUri, LocalUri, Uri, UriFactory @@ -250,7 +249,7 @@ def __get_schema_for_feature_keys( """ all_features_in_feature_spec = list(feature_spec.keys()) return { - feature: SchemaFeature(name=feature) + feature: Feature(name=feature) if feature in synthetic_feature_keys else feature_schema.feature[all_features_in_feature_spec.index(feature)] for feature in feature_keys @@ -277,11 +276,12 @@ def __build_feature_schema( if schema_uri.uri else (None, {}) ) + for synthetic_feature_key in synthetic_feature_keys: if synthetic_feature_key in raw_feature_spec: raise ValueError( f"Synthetic feature key {synthetic_feature_key} already exists " - "in raw feature schema." + "in raw schema; dequantized keys must not overlap normal keys." ) # Dequantized features are materialized later from packed uint8 tensors, diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 9324dbfcb..7a6f513e1 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -114,8 +114,8 @@ class FeatureInfo: @dataclass(frozen=True) -class NodeQuantizationMetadata: - """Metadata needed to unpack/dequantize append-only node features.""" +class FeatureQuantizationMetadata: + """Metadata needed to unpack/dequantize append-only packed features.""" bits: int packed_feature_dim: int diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index 5691025c6..53092a5b5 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -19,7 +19,7 @@ message PreprocessedMetadata{ float bucket_1_value = 2; } - message NodeQuantizedFeatureMetadata{ + message FeatureQuantizationMetadata{ // Packed uint8 fields in output TFRecords. repeated string packed_feature_keys = 1; // Model-visible feature names produced after dequantizing and appending. @@ -58,7 +58,7 @@ message PreprocessedMetadata{ // Contains categorical feature vocabularies string transform_fn_assets_uri = 9; // Optional append-only quantized node feature sidecar metadata. - NodeQuantizedFeatureMetadata quantized_feature_metadata = 10; + FeatureQuantizationMetadata quantized_feature_metadata = 10; } // Houses metadata of edge features output from DataPreprocessor diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.py b/snapchat/research/gbml/preprocessed_metadata_pb2.py index 23dd6fb05..fbd2dc7be 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.py +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.py @@ -14,14 +14,14 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xd0\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1aK\n\x19\x43\x65ntroidQuantizationState\x12\x16\n\x0e\x62ucket_0_value\x18\x01 \x01(\x02\x12\x16\n\x0e\x62ucket_1_value\x18\x02 \x01(\x02\x1a\xe5\x02\n\x1cNodeQuantizedFeatureMetadata\x12\x1b\n\x13packed_feature_keys\x18\x01 \x03(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8b\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12m\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32I.snapchat.research.gbml.PreprocessedMetadata.NodeQuantizedFeatureMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xce\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1aK\n\x19\x43\x65ntroidQuantizationState\x12\x16\n\x0e\x62ucket_0_value\x18\x01 \x01(\x02\x12\x16\n\x0e\x62ucket_1_value\x18\x02 \x01(\x02\x1a\xe4\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1b\n\x13packed_feature_keys\x18\x01 \x03(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') _PREPROCESSEDMETADATA = DESCRIPTOR.message_types_by_name['PreprocessedMetadata'] _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE = _PREPROCESSEDMETADATA.nested_types_by_name['LinearQuantizationState'] _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE = _PREPROCESSEDMETADATA.nested_types_by_name['CentroidQuantizationState'] -_PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA = _PREPROCESSEDMETADATA.nested_types_by_name['NodeQuantizedFeatureMetadata'] +_PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA = _PREPROCESSEDMETADATA.nested_types_by_name['FeatureQuantizationMetadata'] _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT = _PREPROCESSEDMETADATA.nested_types_by_name['NodeMetadataOutput'] _PREPROCESSEDMETADATA_EDGEMETADATAINFO = _PREPROCESSEDMETADATA.nested_types_by_name['EdgeMetadataInfo'] _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT = _PREPROCESSEDMETADATA.nested_types_by_name['EdgeMetadataOutput'] @@ -43,10 +43,10 @@ }) , - 'NodeQuantizedFeatureMetadata' : _reflection.GeneratedProtocolMessageType('NodeQuantizedFeatureMetadata', (_message.Message,), { - 'DESCRIPTOR' : _PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA, + 'FeatureQuantizationMetadata' : _reflection.GeneratedProtocolMessageType('FeatureQuantizationMetadata', (_message.Message,), { + 'DESCRIPTOR' : _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA, '__module__' : 'snapchat.research.gbml.preprocessed_metadata_pb2' - # @@protoc_insertion_point(class_scope:snapchat.research.gbml.PreprocessedMetadata.NodeQuantizedFeatureMetadata) + # @@protoc_insertion_point(class_scope:snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadata) }) , @@ -91,7 +91,7 @@ _sym_db.RegisterMessage(PreprocessedMetadata) _sym_db.RegisterMessage(PreprocessedMetadata.LinearQuantizationState) _sym_db.RegisterMessage(PreprocessedMetadata.CentroidQuantizationState) -_sym_db.RegisterMessage(PreprocessedMetadata.NodeQuantizedFeatureMetadata) +_sym_db.RegisterMessage(PreprocessedMetadata.FeatureQuantizationMetadata) _sym_db.RegisterMessage(PreprocessedMetadata.NodeMetadataOutput) _sym_db.RegisterMessage(PreprocessedMetadata.EdgeMetadataInfo) _sym_db.RegisterMessage(PreprocessedMetadata.EdgeMetadataOutput) @@ -106,21 +106,21 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._options = None _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_options = b'8\001' _PREPROCESSEDMETADATA._serialized_start=79 - _PREPROCESSEDMETADATA._serialized_end=2207 + _PREPROCESSEDMETADATA._serialized_end=2205 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_start=395 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_end=456 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_start=458 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=533 - _PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA._serialized_start=536 - _PREPROCESSEDMETADATA_NODEQUANTIZEDFEATUREMETADATA._serialized_end=893 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=896 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1291 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1294 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1517 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1520 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1915 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1918 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2061 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2064 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2207 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_start=536 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=892 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=895 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1289 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1292 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1515 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1518 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1913 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1916 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2059 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2062 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2205 # @@protoc_insertion_point(module_scope) diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index a59fecd0a..8a6190fa2 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -27,10 +27,12 @@ class PreprocessedMetadata(google.protobuf.message.Message): CLIP_MAX_FIELD_NUMBER: builtins.int clip_min: builtins.float """Value produced by quantized code 0. Inputs are clipped to this value - before linear quantization.""" + before linear quantization. + """ clip_max: builtins.float """Value produced by the largest quantized code. Inputs are clipped to this - value before linear quantization.""" + value before linear quantization. + """ def __init__( self, *, @@ -56,7 +58,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["bucket_0_value", b"bucket_0_value", "bucket_1_value", b"bucket_1_value"]) -> None: ... - class NodeQuantizedFeatureMetadata(google.protobuf.message.Message): + class FeatureQuantizationMetadata(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor PACKED_FEATURE_KEYS_FIELD_NUMBER: builtins.int @@ -93,8 +95,8 @@ class PreprocessedMetadata(google.protobuf.message.Message): linear: global___PreprocessedMetadata.LinearQuantizationState | None = ..., centroid: global___PreprocessedMetadata.CentroidQuantizationState | None = ..., ) -> None: ... - def HasField(self, field_name: typing_extensions.Literal["centroid", b"centroid", "linear", b"linear"]) -> builtins.bool: ... - def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_dim", b"packed_feature_dim", "packed_feature_keys", b"packed_feature_keys"]) -> None: ... + def HasField(self, field_name: typing_extensions.Literal["centroid", b"centroid", "linear", b"linear", "state", b"state"]) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_dim", b"packed_feature_dim", "packed_feature_keys", b"packed_feature_keys", "state", b"state"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["state", b"state"]) -> typing_extensions.Literal["linear", "centroid"] | None: ... class NodeMetadataOutput(google.protobuf.message.Message): @@ -133,7 +135,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): transform_fn_assets_uri: builtins.str """Contains categorical feature vocabularies""" @property - def quantized_feature_metadata(self) -> global___PreprocessedMetadata.NodeQuantizedFeatureMetadata: + def quantized_feature_metadata(self) -> global___PreprocessedMetadata.FeatureQuantizationMetadata: """Optional append-only quantized node feature sidecar metadata.""" def __init__( self, @@ -147,7 +149,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): enumerated_node_data_bq_table: builtins.str = ..., feature_dim: builtins.int | None = ..., transform_fn_assets_uri: builtins.str = ..., - quantized_feature_metadata: global___PreprocessedMetadata.NodeQuantizedFeatureMetadata | None = ..., + quantized_feature_metadata: global___PreprocessedMetadata.FeatureQuantizationMetadata | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["_feature_dim", b"_feature_dim", "feature_dim", b"feature_dim", "quantized_feature_metadata", b"quantized_feature_metadata"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["_feature_dim", b"_feature_dim", "enumerated_node_data_bq_table", b"enumerated_node_data_bq_table", "enumerated_node_ids_bq_table", b"enumerated_node_ids_bq_table", "feature_dim", b"feature_dim", "feature_keys", b"feature_keys", "label_keys", b"label_keys", "node_id_key", b"node_id_key", "quantized_feature_metadata", b"quantized_feature_metadata", "schema_uri", b"schema_uri", "tfrecord_uri_prefix", b"tfrecord_uri_prefix", "transform_fn_assets_uri", b"transform_fn_assets_uri"]) -> None: ... From 9c9df99663e7203fec7633a9e139b9ced6e4001f Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 16:40:45 +0000 Subject: [PATCH 04/44] Refine translator --- .../serialized_graph_metadata_translator.py | 37 ++++++++++--------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index 1e7554317..e7e445fe8 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -44,6 +44,9 @@ def _build_serialized_tfrecord_entity_info( stored_keys = set(preprocessed_metadata.feature_keys) stored_keys.update(preprocessed_metadata.label_keys) stored_keys.update(quantized_feature_keys) + # PreprocessedMetadataPbWrapper may add synthetic dequantized feature keys + # to the logical feature schema. Only keep fields physically stored in + # TFRecords here; dequantized features are reconstructed later from uint8. filtered_feature_spec_dict = { feature_key: feature_spec for feature_key, feature_spec in feature_spec_dict.items() @@ -68,33 +71,31 @@ def _build_serialized_tfrecord_entity_info( def _build_feature_quantization_metadata( quantized_metadata: PreprocessedMetadata.FeatureQuantizationMetadata, ) -> FeatureQuantizationMetadata: + bits = quantized_metadata.bits state = quantized_metadata.WhichOneof("state") - if quantized_metadata.bits == 1: - if state != "centroid": - raise ValueError("Expected centroid quantization state for 1-bit features.") - return FeatureQuantizationMetadata( - bits=quantized_metadata.bits, - packed_feature_dim=quantized_metadata.packed_feature_dim, - dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, - dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), - bucket_0_value=quantized_metadata.centroid.bucket_0_value, - bucket_1_value=quantized_metadata.centroid.bucket_1_value, - ) - if quantized_metadata.bits not in {2, 4, 8}: + if bits not in {1, 2, 4, 8}: raise ValueError( - f"Expected quantized feature bits to be one of 1, 2, 4, or 8, got {quantized_metadata.bits}." + f"Expected quantized feature bits to be one of 1, 2, 4, or 8, got {bits}." ) - if state != "linear": + expected_state = "centroid" if bits == 1 else "linear" + if state != expected_state: raise ValueError( - f"Expected linear quantization state for {quantized_metadata.bits}-bit features." + f"Expected {expected_state} quantization state for {bits}-bit features." ) + return FeatureQuantizationMetadata( - bits=quantized_metadata.bits, + bits=bits, packed_feature_dim=quantized_metadata.packed_feature_dim, dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), - clip_min=quantized_metadata.linear.clip_min, - clip_max=quantized_metadata.linear.clip_max, + clip_min=quantized_metadata.linear.clip_min if bits != 1 else None, + clip_max=quantized_metadata.linear.clip_max if bits != 1 else None, + bucket_0_value=quantized_metadata.centroid.bucket_0_value + if bits == 1 + else None, + bucket_1_value=quantized_metadata.centroid.bucket_1_value + if bits == 1 + else None, ) From 631bad8c530af2ff636b7201112ba0042d69c8fd Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 17:42:59 +0000 Subject: [PATCH 05/44] Custom dtype in _concatenate_features_by_names --- gigl/common/data/dataloaders.py | 44 ++++++++++++--------------------- 1 file changed, 16 insertions(+), 28 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index df9d1b436..36bf50427 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -139,6 +139,7 @@ def _concatenate_features_by_names( feature_key_to_tf_tensor: dict[str, tf.Tensor], feature_keys: Sequence[str], label_keys: Sequence[str], + dtype: tf.dtypes.DType = tf.float32, ) -> tuple[Optional[tf.Tensor], Optional[tf.Tensor]]: """ Concatenates feature tensors in the order specified by feature names. @@ -150,6 +151,7 @@ def _concatenate_features_by_names( feature_key_to_tf_tensor (dict[str, tf.Tensor]): A dictionary mapping feature names to their corresponding tf tensors. feature_keys (Sequence[str]): A list of feature names specifying the order in which tensors should be concatenated. label_keys (Sequence[str]): Name of the label columns for the current entity. + dtype (tf.dtypes.DType): Type to cast concatenated tensors to. Returns: Tuple[ @@ -169,13 +171,15 @@ def _concatenate_features_by_names( for feature_key in feature_iterable: tensor = feature_key_to_tf_tensor[feature_key] - # TODO(kmonte, xgao, zfan): We will need to add support for this if we're trying to scale up. - # Features may be stored as int type. We cast it to float here and will need to subsequently convert - # it back to int. Note that this is ok for small int values (less than 2^24, or ~16 million). - # For large int values, we will need to round it when converting back - # from float, as otherwise there will be precision loss. - if tensor.dtype != tf.float32: - tensor = tf.cast(tensor, tf.float32) + if tensor.dtype != dtype: + # TODO(kmonte, xgao, zfan): We will need to add support for this + # if we're trying to scale up. Features may be stored as int type. + # We cast it to float here and will need to subsequently convert + # it back to int. Note that this is ok for small int values (less + # than 2^24, or ~16 million). + # For large int values, we will need to round it when converting back + # from float, as otherwise there will be precision loss. + tensor = tf.cast(tensor, dtype) # Reshape 1D tensor to column vector if len(tensor.shape) == 1: @@ -192,25 +196,6 @@ def _concatenate_features_by_names( return _get_labels_from_features(combined_feature_tensor, label_dim) -def _concatenate_quantized_features_by_names( - feature_key_to_tf_tensor: dict[str, tf.Tensor], - feature_keys: Sequence[str], -) -> Optional[tf.Tensor]: - """Concatenate packed quantized feature tensors as uint8 columns.""" - if not feature_keys: - return None - - features: list[tf.Tensor] = [] - for feature_key in feature_keys: - tensor = feature_key_to_tf_tensor[feature_key] - if tensor.dtype != tf.uint8: - tensor = tf.cast(tensor, tf.uint8) - if len(tensor.shape) == 1: - tensor = tf.expand_dims(tensor, axis=-1) - features.append(tensor) - return tf.concat(features, axis=1) - - def _tf_tensor_to_torch_tensor(tf_tensor: tf.Tensor) -> torch.Tensor: """ Converts a TensorFlow tensor to a PyTorch tensor using DLPack to ensure zero-copy conversion. @@ -516,8 +501,11 @@ def load_as_torch_tensors( if label_tensor is not None: label_tensors.append(label_tensor) if quantized_feature_keys: - quantized_feature_tensor = _concatenate_quantized_features_by_names( - batch, quantized_feature_keys + quantized_feature_tensor, _ = _concatenate_features_by_names( + batch, + quantized_feature_keys, + label_keys=[], + dtype=tf.uint8, ) if quantized_feature_tensor is not None: quantized_feature_tensors.append(quantized_feature_tensor) From e4dafa044c90b7675cf5ba5dd164e471fc3de1b4 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 17:46:37 +0000 Subject: [PATCH 06/44] Simplify cast --- gigl/common/data/dataloaders.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 36bf50427..dde351e7f 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -171,14 +171,12 @@ def _concatenate_features_by_names( for feature_key in feature_iterable: tensor = feature_key_to_tf_tensor[feature_key] + # TODO(kmonte, xgao, zfan): We will need to add support for this if we're trying to scale up. + # Features may be stored as int type. We cast it to float here and will need to subsequently convert + # it back to int. Note that this is ok for small int values (less than 2^24, or ~16 million). + # For large int values, we will need to round it when converting back + # from float, as otherwise there will be precision loss. if tensor.dtype != dtype: - # TODO(kmonte, xgao, zfan): We will need to add support for this - # if we're trying to scale up. Features may be stored as int type. - # We cast it to float here and will need to subsequently convert - # it back to int. Note that this is ok for small int values (less - # than 2^24, or ~16 million). - # For large int values, we will need to round it when converting back - # from float, as otherwise there will be precision loss. tensor = tf.cast(tensor, dtype) # Reshape 1D tensor to column vector From bc639add4b163a820a0ecfdc6ff359bfc9d7ba0d Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 18:11:51 +0000 Subject: [PATCH 07/44] Refine error checks --- gigl/common/data/dataloaders.py | 2 ++ gigl/common/data/load_torch_tensors.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index dde351e7f..0b29febdd 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -410,6 +410,8 @@ def load_as_torch_tensors( logger.info( f"Injecting quantized feature key {feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape={feature_shape}, dtype=tf.int64)`" ) + # TODO(jchmura-sc): Serialize uint8 packed features as raw + # bytes, then decode string as uint8 to avoid int64. feature_spec_dict[feature_key] = tf.io.FixedLenFeature( shape=feature_shape, dtype=tf.int64 ) diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index 5348f96ab..e0464ce1d 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -201,6 +201,13 @@ def _data_loading_process( raise NotImplementedError( "Label keys are not supported for edge entities" ) + if ( + serialized_entity_tf_record_info.quantized_feature_keys + and not serialized_entity_tf_record_info.is_node_entity + ): + raise NotImplementedError( + "Quantized 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, From c3151f35bd568df820f2e1e0d5e0042b527248a9 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 18:14:36 +0000 Subject: [PATCH 08/44] Rename node quant key --- gigl/common/data/load_torch_tensors.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index e0464ce1d..7dbd61608 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -30,7 +30,7 @@ _LABEL_FMT = "{entity}_labels" _EDGE_WEIGHTS_KEY = "edge_weights" _NODE_KEY = "node" -_NODE_QUANTIZED_KEY = "node_quantized" +_NODE_QUANTIZED_KEY = "node_quantized_features" def _extract_weight_col( @@ -344,7 +344,7 @@ def _data_loading_process( list(features.values())[0] if is_input_homogeneous else features ) if quantized_features and entity_type == _NODE_KEY: - output_dict[_FEATURE_FMT.format(entity=_NODE_QUANTIZED_KEY)] = ( + output_dict[_NODE_QUANTIZED_KEY] = ( list(quantized_features.values())[0] if is_input_homogeneous else quantized_features @@ -519,9 +519,7 @@ 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( - _FEATURE_FMT.format(entity=_NODE_QUANTIZED_KEY), None - ) + node_quantized_features = node_output_dict.get(_NODE_QUANTIZED_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)] From 5dc158a86abbe840341af74927eef4e6b673ba40 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 18:49:50 +0000 Subject: [PATCH 09/44] Considering setting metadata with setter to keep build() partition-focused --- gigl/distributed/dataset_factory.py | 4 +++- gigl/distributed/dist_dataset.py | 20 ++++++++++++-------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/gigl/distributed/dataset_factory.py b/gigl/distributed/dataset_factory.py index dc9e154ff..70e8567e6 100644 --- a/gigl/distributed/dataset_factory.py +++ b/gigl/distributed/dataset_factory.py @@ -227,7 +227,9 @@ def _load_and_build_partitioned_dataset( dataset.build( partition_output=partition_output, splitter=splitter, - node_quantization_metadata=serialized_graph_metadata.node_quantization_metadata, + ) + dataset.node_quantization_metadata = ( + serialized_graph_metadata.node_quantization_metadata ) return dataset diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index a951fddc4..3c9759737 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -353,6 +353,18 @@ def node_quantization_metadata( ]: return self._node_quantization_metadata + @node_quantization_metadata.setter + def node_quantization_metadata( + self, + new_node_quantization_metadata: Optional[ + Union[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] + ], + ) -> None: + self._node_quantization_metadata = new_node_quantization_metadata + @property def edge_feature_info( self, @@ -911,12 +923,6 @@ def build( self, partition_output: PartitionOutput, splitter: Optional[Union[NodeSplitter, NodeAnchorLinkSplitter]] = None, - node_quantization_metadata: Optional[ - Union[ - FeatureQuantizationMetadata, - dict[NodeType, FeatureQuantizationMetadata], - ] - ] = None, ) -> None: """ Provided some partition graph information, this method stores these tensors inside of the class for @@ -931,7 +937,6 @@ def build( * a tuple of train, val, and test node ids, if heterogeneous * a dict[NodeType, tuple[train, val, test]] of node ids, if homogeneous Optional as not all datasets need to be split on, e.g. if we're doing inference. - node_quantization_metadata: Optional metadata for append-only packed node features. """ logger.info( f"Rank {self._rank} starting building dataset class from partitioned graph ..." @@ -1009,7 +1014,6 @@ def build( node_partition_book=partition_output.node_partition_book, partitioned_node_quantized_features=partition_output.partitioned_node_quantized_features, ) - self._node_quantization_metadata = node_quantization_metadata partition_output.partitioned_node_quantized_features = None gc.collect() From 0041a141285ded031fbf85fdc2d81263037b2824 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 18:55:41 +0000 Subject: [PATCH 10/44] Pass metdata through init instead of setter --- gigl/distributed/dataset_factory.py | 10 ++++++---- gigl/distributed/dist_dataset.py | 12 ------------ 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/gigl/distributed/dataset_factory.py b/gigl/distributed/dataset_factory.py index 70e8567e6..1a5f859b5 100644 --- a/gigl/distributed/dataset_factory.py +++ b/gigl/distributed/dataset_factory.py @@ -222,15 +222,17 @@ def _load_and_build_partitioned_dataset( partition_output = partitioner.partition() logger.info(f"Initializing DistDataset instance with edge direction {edge_dir}") - dataset = DistDataset(rank=rank, world_size=world_size, edge_dir=edge_dir) + dataset = DistDataset( + rank=rank, + world_size=world_size, + edge_dir=edge_dir, + node_quantization_metadata=serialized_graph_metadata.node_quantization_metadata, + ) dataset.build( partition_output=partition_output, splitter=splitter, ) - dataset.node_quantization_metadata = ( - serialized_graph_metadata.node_quantization_metadata - ) return dataset diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index 3c9759737..0af1efd64 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -353,18 +353,6 @@ def node_quantization_metadata( ]: return self._node_quantization_metadata - @node_quantization_metadata.setter - def node_quantization_metadata( - self, - new_node_quantization_metadata: Optional[ - Union[ - FeatureQuantizationMetadata, - dict[NodeType, FeatureQuantizationMetadata], - ] - ], - ) -> None: - self._node_quantization_metadata = new_node_quantization_metadata - @property def edge_feature_info( self, From 16182b22104723604fa2ac42d0908df2b7b27eb8 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 20:21:33 +0000 Subject: [PATCH 11/44] Argument order --- gigl/distributed/dist_dataset.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index 0af1efd64..25f71de6f 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -107,6 +107,8 @@ def __init__( rank (int): Rank of the current process world_size (int): World size of the current process edge_dir (Literal["in", "out"]): Edge direction of the provied graph + node_quantization_metadata: Metadata for packed node features. + May be provided during initial construction or IPC rebuild. The below arguments are only expected to be provided when re-serializing an instance of the DistDataset class after build() has been called graph_partition (Optional[Union[Graph, dict[EdgeType, Graph]]]): Partitioned Graph Data node_feature_partition (Optional[Union[Feature, dict[NodeType, Feature]]]): Partitioned Node Feature Data @@ -124,7 +126,6 @@ def __init__( node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension of node features and its data type, will be a dict if heterogeneous. Note this will be None in the homogeneous case if the data has no node features, or will only contain node types with node features in the heterogeneous case. node_quantized_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Dimension and dtype for packed uint8 node features. - node_quantization_metadata: Metadata for packed node features. edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Dimension of edge features and its data type, will be a dict if heterogeneous. Note this will be None in the homogeneous case if the data has no edge features, or will only contain edge types with edge features in the heterogeneous case. degree_tensor: Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Pre-computed degree tensor. Lazily computed on first access via the degree_tensor property. From 6176e3f0193d916733c970cde87b7cb4418f34c7 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 20:54:47 +0000 Subject: [PATCH 12/44] Factor out feature quant specific ops --- gigl/common/utils/feature_quantization.py | 76 ++++++++++ gigl/distributed/base_sampler.py | 4 +- gigl/distributed/dist_ablp_neighborloader.py | 4 +- .../distributed/distributed_neighborloader.py | 4 +- gigl/distributed/sampler.py | 1 + gigl/distributed/utils/neighborloader.py | 140 ++++-------------- 6 files changed, 114 insertions(+), 115 deletions(-) create mode 100644 gigl/common/utils/feature_quantization.py diff --git a/gigl/common/utils/feature_quantization.py b/gigl/common/utils/feature_quantization.py new file mode 100644 index 000000000..73a2d888d --- /dev/null +++ b/gigl/common/utils/feature_quantization.py @@ -0,0 +1,76 @@ +"""Tensor utilities for feature quantization.""" + +import torch + +from gigl.types.graph import FeatureQuantizationMetadata + + +def _unpack_quantized_features( + packed_features: torch.Tensor, + quantization_metadata: FeatureQuantizationMetadata, +) -> torch.Tensor: + bits = quantization_metadata.bits + dequantized_feature_dim = quantization_metadata.dequantized_feature_dim + if bits not in {1, 2, 4, 8}: + raise ValueError(f"Expected bits to be one of 1, 2, 4, or 8, got {bits}.") + + if bits == 8: + if packed_features.size(1) < dequantized_feature_dim: + raise ValueError( + f"Packed feature dim {packed_features.size(1)} is smaller than " + f"dequantized dim {dequantized_feature_dim}." + ) + return packed_features[:, :dequantized_feature_dim].to(torch.float32) + + values_per_byte = 8 // bits + feature_indices = torch.arange( + dequantized_feature_dim, device=packed_features.device + ) + byte_indices = torch.div(feature_indices, values_per_byte, rounding_mode="floor") + bit_offsets = ((feature_indices % values_per_byte) * bits).to(torch.int16) + selected_bytes = packed_features[:, byte_indices].to(torch.int16) + return torch.bitwise_and( + torch.bitwise_right_shift(selected_bytes, bit_offsets), + (1 << bits) - 1, + ).to(torch.float32) + + +def dequantize_feature_tensor( + packed_features: torch.Tensor, + quantization_metadata: FeatureQuantizationMetadata, +) -> torch.Tensor: + codes = _unpack_quantized_features(packed_features, quantization_metadata) + if quantization_metadata.bits == 1: + if ( + quantization_metadata.bucket_0_value is None + or quantization_metadata.bucket_1_value is None + ): + raise ValueError("Centroid quantization requires both bucket values.") + return torch.where( + codes.bool(), + torch.tensor( + quantization_metadata.bucket_1_value, + dtype=torch.float32, + device=packed_features.device, + ), + torch.tensor( + quantization_metadata.bucket_0_value, + dtype=torch.float32, + device=packed_features.device, + ), + ) + + if quantization_metadata.clip_min is None or quantization_metadata.clip_max is None: + raise ValueError("Linear quantization requires both clip bounds.") + levels = (1 << quantization_metadata.bits) - 1 + clip_min = torch.tensor( + quantization_metadata.clip_min, + dtype=torch.float32, + device=packed_features.device, + ) + clip_max = torch.tensor( + quantization_metadata.clip_max, + dtype=torch.float32, + device=packed_features.device, + ) + return clip_min + (codes / levels) * (clip_max - clip_min) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 1e3a97d77..05638b43e 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -16,11 +16,11 @@ from graphlearn_torch.utils import reverse_edge_type from gigl.distributed.sampler import ( + NODE_QUANTIZED_FEATURES_METADATA_KEY, NEGATIVE_LABEL_METADATA_KEY, POSITIVE_LABEL_METADATA_KEY, ABLPNodeSamplerInput, ) -from gigl.distributed.utils.neighborloader import NODE_QUANTIZED_FEATURES_METADATA_KEY from gigl.utils.data_splitters import PADDING_NODE @@ -101,6 +101,8 @@ def __init__(self, *args, **kwargs): and data is not None and getattr(data, "node_quantized_features", None) is not None ): + # Mirrors GLT's dist_node_feature initialization: + # https://github.com/alibaba/graphlearn-for-pytorch/blob/88ff111ac0d9e45c6c9d2d18cfc5883dca07e9f9/graphlearn_torch/python/distributed/dist_neighbor_sampler.py#L162-L167 self.dist_node_quantized_feature = DistFeature( data.num_partitions, data.partition_idx, diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 757eefcd7..c4ca79669 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -30,10 +30,10 @@ from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, - append_dequantized_node_features, extract_edge_type_metadata, extract_metadata, labeled_to_homogeneous, + materialize_quantized_node_features, set_missing_features, shard_nodes_by_process, strip_label_edges, @@ -881,7 +881,7 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: data = self._set_labels(data, positive_labels, negative_labels) data, metadata = self._apply_ppr_outputs(data, metadata) - data, metadata = append_dequantized_node_features( + data, metadata = materialize_quantized_node_features( data=data, metadata=metadata, node_quantization_metadata=self._node_quantization_metadata, diff --git a/gigl/distributed/distributed_neighborloader.py b/gigl/distributed/distributed_neighborloader.py index e36f291c9..a2486af17 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -27,9 +27,9 @@ from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, - append_dequantized_node_features, extract_metadata, labeled_to_homogeneous, + materialize_quantized_node_features, set_missing_features, shard_nodes_by_process, strip_label_edges, @@ -555,7 +555,7 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: data = labeled_to_homogeneous(DEFAULT_HOMOGENEOUS_EDGE_TYPE, data) data, metadata = self._apply_ppr_outputs(data, metadata) - data, metadata = append_dequantized_node_features( + data, metadata = materialize_quantized_node_features( data=data, metadata=metadata, node_quantization_metadata=self._node_quantization_metadata, diff --git a/gigl/distributed/sampler.py b/gigl/distributed/sampler.py index e99dd65dc..a216fa41c 100644 --- a/gigl/distributed/sampler.py +++ b/gigl/distributed/sampler.py @@ -8,6 +8,7 @@ POSITIVE_LABEL_METADATA_KEY: Final[str] = "gigl_positive_labels." NEGATIVE_LABEL_METADATA_KEY: Final[str] = "gigl_negative_labels." +NODE_QUANTIZED_FEATURES_METADATA_KEY: Final[str] = "node_quantized_features" class ABLPNodeSamplerInput(NodeSamplerInput): diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 9d8d04709..eb5ef9fb0 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -13,6 +13,8 @@ from torch_geometric.typing import EdgeType, NodeType from gigl.common.logger import Logger +from gigl.common.utils.feature_quantization import dequantize_feature_tensor +from gigl.distributed.sampler import NODE_QUANTIZED_FEATURES_METADATA_KEY from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_NODE_TYPE, FeatureInfo, @@ -23,7 +25,6 @@ logger = Logger() _GraphType = TypeVar("_GraphType", Data, HeteroData) -NODE_QUANTIZED_FEATURES_METADATA_KEY = "node_quantized_features" class SamplingClusterSetup(Enum): @@ -340,79 +341,7 @@ def set_missing_features( return data -def _unpack_quantized_features( - packed_features: torch.Tensor, - quantization_metadata: FeatureQuantizationMetadata, -) -> torch.Tensor: - bits = quantization_metadata.bits - dequantized_feature_dim = quantization_metadata.dequantized_feature_dim - if bits not in {1, 2, 4, 8}: - raise ValueError(f"Expected bits to be one of 1, 2, 4, or 8, got {bits}.") - - if bits == 8: - if packed_features.size(1) < dequantized_feature_dim: - raise ValueError( - f"Packed feature dim {packed_features.size(1)} is smaller than dequantized dim {dequantized_feature_dim}." - ) - return packed_features[:, :dequantized_feature_dim].to(torch.float32) - - values_per_byte = 8 // bits - feature_indices = torch.arange( - dequantized_feature_dim, device=packed_features.device - ) - byte_indices = torch.div( - feature_indices, values_per_byte, rounding_mode="floor" - ) - bit_offsets = ((feature_indices % values_per_byte) * bits).to(torch.int16) - selected_bytes = packed_features[:, byte_indices].to(torch.int16) - return torch.bitwise_and( - torch.bitwise_right_shift(selected_bytes, bit_offsets), - (1 << bits) - 1, - ).to(torch.float32) - - -def _dequantize_node_features( - packed_features: torch.Tensor, - quantization_metadata: FeatureQuantizationMetadata, -) -> torch.Tensor: - codes = _unpack_quantized_features(packed_features, quantization_metadata) - if quantization_metadata.bits == 1: - if ( - quantization_metadata.bucket_0_value is None - or quantization_metadata.bucket_1_value is None - ): - raise ValueError("Centroid quantization requires both bucket values.") - return torch.where( - codes.bool(), - torch.tensor( - quantization_metadata.bucket_1_value, - dtype=torch.float32, - device=packed_features.device, - ), - torch.tensor( - quantization_metadata.bucket_0_value, - dtype=torch.float32, - device=packed_features.device, - ), - ) - - if quantization_metadata.clip_min is None or quantization_metadata.clip_max is None: - raise ValueError("Linear quantization requires both clip bounds.") - levels = (1 << quantization_metadata.bits) - 1 - clip_min = torch.tensor( - quantization_metadata.clip_min, - dtype=torch.float32, - device=packed_features.device, - ) - clip_max = torch.tensor( - quantization_metadata.clip_max, - dtype=torch.float32, - device=packed_features.device, - ) - return clip_min + (codes / levels) * (clip_max - clip_min) - - -def append_dequantized_node_features( +def materialize_quantized_node_features( data: _GraphType, metadata: dict[str, torch.Tensor], node_quantization_metadata: Optional[ @@ -421,13 +350,32 @@ def append_dequantized_node_features( ] ], ) -> tuple[_GraphType, dict[str, torch.Tensor]]: - """Append dequantized packed node features to PyG node feature tensors.""" + """Materialize packed quantized node features into PyG node feature tensors.""" if node_quantization_metadata is None: return data, metadata + def materialize_node_store( + node_store, + packed_features: torch.Tensor, + quantization_metadata: FeatureQuantizationMetadata, + ) -> None: + dequantized_features = dequantize_feature_tensor( + packed_features=packed_features, + quantization_metadata=quantization_metadata, + ) + node_x = getattr(node_store, "x", None) + if node_x is None: + node_store.x = dequantized_features + return + if node_x.size(0) != dequantized_features.size(0): + raise ValueError( + "Cannot materialize quantized features with " + f"{dequantized_features.size(0)} rows into existing x with " + f"{node_x.size(0)} rows." + ) + node_store.x = torch.cat([node_x, dequantized_features], dim=1) + if isinstance(data, Data): - quantization_metadata = node_quantization_metadata - metadata_key = NODE_QUANTIZED_FEATURES_METADATA_KEY if isinstance(node_quantization_metadata, dict): quantization_metadata = node_quantization_metadata.get( DEFAULT_HOMOGENEOUS_NODE_TYPE @@ -436,6 +384,9 @@ def append_dequantized_node_features( f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}." f"{DEFAULT_HOMOGENEOUS_NODE_TYPE}" ) + else: + quantization_metadata = node_quantization_metadata + metadata_key = NODE_QUANTIZED_FEATURES_METADATA_KEY if quantization_metadata is None: return data, metadata packed_features = metadata.pop(metadata_key, None) @@ -443,21 +394,7 @@ def append_dequantized_node_features( packed_features = metadata.pop(NODE_QUANTIZED_FEATURES_METADATA_KEY, None) if packed_features is None: return data, metadata - dequantized_features = _dequantize_node_features( - packed_features=packed_features, - quantization_metadata=quantization_metadata, - ) - data_x = data.x - if data_x is None: - data.x = dequantized_features - else: - if data_x.size(0) != dequantized_features.size(0): - raise ValueError( - "Cannot append dequantized features with " - f"{dequantized_features.size(0)} rows to existing x with " - f"{data_x.size(0)} rows." - ) - data.x = torch.cat([data_x.to(torch.float32), dequantized_features], dim=1) + materialize_node_store(data, packed_features, quantization_metadata) return data, metadata if isinstance(node_quantization_metadata, FeatureQuantizationMetadata): @@ -470,24 +407,7 @@ def append_dequantized_node_features( packed_features = metadata.pop(metadata_key, None) if packed_features is None: continue - dequantized_features = _dequantize_node_features( - packed_features=packed_features, - quantization_metadata=quantization_metadata, - ) - node_data = data[node_type] - node_x = getattr(node_data, "x", None) - if node_x is None: - node_data.x = dequantized_features - else: - if node_x.size(0) != dequantized_features.size(0): - raise ValueError( - "Cannot append dequantized features with " - f"{dequantized_features.size(0)} rows to existing x with " - f"{node_x.size(0)} rows." - ) - node_data.x = torch.cat( - [node_x.to(torch.float32), dequantized_features], dim=1 - ) + materialize_node_store(data[node_type], packed_features, quantization_metadata) return data, metadata From 29bf1d6e8d809228fed57daaecf2623ef9ea41e0 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 21:05:30 +0000 Subject: [PATCH 13/44] Simplify packing utils --- gigl/common/utils/feature_quantization.py | 102 ++++++++-------------- gigl/distributed/utils/neighborloader.py | 24 +++-- 2 files changed, 49 insertions(+), 77 deletions(-) diff --git a/gigl/common/utils/feature_quantization.py b/gigl/common/utils/feature_quantization.py index 73a2d888d..3043498e4 100644 --- a/gigl/common/utils/feature_quantization.py +++ b/gigl/common/utils/feature_quantization.py @@ -1,76 +1,50 @@ -"""Tensor utilities for feature quantization.""" - import torch -from gigl.types.graph import FeatureQuantizationMetadata - -def _unpack_quantized_features( +def dequantize_feature_tensor( packed_features: torch.Tensor, - quantization_metadata: FeatureQuantizationMetadata, + *, + bits: int, + dequantized_dim: int, + clip_min: float | None = None, clip_max: float | None = None, + bucket_0_value: float | None = None, bucket_1_value: float | None = None ) -> torch.Tensor: - bits = quantization_metadata.bits - dequantized_feature_dim = quantization_metadata.dequantized_feature_dim - if bits not in {1, 2, 4, 8}: - raise ValueError(f"Expected bits to be one of 1, 2, 4, or 8, got {bits}.") - - if bits == 8: - if packed_features.size(1) < dequantized_feature_dim: + codes = _unpack_codes(packed_features, bits=bits, dim=dequantized_dim) + if bits == 1: + if bucket_0_value is None or bucket_1_value is None: raise ValueError( - f"Packed feature dim {packed_features.size(1)} is smaller than " - f"dequantized dim {dequantized_feature_dim}." + "Expected bucket_0_value and bucket_1_value for 1-bit centroid " + "dequantization." ) - return packed_features[:, :dequantized_feature_dim].to(torch.float32) + return torch.where(codes.bool(), bucket_1_value, bucket_0_value) - values_per_byte = 8 // bits - feature_indices = torch.arange( - dequantized_feature_dim, device=packed_features.device - ) - byte_indices = torch.div(feature_indices, values_per_byte, rounding_mode="floor") - bit_offsets = ((feature_indices % values_per_byte) * bits).to(torch.int16) - selected_bytes = packed_features[:, byte_indices].to(torch.int16) - return torch.bitwise_and( - torch.bitwise_right_shift(selected_bytes, bit_offsets), - (1 << bits) - 1, - ).to(torch.float32) + if clip_min is None or clip_max is None: + raise ValueError( + f"Expected clip_min and clip_max for {bits}-bit linear dequantization." + ) + levels = (1 << bits) - 1 + return clip_min + (codes / levels) * (clip_max - clip_min) -def dequantize_feature_tensor( - packed_features: torch.Tensor, - quantization_metadata: FeatureQuantizationMetadata, +def _unpack_codes( + packed_features: torch.Tensor, *, bits: int, dim: int ) -> torch.Tensor: - codes = _unpack_quantized_features(packed_features, quantization_metadata) - if quantization_metadata.bits == 1: - if ( - quantization_metadata.bucket_0_value is None - or quantization_metadata.bucket_1_value is None - ): - raise ValueError("Centroid quantization requires both bucket values.") - return torch.where( - codes.bool(), - torch.tensor( - quantization_metadata.bucket_1_value, - dtype=torch.float32, - device=packed_features.device, - ), - torch.tensor( - quantization_metadata.bucket_0_value, - dtype=torch.float32, - device=packed_features.device, - ), - ) + if bits not in {1, 2, 4, 8}: + raise ValueError(f"Expected bits to be one of 1, 2, 4, or 8, got {bits}.") - if quantization_metadata.clip_min is None or quantization_metadata.clip_max is None: - raise ValueError("Linear quantization requires both clip bounds.") - levels = (1 << quantization_metadata.bits) - 1 - clip_min = torch.tensor( - quantization_metadata.clip_min, - dtype=torch.float32, - device=packed_features.device, - ) - clip_max = torch.tensor( - quantization_metadata.clip_max, - dtype=torch.float32, - device=packed_features.device, - ) - return clip_min + (codes / levels) * (clip_max - clip_min) + per_byte = 8 // bits + packed_dim = (dim + per_byte - 1) // per_byte + if packed_features.size(1) != packed_dim: + raise ValueError( + f"Expected packed feature dim {packed_dim} for {dim} {bits}-bit " + f"features, got {packed_features.size(1)}." + ) + if bits == 8: + return packed_features.to(torch.float32) + + idx = torch.arange(dim, device=packed_features.device) + byte_idx = torch.div(idx, per_byte, rounding_mode="floor") + offsets = ((idx % per_byte) * bits).to(torch.int16) + bytes_ = packed_features[:, byte_idx].to(torch.int16) + shifted = torch.bitwise_right_shift(bytes_, offsets) + return torch.bitwise_and(shifted, (1 << bits) - 1).to(torch.float32) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index eb5ef9fb0..1e923d0e9 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -355,25 +355,23 @@ def materialize_quantized_node_features( return data, metadata def materialize_node_store( - node_store, - packed_features: torch.Tensor, - quantization_metadata: FeatureQuantizationMetadata, + node_store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata ) -> None: - dequantized_features = dequantize_feature_tensor( - packed_features=packed_features, - quantization_metadata=quantization_metadata, + dequantized = dequantize_feature_tensor( + packed_features, bits=q.bits, dequantized_dim=q.dequantized_feature_dim, + clip_min=q.clip_min, clip_max=q.clip_max, + bucket_0_value=q.bucket_0_value, bucket_1_value=q.bucket_1_value ) - node_x = getattr(node_store, "x", None) - if node_x is None: - node_store.x = dequantized_features + x = getattr(node_store, "x", None) + if x is None: + node_store.x = dequantized return - if node_x.size(0) != dequantized_features.size(0): + if x.size(0) != dequantized.size(0): raise ValueError( "Cannot materialize quantized features with " - f"{dequantized_features.size(0)} rows into existing x with " - f"{node_x.size(0)} rows." + f"{dequantized.size(0)} rows into existing x with {x.size(0)} rows." ) - node_store.x = torch.cat([node_x, dequantized_features], dim=1) + node_store.x = torch.cat([x, dequantized], dim=1) if isinstance(data, Data): if isinstance(node_quantization_metadata, dict): From 34df1b4d1285bb85206facf52e7d59dc2ccd20fb Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 21:28:36 +0000 Subject: [PATCH 14/44] Clean up quant util --- gigl/common/utils/feature_quantization.py | 48 +++++++++++------------ 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/gigl/common/utils/feature_quantization.py b/gigl/common/utils/feature_quantization.py index 3043498e4..ced03c42b 100644 --- a/gigl/common/utils/feature_quantization.py +++ b/gigl/common/utils/feature_quantization.py @@ -9,42 +9,42 @@ def dequantize_feature_tensor( clip_min: float | None = None, clip_max: float | None = None, bucket_0_value: float | None = None, bucket_1_value: float | None = None ) -> torch.Tensor: - codes = _unpack_codes(packed_features, bits=bits, dim=dequantized_dim) + VALID_BITS = (1, 2, 4, 8) + if bits not in VALID_BITS: + raise ValueError(f"Expected bits to be one of {VALID_BITS}, got {bits}.") if bits == 1: if bucket_0_value is None or bucket_1_value is None: raise ValueError( - "Expected bucket_0_value and bucket_1_value for 1-bit centroid " - "dequantization." + "Expected bucket_0_value and bucket_1_value for 1-bit dequantization." ) - return torch.where(codes.bool(), bucket_1_value, bucket_0_value) - - if clip_min is None or clip_max is None: + elif clip_min is None or clip_max is None: raise ValueError( - f"Expected clip_min and clip_max for {bits}-bit linear dequantization." + f"Expected clip_min and clip_max for {bits}-bit dequantization." ) - levels = (1 << bits) - 1 - return clip_min + (codes / levels) * (clip_max - clip_min) + + codes = _unpack_features(packed_features, dim=dequantized_dim, bits=bits).float() + if bits == 1: + return torch.where(codes.bool(), bucket_1_value, bucket_0_value) + else: + levels = (1 << bits) - 1 + return clip_min + (codes / levels) * (clip_max - clip_min) -def _unpack_codes( - packed_features: torch.Tensor, *, bits: int, dim: int +def _unpack_features( + packed_features: torch.Tensor, *, dim: int, bits: int ) -> torch.Tensor: - if bits not in {1, 2, 4, 8}: - raise ValueError(f"Expected bits to be one of 1, 2, 4, or 8, got {bits}.") + """Unpack codes written high-bits-first within each byte. + Example: 2-bit codes [a, b, c, d] must be packed as byte: ``a << 6 | b << 4 | c << 2 | d``. + """ per_byte = 8 // bits packed_dim = (dim + per_byte - 1) // per_byte - if packed_features.size(1) != packed_dim: + if packed_features.size(-1) != packed_dim: raise ValueError( f"Expected packed feature dim {packed_dim} for {dim} {bits}-bit " - f"features, got {packed_features.size(1)}." + f"features, got {packed_features.size(-1)}." ) - if bits == 8: - return packed_features.to(torch.float32) - - idx = torch.arange(dim, device=packed_features.device) - byte_idx = torch.div(idx, per_byte, rounding_mode="floor") - offsets = ((idx % per_byte) * bits).to(torch.int16) - bytes_ = packed_features[:, byte_idx].to(torch.int16) - shifted = torch.bitwise_right_shift(bytes_, offsets) - return torch.bitwise_and(shifted, (1 << bits) - 1).to(torch.float32) + mask = (1 << bits) - 1 + 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) From 850abfccc8fce031d2b7ea2aab053de0f0ff7cde Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 21:39:31 +0000 Subject: [PATCH 15/44] simplify --- gigl/distributed/utils/neighborloader.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 1e923d0e9..6b5f11963 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -375,9 +375,9 @@ def materialize_node_store( if isinstance(data, Data): if isinstance(node_quantization_metadata, dict): - quantization_metadata = node_quantization_metadata.get( + quantization_metadata = node_quantization_metadata[ DEFAULT_HOMOGENEOUS_NODE_TYPE - ) + ] metadata_key = ( f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}." f"{DEFAULT_HOMOGENEOUS_NODE_TYPE}" @@ -385,21 +385,14 @@ def materialize_node_store( else: quantization_metadata = node_quantization_metadata metadata_key = NODE_QUANTIZED_FEATURES_METADATA_KEY - if quantization_metadata is None: - return data, metadata packed_features = metadata.pop(metadata_key, None) if packed_features is None: - packed_features = metadata.pop(NODE_QUANTIZED_FEATURES_METADATA_KEY, None) - if packed_features is None: - return data, metadata + raise ValueError( + f"Missing packed quantized node features in metadata key {metadata_key}." + ) materialize_node_store(data, packed_features, quantization_metadata) return data, metadata - if isinstance(node_quantization_metadata, FeatureQuantizationMetadata): - raise ValueError( - "Expected per-node-type quantization metadata for heterogeneous data." - ) - for node_type, quantization_metadata in node_quantization_metadata.items(): metadata_key = f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{node_type}" packed_features = metadata.pop(metadata_key, None) From b1f75ca5b611e2de1044bcd1323123f8cae9d6d9 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 30 Jun 2026 22:15:39 +0000 Subject: [PATCH 16/44] Type check --- gigl/common/data/dataloaders.py | 2 +- gigl/common/data/load_torch_tensors.py | 6 ++-- gigl/common/utils/feature_quantization.py | 21 ++++++------ gigl/distributed/base_sampler.py | 14 ++++---- gigl/distributed/dist_dataset.py | 12 +++---- gigl/distributed/dist_partitioner.py | 14 +++----- gigl/distributed/utils/neighborloader.py | 32 ++++++++++++------- .../pb_wrappers/preprocessed_metadata.py | 5 ++- .../run_distributed_partitioner.py | 1 + tests/unit/common/data/dataloaders_test.py | 6 ++-- 10 files changed, 57 insertions(+), 56 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 0b29febdd..2026167b2 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -374,7 +374,7 @@ 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 diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index 7dbd61608..e8ae33a4d 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -18,8 +18,8 @@ from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, - LoadedGraphTensors, FeatureQuantizationMetadata, + LoadedGraphTensors, ) from gigl.utils.share_memory import share_memory @@ -117,9 +117,7 @@ class SerializedGraphMetadata: ] = None # Optional node quantization metadata. node_quantization_metadata: Optional[ - Union[ - FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] - ] + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ] = None diff --git a/gigl/common/utils/feature_quantization.py b/gigl/common/utils/feature_quantization.py index ced03c42b..1f082ac34 100644 --- a/gigl/common/utils/feature_quantization.py +++ b/gigl/common/utils/feature_quantization.py @@ -6,26 +6,27 @@ def dequantize_feature_tensor( *, bits: int, dequantized_dim: int, - clip_min: float | None = None, clip_max: float | None = None, - bucket_0_value: float | None = None, bucket_1_value: float | None = None + clip_min: float | None = None, + clip_max: float | None = None, + bucket_0_value: float | None = None, + bucket_1_value: float | None = None, ) -> torch.Tensor: VALID_BITS = (1, 2, 4, 8) if bits not in VALID_BITS: raise ValueError(f"Expected bits to be one of {VALID_BITS}, got {bits}.") + + codes = _unpack_features(packed_features, dim=dequantized_dim, bits=bits).float() if bits == 1: if bucket_0_value is None or bucket_1_value is None: raise ValueError( - "Expected bucket_0_value and bucket_1_value for 1-bit dequantization." + "bucket_0_value and bucket_1_value required for 1-bit dequantization." ) - elif clip_min is None or clip_max is None: - raise ValueError( - f"Expected clip_min and clip_max for {bits}-bit dequantization." - ) - - codes = _unpack_features(packed_features, dim=dequantized_dim, bits=bits).float() - if bits == 1: return torch.where(codes.bool(), bucket_1_value, bucket_0_value) else: + if clip_min is None or clip_max is None: + raise ValueError( + f"clip_min and clip_max required for {bits}-bit dequantization." + ) levels = (1 << bits) - 1 return clip_min + (codes / levels) * (clip_max - clip_min) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 05638b43e..91aec64c7 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -16,8 +16,8 @@ from graphlearn_torch.utils import reverse_edge_type from gigl.distributed.sampler import ( - NODE_QUANTIZED_FEATURES_METADATA_KEY, NEGATIVE_LABEL_METADATA_KEY, + NODE_QUANTIZED_FEATURES_METADATA_KEY, POSITIVE_LABEL_METADATA_KEY, ABLPNodeSamplerInput, ) @@ -349,10 +349,8 @@ async def _collate_fn( sorted_ntype = sorted( self.dist_node_quantized_feature.feature_pb.keys() ) - quantized_nfeat_dict = ( - self.dist_node_quantized_feature.get_all2all( - output, sorted_ntype - ) + quantized_nfeat_dict = self.dist_node_quantized_feature.get_all2all( + output, sorted_ntype ) for ntype, quantized_nfeats in quantized_nfeat_dict.items(): result_map[ @@ -433,9 +431,9 @@ async def _collate_fn( if self.dist_node_quantized_feature is not None: fut = self.dist_node_quantized_feature.async_get(output.node) quantized_nfeats = await wrap_torch_future(fut) - result_map[ - f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}" - ] = quantized_nfeats + result_map[f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}"] = ( + quantized_nfeats + ) if self.dist_edge_feature is not None: eids = result_map["eids"] fut = self.dist_edge_feature.async_get(eids) diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index 25f71de6f..e44af0c17 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -23,8 +23,8 @@ from gigl.types.graph import ( FeatureInfo, FeaturePartitionData, - GraphPartitionData, FeatureQuantizationMetadata, + GraphPartitionData, PartitionOutput, ) from gigl.utils.data_splitters import ( @@ -348,9 +348,7 @@ def node_quantized_feature_info( def node_quantization_metadata( self, ) -> Optional[ - Union[ - FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] - ] + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ]: return self._node_quantization_metadata @@ -810,8 +808,8 @@ def _initialize_node_quantized_features( for node_type, features_per_node_type in node_quantized_features.items(): assert not isinstance(node_type, EdgeType) self._node_quantized_feature_info[node_type] = FeatureInfo( - dim=features_per_node_type.size(1), - dtype=features_per_node_type.dtype, + dim=features_per_node_type.size(1), # ty: ignore[unresolved-attribute] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + dtype=features_per_node_type.dtype, # ty: ignore[unresolved-attribute] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. ) logger.info( f"Initialized node quantized features for heterogeneous graph to dataset with node types: {node_quantized_features.keys()}" @@ -1347,7 +1345,7 @@ def _build_feature_store( return { entity_key: Feature( feature_tensor=features, - id2index=id2idx[entity_key], + id2index=id2idx[entity_key], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. with_gpu=False, dtype=dtype, ) diff --git a/gigl/distributed/dist_partitioner.py b/gigl/distributed/dist_partitioner.py index 3728740d3..04de8ce72 100644 --- a/gigl/distributed/dist_partitioner.py +++ b/gigl/distributed/dist_partitioner.py @@ -582,9 +582,9 @@ def register_node_quantized_features( ) self._node_quantized_feat_dim = {} for node_type in input_node_quantized_features: - self._node_quantized_feat_dim[node_type] = ( - input_node_quantized_features[node_type].shape[1] - ) + self._node_quantized_feat_dim[node_type] = input_node_quantized_features[ + node_type + ].shape[1] def register_node_labels( self, node_labels: Union[torch.Tensor, dict[NodeType, torch.Tensor]] @@ -1129,9 +1129,7 @@ def _node_feature_partition_fn(node_feature_ids, _): if has_node_features: assert node_feature_ind is not None node_feature_partition_data = FeaturePartitionData( - feats=torch.cat( - [r[node_feature_ind] for r in partitioned_results] - ), + feats=torch.cat([r[node_feature_ind] for r in partitioned_results]), ids=partitioned_ids, ) @@ -1170,9 +1168,7 @@ def _node_feature_partition_fn(node_feature_ids, _): if node_quantized_feat_dim is not None: node_quantized_feature_partition_data = FeaturePartitionData( - feats=torch.empty( - (0, node_quantized_feat_dim), dtype=torch.uint8 - ), + feats=torch.empty((0, node_quantized_feat_dim), dtype=torch.uint8), ids=torch.empty(0), ) else: diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 6b5f11963..c132dce48 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -5,7 +5,7 @@ from copy import deepcopy from dataclasses import dataclass from enum import Enum -from typing import Literal, Optional, TypeVar, Union +from typing import Literal, Optional, TypeVar, Union, cast import torch from graphlearn_torch.channel import SampleMessage @@ -55,9 +55,7 @@ class DatasetSchema: ] # Quantization metadata for append-only packed node features. node_quantization_metadata: Optional[ - Union[ - FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] - ] + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ] # Edge feature info. edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]] @@ -345,9 +343,7 @@ def materialize_quantized_node_features( data: _GraphType, metadata: dict[str, torch.Tensor], node_quantization_metadata: Optional[ - Union[ - FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata] - ] + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ], ) -> tuple[_GraphType, dict[str, torch.Tensor]]: """Materialize packed quantized node features into PyG node feature tensors.""" @@ -358,9 +354,13 @@ def materialize_node_store( node_store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata ) -> None: dequantized = dequantize_feature_tensor( - packed_features, bits=q.bits, dequantized_dim=q.dequantized_feature_dim, - clip_min=q.clip_min, clip_max=q.clip_max, - bucket_0_value=q.bucket_0_value, bucket_1_value=q.bucket_1_value + packed_features, + bits=q.bits, + dequantized_dim=q.dequantized_feature_dim, + clip_min=q.clip_min, + clip_max=q.clip_max, + bucket_0_value=q.bucket_0_value, + bucket_1_value=q.bucket_1_value, ) x = getattr(node_store, "x", None) if x is None: @@ -375,7 +375,11 @@ def materialize_node_store( if isinstance(data, Data): if isinstance(node_quantization_metadata, dict): - quantization_metadata = node_quantization_metadata[ + homogeneous_quantization_metadata = cast( + dict[NodeType, FeatureQuantizationMetadata], + node_quantization_metadata, + ) + quantization_metadata = homogeneous_quantization_metadata[ DEFAULT_HOMOGENEOUS_NODE_TYPE ] metadata_key = ( @@ -393,6 +397,12 @@ def materialize_node_store( materialize_node_store(data, packed_features, quantization_metadata) return data, metadata + assert isinstance(node_quantization_metadata, dict), ( + "Expected per-node-type quantization metadata for heterogeneous data." + ) + node_quantization_metadata = cast( + dict[NodeType, FeatureQuantizationMetadata], node_quantization_metadata + ) for node_type, quantization_metadata in node_quantization_metadata.items(): metadata_key = f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{node_type}" packed_features = metadata.pop(metadata_key, None) diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index 4448ce97d..6e13a9422 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -72,9 +72,8 @@ def __post_init__(self): dequantized_feature_keys = list( quantized_metadata.dequantized_feature_keys ) - if ( - quantized_metadata.dequantized_feature_dim - != len(dequantized_feature_keys) + if quantized_metadata.dequantized_feature_dim != len( + dequantized_feature_keys ): raise ValueError( f"Expected quantized feature metadata for condensed node type {condensed_node_type} " diff --git a/tests/test_assets/distributed/run_distributed_partitioner.py b/tests/test_assets/distributed/run_distributed_partitioner.py index 3bbc3406c..f11014d77 100644 --- a/tests/test_assets/distributed/run_distributed_partitioner.py +++ b/tests/test_assets/distributed/run_distributed_partitioner.py @@ -120,6 +120,7 @@ def run_distributed_partitioner( del node_features ( output_node_features, + output_quantized_node_features, # TODO(jchmura-sc): Update tests once API is decided on output_node_labels, ) = dist_partitioner.partition_node_features_and_labels( node_partition_book=output_node_partition_book diff --git a/tests/unit/common/data/dataloaders_test.py b/tests/unit/common/data/dataloaders_test.py index bb7bad002..d52492d0c 100644 --- a/tests/unit/common/data/dataloaders_test.py +++ b/tests/unit/common/data/dataloaders_test.py @@ -291,7 +291,7 @@ def test_load_as_torch_tensors( ): """Test TFRecordDataLoader's ability to load features and optionally labels.""" loader = TFRecordDataLoader(rank=0, world_size=1) - node_ids, feature_tensor, label_tensor = loader.load_as_torch_tensors( + node_ids, feature_tensor, quantized_feature_tensor, label_tensor = loader.load_as_torch_tensors( serialized_tf_record_info=SerializedTFRecordInfo( tfrecord_uri_prefix=UriFactory.create_uri(self.data_dir), feature_spec=feature_spec, @@ -396,7 +396,7 @@ def test_load_empty_directory( self.addCleanup(temp_dir.cleanup) loader = TFRecordDataLoader(rank=0, world_size=1) - node_ids, feature_tensor, label_tensor = loader.load_as_torch_tensors( + node_ids, feature_tensor, quantized_feature_tensor, label_tensor = loader.load_as_torch_tensors( serialized_tf_record_info=SerializedTFRecordInfo( tfrecord_uri_prefix=UriFactory.create_uri(temp_dir.name), feature_spec={}, # Doesn't matter what this is. @@ -470,7 +470,7 @@ def test_load_labels_from_pb(self): condensed_node_type ] loader = TFRecordDataLoader(rank=0, world_size=1) - _, feature_tensor, label_tensor = loader.load_as_torch_tensors( + _, feature_tensor, quantized_feature_tensor, label_tensor = loader.load_as_torch_tensors( serialized_tf_record_info=SerializedTFRecordInfo( tfrecord_uri_prefix=UriFactory.create_uri( node_metadata.tfrecord_uri_prefix From 490c7179bda8c81fdb8d3e2931d933eff31c0726 Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 1 Jul 2026 19:22:10 +0000 Subject: [PATCH 17/44] Write path --- .../utils/feature_quantization/__init__.py | 5 + .../feature_quantization/tensorflow_ops.py | 100 ++++++++++++++++++ .../torch_ops.py} | 39 ++++--- gigl/distributed/dist_dataset.py | 2 +- gigl/distributed/utils/neighborloader.py | 10 +- .../run_distributed_partitioner.py | 2 +- tests/unit/common/data/dataloaders_test.py | 76 +++++++------ 7 files changed, 172 insertions(+), 62 deletions(-) create mode 100644 gigl/common/utils/feature_quantization/__init__.py create mode 100644 gigl/common/utils/feature_quantization/tensorflow_ops.py rename gigl/common/utils/{feature_quantization.py => feature_quantization/torch_ops.py} (55%) diff --git a/gigl/common/utils/feature_quantization/__init__.py b/gigl/common/utils/feature_quantization/__init__.py new file mode 100644 index 000000000..e124405a9 --- /dev/null +++ b/gigl/common/utils/feature_quantization/__init__.py @@ -0,0 +1,5 @@ +"""Feature quantization utilities.""" + +from gigl.common.utils.feature_quantization.torch_ops import dequantize_feature_tensor + +__all__ = ["dequantize_feature_tensor"] diff --git a/gigl/common/utils/feature_quantization/tensorflow_ops.py b/gigl/common/utils/feature_quantization/tensorflow_ops.py new file mode 100644 index 000000000..5a5e2a09b --- /dev/null +++ b/gigl/common/utils/feature_quantization/tensorflow_ops.py @@ -0,0 +1,100 @@ +from collections.abc import Sequence + +import tensorflow as tf +import tensorflow_transform as tft + +from gigl.types.graph import FeatureQuantizationMetadata + + +def quantize_tft_tensor( + x: tft.common_types.ConsistentTensorType, + bits: int, + dequantized_feature_keys: Sequence[str] = (), +) -> tuple[tf.Tensor, FeatureQuantizationMetadata]: + """Quantize a dense TFT tensor and pack codes along its last dimension. + + Quantization runs in TensorFlow/TFT preprocessing. GiGL training later unpacks + and dequantizes the emitted packed tensor with Torch. By convention, 1-bit + uses two centroids and 2/4/8-bit uses linear buckets. + """ + VALID_BITS = (1, 2, 4, 8) + if bits not in VALID_BITS: + raise ValueError(f"bits must be one of {VALID_BITS}.") + if isinstance(x, tf.SparseTensor): + raise ValueError("Feature quantization expects a dense numerical tensor.") + x = tf.cast(x, tf.float32) + q = compute_feature_quantization_metadata( + x=x, bits=bits, dequantized_feature_keys=dequantized_feature_keys + ) + if bits == 1: + codes = tf.cast(x > 0, tf.uint8) + else: + levels = (1 << bits) - 1 + clipped = tf.clip_by_value(x, q.clip_min, q.clip_max) + scaled = (clipped - q.clip_min) / (q.clip_max - q.clip_min) + codes = tf.cast(tf.round(scaled * levels), tf.uint8) + return pack_tft_feature_codes(codes, bits), q + + +def compute_feature_quantization_metadata( + x: tft.common_types.ConsistentTensorType, + bits: int, + dequantized_feature_keys: Sequence[str] = (), +) -> FeatureQuantizationMetadata: + """Compute global metadata: 1-bit centroid, 2/4/8-bit linear buckets.""" + dim = x.shape[-1] + if dim is None: + raise ValueError("Feature quantization expects a known final dimension.") + dim = int(dim) + per_byte = 8 // bits + kwargs = dict( + bits=bits, + packed_feature_dim=(dim + per_byte - 1) // per_byte, + dequantized_feature_dim=dim, + dequantized_feature_keys=tuple(dequantized_feature_keys) + ) + if bits == 1: + x_flat = tf.reshape(x, [-1]) + pos = tf.reshape(x > 0, [-1]) + neg = tf.logical_not(pos) + zeros = tf.zeros_like(x_flat) + return FeatureQuantizationMetadata( + **kwargs, + bucket_0_value=tf.math.divide_no_nan( + tft.sum(tf.where(neg, x_flat, zeros)), tft.sum(tf.cast(neg, tf.float32)) + ), + bucket_1_value=tf.math.divide_no_nan( + tft.sum(tf.where(pos, x_flat, zeros)), tft.sum(tf.cast(pos, tf.float32)) + ) + ) + + x_abs = tf.reshape(tf.abs(x), [-1]) + bounds = tft.quantiles(x_abs, num_buckets=1000, epsilon=0.001) + # TFT returns boundaries for k / 1000, k=1..999; 0.995 is index 994. + fmax = tf.maximum(bounds[..., 994], 1e-5) + clip_abs = tf.maximum(fmax, 1e-5) + return FeatureQuantizationMetadata(**kwargs, clip_min=-clip_abs, clip_max=clip_abs) + + +def pack_tft_feature_codes(codes: tf.Tensor, bits: int) -> tf.Tensor: + """Pack integer codes along the last dim, high-bits-first within each byte. + + This must match Torch unpacking in ``torch_ops.py``. Example: 2-bit codes + [a, b, c, d] are packed as byte ``a << 6 | b << 4 | c << 2 | d``. + """ + codes = tf.cast(codes, tf.int32) + per_byte = 8 // bits + pad = tf.math.floormod(-tf.shape(codes)[-1], per_byte) + rank = tf.rank(codes) + leading_padding = tf.zeros(tf.stack([rank - 1, 2]), dtype=tf.int32) + last_padding = tf.expand_dims(tf.stack([0, pad]), axis=0) + codes = tf.pad(codes, tf.concat([leading_padding, last_padding], axis=0)) + + padded_shape = tf.shape(codes) + new_shape = tf.concat( + [padded_shape[:-1], tf.constant([-1, per_byte], dtype=tf.int32)], axis=0 + ) + codes = tf.reshape(codes, new_shape) + shifts = bits * tf.range(per_byte - 1, -1, -1, dtype=tf.int32) + weights = tf.bitwise.left_shift(tf.ones_like(shifts), shifts) + return tf.cast(tf.reduce_sum(codes * weights, axis=-1), tf.uint8) diff --git a/gigl/common/utils/feature_quantization.py b/gigl/common/utils/feature_quantization/torch_ops.py similarity index 55% rename from gigl/common/utils/feature_quantization.py rename to gigl/common/utils/feature_quantization/torch_ops.py index 1f082ac34..423647f2f 100644 --- a/gigl/common/utils/feature_quantization.py +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -1,34 +1,41 @@ import torch +from gigl.types.graph import FeatureQuantizationMetadata + def dequantize_feature_tensor( packed_features: torch.Tensor, *, - bits: int, - dequantized_dim: int, - clip_min: float | None = None, - clip_max: float | None = None, - bucket_0_value: float | None = None, - bucket_1_value: float | None = None, + metadata: FeatureQuantizationMetadata, ) -> torch.Tensor: VALID_BITS = (1, 2, 4, 8) - if bits not in VALID_BITS: - raise ValueError(f"Expected bits to be one of {VALID_BITS}, got {bits}.") + if metadata.bits not in VALID_BITS: + raise ValueError( + f"Expected bits to be one of {VALID_BITS}, got {metadata.bits}." + ) - codes = _unpack_features(packed_features, dim=dequantized_dim, bits=bits).float() - if bits == 1: - if bucket_0_value is None or bucket_1_value is None: + codes = _unpack_features( + packed_features, + dim=metadata.dequantized_feature_dim, + bits=metadata.bits, + ).float() + if metadata.bits == 1: + if metadata.bucket_0_value is None or metadata.bucket_1_value is None: raise ValueError( "bucket_0_value and bucket_1_value required for 1-bit dequantization." ) - return torch.where(codes.bool(), bucket_1_value, bucket_0_value) + return torch.where( + codes.bool(), metadata.bucket_1_value, metadata.bucket_0_value + ) else: - if clip_min is None or clip_max is None: + if metadata.clip_min is None or metadata.clip_max is None: raise ValueError( - f"clip_min and clip_max required for {bits}-bit dequantization." + f"clip_min and clip_max required for {metadata.bits}-bit dequantization." ) - levels = (1 << bits) - 1 - return clip_min + (codes / levels) * (clip_max - clip_min) + levels = (1 << metadata.bits) - 1 + return metadata.clip_min + (codes / levels) * ( + metadata.clip_max - metadata.clip_min + ) def _unpack_features( diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index e44af0c17..066aa3caf 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -1345,7 +1345,7 @@ def _build_feature_store( return { entity_key: Feature( feature_tensor=features, - id2index=id2idx[entity_key], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + id2index=id2idx[entity_key], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. with_gpu=False, dtype=dtype, ) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index c132dce48..9da402040 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -353,15 +353,7 @@ def materialize_quantized_node_features( def materialize_node_store( node_store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata ) -> None: - dequantized = dequantize_feature_tensor( - packed_features, - bits=q.bits, - dequantized_dim=q.dequantized_feature_dim, - clip_min=q.clip_min, - clip_max=q.clip_max, - bucket_0_value=q.bucket_0_value, - bucket_1_value=q.bucket_1_value, - ) + dequantized = dequantize_feature_tensor(packed_features, metadata=q) x = getattr(node_store, "x", None) if x is None: node_store.x = dequantized diff --git a/tests/test_assets/distributed/run_distributed_partitioner.py b/tests/test_assets/distributed/run_distributed_partitioner.py index f11014d77..07dce41d6 100644 --- a/tests/test_assets/distributed/run_distributed_partitioner.py +++ b/tests/test_assets/distributed/run_distributed_partitioner.py @@ -120,7 +120,7 @@ def run_distributed_partitioner( del node_features ( output_node_features, - output_quantized_node_features, # TODO(jchmura-sc): Update tests once API is decided on + output_quantized_node_features, # TODO(jchmura-sc): Update tests once API is decided on output_node_labels, ) = dist_partitioner.partition_node_features_and_labels( node_partition_book=output_node_partition_book diff --git a/tests/unit/common/data/dataloaders_test.py b/tests/unit/common/data/dataloaders_test.py index d52492d0c..de37db23b 100644 --- a/tests/unit/common/data/dataloaders_test.py +++ b/tests/unit/common/data/dataloaders_test.py @@ -291,17 +291,19 @@ def test_load_as_torch_tensors( ): """Test TFRecordDataLoader's ability to load features and optionally labels.""" loader = TFRecordDataLoader(rank=0, world_size=1) - node_ids, feature_tensor, quantized_feature_tensor, label_tensor = loader.load_as_torch_tensors( - serialized_tf_record_info=SerializedTFRecordInfo( - tfrecord_uri_prefix=UriFactory.create_uri(self.data_dir), - feature_spec=feature_spec, - feature_keys=feature_keys, - feature_dim=feature_dim, - entity_key="node_id", - label_keys=label_keys, - tfrecord_uri_pattern="100.tfrecord", - ), - tf_dataset_options=TFDatasetOptions(deterministic=True), + node_ids, feature_tensor, quantized_feature_tensor, label_tensor = ( + loader.load_as_torch_tensors( + serialized_tf_record_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=UriFactory.create_uri(self.data_dir), + feature_spec=feature_spec, + feature_keys=feature_keys, + feature_dim=feature_dim, + entity_key="node_id", + label_keys=label_keys, + tfrecord_uri_pattern="100.tfrecord", + ), + tf_dataset_options=TFDatasetOptions(deterministic=True), + ) ) # Verify entity IDs are loaded correctly @@ -396,16 +398,18 @@ def test_load_empty_directory( self.addCleanup(temp_dir.cleanup) loader = TFRecordDataLoader(rank=0, world_size=1) - node_ids, feature_tensor, quantized_feature_tensor, label_tensor = loader.load_as_torch_tensors( - serialized_tf_record_info=SerializedTFRecordInfo( - tfrecord_uri_prefix=UriFactory.create_uri(temp_dir.name), - feature_spec={}, # Doesn't matter what this is. - feature_keys=feature_keys, - feature_dim=feature_dim, - entity_key=entity_key, - label_keys=label_keys, - ), - tf_dataset_options=TFDatasetOptions(deterministic=True), + node_ids, feature_tensor, quantized_feature_tensor, label_tensor = ( + loader.load_as_torch_tensors( + serialized_tf_record_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=UriFactory.create_uri(temp_dir.name), + feature_spec={}, # Doesn't matter what this is. + feature_keys=feature_keys, + feature_dim=feature_dim, + entity_key=entity_key, + label_keys=label_keys, + ), + tf_dataset_options=TFDatasetOptions(deterministic=True), + ) ) assert_close(node_ids, expected_node_ids) @@ -470,21 +474,23 @@ def test_load_labels_from_pb(self): condensed_node_type ] loader = TFRecordDataLoader(rank=0, world_size=1) - _, feature_tensor, quantized_feature_tensor, label_tensor = loader.load_as_torch_tensors( - serialized_tf_record_info=SerializedTFRecordInfo( - tfrecord_uri_prefix=UriFactory.create_uri( - node_metadata.tfrecord_uri_prefix + _, feature_tensor, quantized_feature_tensor, label_tensor = ( + loader.load_as_torch_tensors( + serialized_tf_record_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=UriFactory.create_uri( + node_metadata.tfrecord_uri_prefix + ), + feature_spec=preprocessed_metadata_pb_wrapper.condensed_node_type_to_feature_schema_map[ + condensed_node_type + ].feature_spec, + feature_keys=node_metadata.feature_keys, + feature_dim=node_metadata.feature_dim, + entity_key=node_metadata.node_id_key, + label_keys=node_metadata.label_keys, + tfrecord_uri_pattern=".*\.tfrecord$", ), - feature_spec=preprocessed_metadata_pb_wrapper.condensed_node_type_to_feature_schema_map[ - condensed_node_type - ].feature_spec, - feature_keys=node_metadata.feature_keys, - feature_dim=node_metadata.feature_dim, - entity_key=node_metadata.node_id_key, - label_keys=node_metadata.label_keys, - tfrecord_uri_pattern=".*\.tfrecord$", - ), - tf_dataset_options=TFDatasetOptions(deterministic=True), + tf_dataset_options=TFDatasetOptions(deterministic=True), + ) ) # Ensure we have loaded data assert feature_tensor is not None and label_tensor is not None From f42f8851f5bfc1e13d2bb80aa64d85a9abe911b6 Mon Sep 17 00:00:00 2001 From: jchmura Date: Wed, 1 Jul 2026 20:49:30 +0000 Subject: [PATCH 18/44] WIP --- .../utils/feature_quantization/__init__.py | 4 +- .../feature_quantization/tensorflow_ops.py | 67 +++++-------------- .../utils/feature_quantization/torch_ops.py | 38 ++++------- gigl/distributed/utils/neighborloader.py | 4 +- .../serialized_graph_metadata_translator.py | 17 +++-- gigl/types/graph.py | 10 +-- 6 files changed, 49 insertions(+), 91 deletions(-) diff --git a/gigl/common/utils/feature_quantization/__init__.py b/gigl/common/utils/feature_quantization/__init__.py index e124405a9..4f1f23b0b 100644 --- a/gigl/common/utils/feature_quantization/__init__.py +++ b/gigl/common/utils/feature_quantization/__init__.py @@ -1,5 +1,5 @@ """Feature quantization utilities.""" -from gigl.common.utils.feature_quantization.torch_ops import dequantize_feature_tensor +from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor -__all__ = ["dequantize_feature_tensor"] +__all__ = ["dequantize_torch_tensor"] diff --git a/gigl/common/utils/feature_quantization/tensorflow_ops.py b/gigl/common/utils/feature_quantization/tensorflow_ops.py index 5a5e2a09b..3a3d752c8 100644 --- a/gigl/common/utils/feature_quantization/tensorflow_ops.py +++ b/gigl/common/utils/feature_quantization/tensorflow_ops.py @@ -1,5 +1,3 @@ -from collections.abc import Sequence - import tensorflow as tf import tensorflow_transform as tft @@ -9,23 +7,14 @@ def quantize_tft_tensor( x: tft.common_types.ConsistentTensorType, bits: int, - dequantized_feature_keys: Sequence[str] = (), ) -> tuple[tf.Tensor, FeatureQuantizationMetadata]: - """Quantize a dense TFT tensor and pack codes along its last dimension. - - Quantization runs in TensorFlow/TFT preprocessing. GiGL training later unpacks - and dequantizes the emitted packed tensor with Torch. By convention, 1-bit - uses two centroids and 2/4/8-bit uses linear buckets. - """ VALID_BITS = (1, 2, 4, 8) if bits not in VALID_BITS: - raise ValueError(f"bits must be one of {VALID_BITS}.") + raise ValueError(f"bits must be one of {VALID_BITS}, got {bits}") if isinstance(x, tf.SparseTensor): raise ValueError("Feature quantization expects a dense numerical tensor.") x = tf.cast(x, tf.float32) - q = compute_feature_quantization_metadata( - x=x, bits=bits, dequantized_feature_keys=dequantized_feature_keys - ) + q = _build_quantization_metadata(x, bits) if bits == 1: codes = tf.cast(x > 0, tf.uint8) else: @@ -33,55 +22,35 @@ def quantize_tft_tensor( clipped = tf.clip_by_value(x, q.clip_min, q.clip_max) scaled = (clipped - q.clip_min) / (q.clip_max - q.clip_min) codes = tf.cast(tf.round(scaled * levels), tf.uint8) - return pack_tft_feature_codes(codes, bits), q + return _pack_tft_tensor(codes, bits), q -def compute_feature_quantization_metadata( +def _build_quantization_metadata( x: tft.common_types.ConsistentTensorType, bits: int, - dequantized_feature_keys: Sequence[str] = (), ) -> FeatureQuantizationMetadata: - """Compute global metadata: 1-bit centroid, 2/4/8-bit linear buckets.""" dim = x.shape[-1] if dim is None: raise ValueError("Feature quantization expects a known final dimension.") dim = int(dim) per_byte = 8 // bits - kwargs = dict( - bits=bits, - packed_feature_dim=(dim + per_byte - 1) // per_byte, - dequantized_feature_dim=dim, - dequantized_feature_keys=tuple(dequantized_feature_keys) - ) + q = FeatureQuantizationMetadata(bits, (dim + per_byte - 1) // per_byte, dim, ()) if bits == 1: - x_flat = tf.reshape(x, [-1]) - pos = tf.reshape(x > 0, [-1]) - neg = tf.logical_not(pos) - zeros = tf.zeros_like(x_flat) - return FeatureQuantizationMetadata( - **kwargs, - bucket_0_value=tf.math.divide_no_nan( - tft.sum(tf.where(neg, x_flat, zeros)), tft.sum(tf.cast(neg, tf.float32)) - ), - bucket_1_value=tf.math.divide_no_nan( - tft.sum(tf.where(pos, x_flat, zeros)), tft.sum(tf.cast(pos, tf.float32)) - ) - ) - - x_abs = tf.reshape(tf.abs(x), [-1]) - bounds = tft.quantiles(x_abs, num_buckets=1000, epsilon=0.001) - # TFT returns boundaries for k / 1000, k=1..999; 0.995 is index 994. - fmax = tf.maximum(bounds[..., 994], 1e-5) - clip_abs = tf.maximum(fmax, 1e-5) - return FeatureQuantizationMetadata(**kwargs, clip_min=-clip_abs, clip_max=clip_abs) - + flat = tf.reshape(x, [-1]) + neg = tf.cast(flat <= 0, tf.float32) + pos = tf.cast(flat > 0, tf.float32) + q.bucket_0_value = tf.math.divide_no_nan(tft.sum(flat * neg), tft.sum(neg)) + q.bucket_1_value = tf.math.divide_no_nan(tft.sum(flat * pos), tft.sum(pos)) + else: + x_abs = tf.reshape(tf.abs(x), [-1]) + bounds = tft.quantiles(x_abs, num_buckets=1000, epsilon=0.001) + # TFT returns boundaries for k / 1000, k=1..999; 0.995 is index 994. + q.clip_max = tf.maximum(bounds[..., 994], 1e-5) + q.clip_min = -q.clip_max + return q -def pack_tft_feature_codes(codes: tf.Tensor, bits: int) -> tf.Tensor: - """Pack integer codes along the last dim, high-bits-first within each byte. - This must match Torch unpacking in ``torch_ops.py``. Example: 2-bit codes - [a, b, c, d] are packed as byte ``a << 6 | b << 4 | c << 2 | d``. - """ +def _pack_tft_tensor(codes: tf.Tensor, bits: int) -> tf.Tensor: codes = tf.cast(codes, tf.int32) per_byte = 8 // bits pad = tf.math.floormod(-tf.shape(codes)[-1], per_byte) diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py index 423647f2f..fdc23fdc3 100644 --- a/gigl/common/utils/feature_quantization/torch_ops.py +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -1,50 +1,40 @@ +import math + import torch from gigl.types.graph import FeatureQuantizationMetadata -def dequantize_feature_tensor( +def dequantize_torch_tensor( packed_features: torch.Tensor, - *, metadata: FeatureQuantizationMetadata, ) -> torch.Tensor: VALID_BITS = (1, 2, 4, 8) if metadata.bits not in VALID_BITS: - raise ValueError( - f"Expected bits to be one of {VALID_BITS}, got {metadata.bits}." - ) + raise ValueError(f"bits must be one of {VALID_BITS}, got {metadata.bits}") - codes = _unpack_features( + codes = _unpack_torch_tensor( packed_features, dim=metadata.dequantized_feature_dim, bits=metadata.bits, ).float() if metadata.bits == 1: - if metadata.bucket_0_value is None or metadata.bucket_1_value is None: - raise ValueError( - "bucket_0_value and bucket_1_value required for 1-bit dequantization." - ) + if math.isnan(metadata.bucket_0_value) or math.isnan(metadata.bucket_1_value): + raise ValueError("1-bit dequantization requires bucket values.") return torch.where( codes.bool(), metadata.bucket_1_value, metadata.bucket_0_value ) - else: - if metadata.clip_min is None or metadata.clip_max is None: - raise ValueError( - f"clip_min and clip_max required for {metadata.bits}-bit dequantization." - ) - levels = (1 << metadata.bits) - 1 - return metadata.clip_min + (codes / levels) * ( - metadata.clip_max - metadata.clip_min - ) + if math.isnan(metadata.clip_min) or math.isnan(metadata.clip_max): + raise ValueError(f"{metadata.bits}-bit dequantization requires clip values.") + levels = (1 << metadata.bits) - 1 + return metadata.clip_min + (codes / levels) * ( + metadata.clip_max - metadata.clip_min + ) -def _unpack_features( +def _unpack_torch_tensor( packed_features: torch.Tensor, *, dim: int, bits: int ) -> torch.Tensor: - """Unpack codes written high-bits-first within each byte. - - Example: 2-bit codes [a, b, c, d] must be packed as byte: ``a << 6 | b << 4 | c << 2 | d``. - """ per_byte = 8 // bits packed_dim = (dim + per_byte - 1) // per_byte if packed_features.size(-1) != packed_dim: diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 9da402040..c848c868f 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -13,7 +13,7 @@ from torch_geometric.typing import EdgeType, NodeType from gigl.common.logger import Logger -from gigl.common.utils.feature_quantization import dequantize_feature_tensor +from gigl.common.utils.feature_quantization import dequantize_torch_tensor from gigl.distributed.sampler import NODE_QUANTIZED_FEATURES_METADATA_KEY from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_NODE_TYPE, @@ -353,7 +353,7 @@ def materialize_quantized_node_features( def materialize_node_store( node_store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata ) -> None: - dequantized = dequantize_feature_tensor(packed_features, metadata=q) + dequantized = dequantize_torch_tensor(packed_features, metadata=q) x = getattr(node_store, "x", None) if x is None: node_store.x = dequantized diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index e7e445fe8..ce34d19f8 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -83,20 +83,19 @@ def _build_feature_quantization_metadata( f"Expected {expected_state} quantization state for {bits}-bit features." ) - return FeatureQuantizationMetadata( + metadata = FeatureQuantizationMetadata( bits=bits, packed_feature_dim=quantized_metadata.packed_feature_dim, dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), - clip_min=quantized_metadata.linear.clip_min if bits != 1 else None, - clip_max=quantized_metadata.linear.clip_max if bits != 1 else None, - bucket_0_value=quantized_metadata.centroid.bucket_0_value - if bits == 1 - else None, - bucket_1_value=quantized_metadata.centroid.bucket_1_value - if bits == 1 - else None, ) + if bits == 1: + metadata.bucket_0_value = quantized_metadata.centroid.bucket_0_value + metadata.bucket_1_value = quantized_metadata.centroid.bucket_1_value + else: + metadata.clip_min = quantized_metadata.linear.clip_min + metadata.clip_max = quantized_metadata.linear.clip_max + return metadata def convert_pb_to_serialized_graph_metadata( diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 7a6f513e1..6b0db1eb3 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -113,7 +113,7 @@ class FeatureInfo: dtype: torch.dtype -@dataclass(frozen=True) +@dataclass class FeatureQuantizationMetadata: """Metadata needed to unpack/dequantize append-only packed features.""" @@ -121,10 +121,10 @@ class FeatureQuantizationMetadata: packed_feature_dim: int dequantized_feature_dim: int dequantized_feature_keys: tuple[str, ...] - clip_min: Optional[float] = None - clip_max: Optional[float] = None - bucket_0_value: Optional[float] = None - bucket_1_value: Optional[float] = None + clip_min: float = float("nan") + clip_max: float = float("nan") + bucket_0_value: float = float("nan") + bucket_1_value: float = float("nan") def _get_label_edges( From b2a852df2300b2d3b202ba0b58e5bb3f2a13ef02 Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 15:20:49 +0000 Subject: [PATCH 19/44] WIP --- .../utils/feature_quantization/__init__.py | 5 +-- .../feature_quantization/tensorflow_ops.py | 37 ++++++++++++++++--- .../utils/feature_quantization/torch_ops.py | 29 +++++++-------- gigl/types/graph.py | 2 +- 4 files changed, 47 insertions(+), 26 deletions(-) diff --git a/gigl/common/utils/feature_quantization/__init__.py b/gigl/common/utils/feature_quantization/__init__.py index 4f1f23b0b..f500f7a93 100644 --- a/gigl/common/utils/feature_quantization/__init__.py +++ b/gigl/common/utils/feature_quantization/__init__.py @@ -1,5 +1,4 @@ -"""Feature quantization utilities.""" - +from gigl.common.utils.feature_quantization.tensorflow_ops import quantize_tft_tensor from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor -__all__ = ["dequantize_torch_tensor"] +__all__ = ["dequantize_torch_tensor", "quantize_tft_tensor"] diff --git a/gigl/common/utils/feature_quantization/tensorflow_ops.py b/gigl/common/utils/feature_quantization/tensorflow_ops.py index 3a3d752c8..a42c7267d 100644 --- a/gigl/common/utils/feature_quantization/tensorflow_ops.py +++ b/gigl/common/utils/feature_quantization/tensorflow_ops.py @@ -1,8 +1,11 @@ import tensorflow as tf import tensorflow_transform as tft +from gigl.common.logger import Logger from gigl.types.graph import FeatureQuantizationMetadata +logger = Logger() + def quantize_tft_tensor( x: tft.common_types.ConsistentTensorType, @@ -13,6 +16,7 @@ def quantize_tft_tensor( raise ValueError(f"bits must be one of {VALID_BITS}, got {bits}") if isinstance(x, tf.SparseTensor): raise ValueError("Feature quantization expects a dense numerical tensor.") + x = tf.cast(x, tf.float32) q = _build_quantization_metadata(x, bits) if bits == 1: @@ -34,19 +38,40 @@ def _build_quantization_metadata( raise ValueError("Feature quantization expects a known final dimension.") dim = int(dim) per_byte = 8 // bits - q = FeatureQuantizationMetadata(bits, (dim + per_byte - 1) // per_byte, dim, ()) + q = FeatureQuantizationMetadata( + bits=bits, + packed_feature_dim=(dim + per_byte - 1) // per_byte, + dequantized_feature_dim=dim, + ) if bits == 1: - flat = tf.reshape(x, [-1]) - neg = tf.cast(flat <= 0, tf.float32) - pos = tf.cast(flat > 0, tf.float32) - q.bucket_0_value = tf.math.divide_no_nan(tft.sum(flat * neg), tft.sum(neg)) - q.bucket_1_value = tf.math.divide_no_nan(tft.sum(flat * pos), tft.sum(pos)) + x = tf.reshape(x, [-1]) + neg = tf.cast(x <= 0, tf.float32) + pos = tf.cast(x > 0, tf.float32) + q.bucket_0_value = tf.math.divide_no_nan(tft.sum(x * neg), tft.sum(neg)) + q.bucket_1_value = tf.math.divide_no_nan(tft.sum(x * pos), tft.sum(pos)) + q.bucket_0_value = tf.debugging.check_numerics( + q.bucket_0_value, "non-finite bucket_0_value." + ) + q.bucket_1_value = tf.debugging.check_numerics( + q.bucket_1_value, "non-finite bucket_1_value." + ) + logger.info( + "Computed 1-bit centroid stats: " + f"bucket_0_value={q.bucket_0_value}, bucket_1_value={q.bucket_1_value}" + ) else: x_abs = tf.reshape(tf.abs(x), [-1]) bounds = tft.quantiles(x_abs, num_buckets=1000, epsilon=0.001) # TFT returns boundaries for k / 1000, k=1..999; 0.995 is index 994. q.clip_max = tf.maximum(bounds[..., 994], 1e-5) q.clip_min = -q.clip_max + q.clip_min = tf.debugging.check_numerics(q.clip_min, "non-finite clip_min.") + q.clip_max = tf.debugging.check_numerics(q.clip_max, "non-finite clip_max.") + logger.info( + f"Computed {bits}-bit linear stats: " + f"clip_min={q.clip_min}, clip_max={q.clip_max}" + ) + logger.info(f"Built feature quantization metadata: {q}") return q diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py index fdc23fdc3..461ec60d7 100644 --- a/gigl/common/utils/feature_quantization/torch_ops.py +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -9,27 +9,24 @@ def dequantize_torch_tensor( packed_features: torch.Tensor, metadata: FeatureQuantizationMetadata, ) -> torch.Tensor: + q = metadata + VALID_BITS = (1, 2, 4, 8) - if metadata.bits not in VALID_BITS: - raise ValueError(f"bits must be one of {VALID_BITS}, got {metadata.bits}") + if q.bits not in VALID_BITS: + raise ValueError(f"bits must be one of {VALID_BITS}, got {q.bits}") codes = _unpack_torch_tensor( - packed_features, - dim=metadata.dequantized_feature_dim, - bits=metadata.bits, + packed_features, dim=q.dequantized_feature_dim, bits=q.bits ).float() - if metadata.bits == 1: - if math.isnan(metadata.bucket_0_value) or math.isnan(metadata.bucket_1_value): + if q.bits == 1: + if math.isnan(q.bucket_0_value) or math.isnan(q.bucket_1_value): raise ValueError("1-bit dequantization requires bucket values.") - return torch.where( - codes.bool(), metadata.bucket_1_value, metadata.bucket_0_value - ) - if math.isnan(metadata.clip_min) or math.isnan(metadata.clip_max): - raise ValueError(f"{metadata.bits}-bit dequantization requires clip values.") - levels = (1 << metadata.bits) - 1 - return metadata.clip_min + (codes / levels) * ( - metadata.clip_max - metadata.clip_min - ) + return torch.where(codes.bool(), q.bucket_1_value, q.bucket_0_value) + else: + if math.isnan(q.clip_min) or math.isnan(q.clip_max): + raise ValueError(f"{q.bits}-bit dequantization requires clip values.") + levels = (1 << q.bits) - 1 + return q.clip_min + (codes / levels) * (q.clip_max - q.clip_min) def _unpack_torch_tensor( diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 6b0db1eb3..50e3a4959 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -120,7 +120,7 @@ class FeatureQuantizationMetadata: bits: int packed_feature_dim: int dequantized_feature_dim: int - dequantized_feature_keys: tuple[str, ...] + dequantized_feature_keys: tuple[str, ...] = () clip_min: float = float("nan") clip_max: float = float("nan") bucket_0_value: float = float("nan") From 6ce395f737955db2d564a3df7f912e019cd15141 Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 16:53:03 +0000 Subject: [PATCH 20/44] Better proto names --- .../feature_quantization/tensorflow_ops.py | 46 +++++++++---------- .../utils/feature_quantization/torch_ops.py | 24 ++++++---- .../serialized_graph_metadata_translator.py | 4 +- gigl/types/graph.py | 4 +- .../research/gbml/preprocessed_metadata.proto | 8 ++-- .../gbml/preprocessed_metadata_pb2.py | 30 ++++++------ .../gbml/preprocessed_metadata_pb2.pyi | 18 ++++---- 7 files changed, 69 insertions(+), 65 deletions(-) diff --git a/gigl/common/utils/feature_quantization/tensorflow_ops.py b/gigl/common/utils/feature_quantization/tensorflow_ops.py index a42c7267d..5608534a0 100644 --- a/gigl/common/utils/feature_quantization/tensorflow_ops.py +++ b/gigl/common/utils/feature_quantization/tensorflow_ops.py @@ -11,21 +11,27 @@ def quantize_tft_tensor( x: tft.common_types.ConsistentTensorType, bits: int, ) -> tuple[tf.Tensor, FeatureQuantizationMetadata]: + """Quantize the final tft tensor dimension into packed uint8 codes.""" VALID_BITS = (1, 2, 4, 8) if bits not in VALID_BITS: raise ValueError(f"bits must be one of {VALID_BITS}, got {bits}") if isinstance(x, tf.SparseTensor): raise ValueError("Feature quantization expects a dense numerical tensor.") + if x.shape.rank != 2: + raise ValueError(f"Feature quantization expects a 2D tensor, got {x.shape}.") x = tf.cast(x, tf.float32) q = _build_quantization_metadata(x, bits) if bits == 1: + # 1-bit quantization keeps only sign; values restore from neg/pos means. codes = tf.cast(x > 0, tf.uint8) else: + # Linearly map clipped values into integer buckets. levels = (1 << bits) - 1 clipped = tf.clip_by_value(x, q.clip_min, q.clip_max) scaled = (clipped - q.clip_min) / (q.clip_max - q.clip_min) codes = tf.cast(tf.round(scaled * levels), tf.uint8) + return _pack_tft_tensor(codes, bits), q @@ -44,31 +50,29 @@ def _build_quantization_metadata( dequantized_feature_dim=dim, ) if bits == 1: + # Store mean values for negative and positive buckets. x = tf.reshape(x, [-1]) neg = tf.cast(x <= 0, tf.float32) pos = tf.cast(x > 0, tf.float32) - q.bucket_0_value = tf.math.divide_no_nan(tft.sum(x * neg), tft.sum(neg)) - q.bucket_1_value = tf.math.divide_no_nan(tft.sum(x * pos), tft.sum(pos)) - q.bucket_0_value = tf.debugging.check_numerics( - q.bucket_0_value, "non-finite bucket_0_value." - ) - q.bucket_1_value = tf.debugging.check_numerics( - q.bucket_1_value, "non-finite bucket_1_value." - ) + q.neg_mean = tf.math.divide_no_nan(tft.sum(x * neg), tft.sum(neg)) + q.pos_mean = tf.math.divide_no_nan(tft.sum(x * pos), tft.sum(pos)) + tf.debugging.check_numerics(q.neg_mean, "non-finite neg_mean.") + tf.debugging.check_numerics(q.pos_mean, "non-finite pos_mean.") logger.info( - "Computed 1-bit centroid stats: " - f"bucket_0_value={q.bucket_0_value}, bucket_1_value={q.bucket_1_value}" + "Computed 1-bit quantization stats: " + f"neg_mean={q.neg_mean}, pos_mean={q.pos_mean}" ) else: + # Store symmetric clip bounds from the 99.5th abs-value percentile. x_abs = tf.reshape(tf.abs(x), [-1]) bounds = tft.quantiles(x_abs, num_buckets=1000, epsilon=0.001) # TFT returns boundaries for k / 1000, k=1..999; 0.995 is index 994. q.clip_max = tf.maximum(bounds[..., 994], 1e-5) q.clip_min = -q.clip_max - q.clip_min = tf.debugging.check_numerics(q.clip_min, "non-finite clip_min.") - q.clip_max = tf.debugging.check_numerics(q.clip_max, "non-finite clip_max.") + tf.debugging.check_numerics(q.clip_min, "non-finite clip_min.") + tf.debugging.check_numerics(q.clip_max, "non-finite clip_max.") logger.info( - f"Computed {bits}-bit linear stats: " + f"Computed {bits}-bit quantization stats: " f"clip_min={q.clip_min}, clip_max={q.clip_max}" ) logger.info(f"Built feature quantization metadata: {q}") @@ -76,19 +80,15 @@ def _build_quantization_metadata( def _pack_tft_tensor(codes: tf.Tensor, bits: int) -> tf.Tensor: + # Cast to int32 for integer shift/multiply/sum packing math. codes = tf.cast(codes, tf.int32) + # Pad only the feature dimension of this 2D [row, feature] tensor. per_byte = 8 // bits pad = tf.math.floormod(-tf.shape(codes)[-1], per_byte) - rank = tf.rank(codes) - leading_padding = tf.zeros(tf.stack([rank - 1, 2]), dtype=tf.int32) - last_padding = tf.expand_dims(tf.stack([0, pad]), axis=0) - codes = tf.pad(codes, tf.concat([leading_padding, last_padding], axis=0)) - - padded_shape = tf.shape(codes) - new_shape = tf.concat( - [padded_shape[:-1], tf.constant([-1, per_byte], dtype=tf.int32)], axis=0 - ) - codes = tf.reshape(codes, new_shape) + codes = tf.pad(codes, [[0, 0], [0, pad]]) + # Group the padded feature dimension into chunks that each form one byte. + codes = tf.reshape(codes, [tf.shape(codes)[0], -1, per_byte]) + # Place the first code in each chunk into the highest bits of the byte. shifts = bits * tf.range(per_byte - 1, -1, -1, dtype=tf.int32) weights = tf.bitwise.left_shift(tf.ones_like(shifts), shifts) return tf.cast(tf.reduce_sum(codes * weights, axis=-1), tf.uint8) diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py index 461ec60d7..488f6764d 100644 --- a/gigl/common/utils/feature_quantization/torch_ops.py +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -9,22 +9,31 @@ def dequantize_torch_tensor( packed_features: torch.Tensor, metadata: FeatureQuantizationMetadata, ) -> torch.Tensor: + """Reconstruct approximate float features from packed uint8 codes.""" q = metadata VALID_BITS = (1, 2, 4, 8) if q.bits not in VALID_BITS: raise ValueError(f"bits must be one of {VALID_BITS}, got {q.bits}") + per_byte = 8 // q.bits + expected_packed_dim = (q.dequantized_feature_dim + per_byte - 1) // per_byte + if packed_features.size(-1) != expected_packed_dim: + raise ValueError( + f"Expected packed feature dim {expected_packed_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 math.isnan(q.bucket_0_value) or math.isnan(q.bucket_1_value): - raise ValueError("1-bit dequantization requires bucket values.") - return torch.where(codes.bool(), q.bucket_1_value, q.bucket_0_value) + if math.isnan(q.neg_mean) or math.isnan(q.pos_mean): + raise ValueError("1-bit dequantization requires pos_mean/neg_mean") + return torch.where(codes.bool(), q.pos_mean, q.neg_mean) else: if math.isnan(q.clip_min) or math.isnan(q.clip_max): - raise ValueError(f"{q.bits}-bit dequantization requires clip values.") + 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) @@ -33,13 +42,8 @@ def _unpack_torch_tensor( packed_features: torch.Tensor, *, dim: int, bits: int ) -> torch.Tensor: per_byte = 8 // bits - packed_dim = (dim + per_byte - 1) // per_byte - if packed_features.size(-1) != packed_dim: - raise ValueError( - f"Expected packed feature dim {packed_dim} for {dim} {bits}-bit " - f"features, got {packed_features.size(-1)}." - ) 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) diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index ce34d19f8..121b29d23 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -90,8 +90,8 @@ def _build_feature_quantization_metadata( dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), ) if bits == 1: - metadata.bucket_0_value = quantized_metadata.centroid.bucket_0_value - metadata.bucket_1_value = quantized_metadata.centroid.bucket_1_value + metadata.neg_mean = quantized_metadata.centroid.neg_mean + metadata.pos_mean = quantized_metadata.centroid.pos_mean else: metadata.clip_min = quantized_metadata.linear.clip_min metadata.clip_max = quantized_metadata.linear.clip_max diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 50e3a4959..91f57545c 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -123,8 +123,8 @@ class FeatureQuantizationMetadata: dequantized_feature_keys: tuple[str, ...] = () clip_min: float = float("nan") clip_max: float = float("nan") - bucket_0_value: float = float("nan") - bucket_1_value: float = float("nan") + neg_mean: float = float("nan") + pos_mean: float = float("nan") def _get_label_edges( diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index 53092a5b5..fef5cc4d8 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -13,10 +13,10 @@ message PreprocessedMetadata{ } message CentroidQuantizationState{ - // Value produced by packed bit/code 0. - float bucket_0_value = 1; - // Value produced by packed bit/code 1. - float bucket_1_value = 2; + // Mean value for negative features, produced by packed bit/code 0. + float neg_mean = 1; + // Mean value for positive features, produced by packed bit/code 1. + float pos_mean = 2; } message FeatureQuantizationMetadata{ diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.py b/snapchat/research/gbml/preprocessed_metadata_pb2.py index fbd2dc7be..d87fae0a8 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.py +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.py @@ -14,7 +14,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xce\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1aK\n\x19\x43\x65ntroidQuantizationState\x12\x16\n\x0e\x62ucket_0_value\x18\x01 \x01(\x02\x12\x16\n\x0e\x62ucket_1_value\x18\x02 \x01(\x02\x1a\xe4\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1b\n\x13packed_feature_keys\x18\x01 \x03(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xc2\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xe4\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1b\n\x13packed_feature_keys\x18\x01 \x03(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') @@ -106,21 +106,21 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._options = None _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_options = b'8\001' _PREPROCESSEDMETADATA._serialized_start=79 - _PREPROCESSEDMETADATA._serialized_end=2205 + _PREPROCESSEDMETADATA._serialized_end=2193 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_start=395 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_end=456 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_start=458 - _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=533 - _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_start=536 - _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=892 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=895 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1289 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1292 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1515 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1518 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1913 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1916 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2059 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2062 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2205 + _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=521 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_start=524 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=880 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=883 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1277 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1280 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1503 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1506 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1901 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1904 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2047 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2050 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2193 # @@protoc_insertion_point(module_scope) diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index 8a6190fa2..2d21db24d 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -44,19 +44,19 @@ class PreprocessedMetadata(google.protobuf.message.Message): class CentroidQuantizationState(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor - BUCKET_0_VALUE_FIELD_NUMBER: builtins.int - BUCKET_1_VALUE_FIELD_NUMBER: builtins.int - bucket_0_value: builtins.float - """Value produced by packed bit/code 0.""" - bucket_1_value: builtins.float - """Value produced by packed bit/code 1.""" + NEG_MEAN_FIELD_NUMBER: builtins.int + POS_MEAN_FIELD_NUMBER: builtins.int + neg_mean: builtins.float + """Mean value for negative features, produced by packed bit/code 0.""" + pos_mean: builtins.float + """Mean value for positive features, produced by packed bit/code 1.""" def __init__( self, *, - bucket_0_value: builtins.float = ..., - bucket_1_value: builtins.float = ..., + neg_mean: builtins.float = ..., + pos_mean: builtins.float = ..., ) -> None: ... - def ClearField(self, field_name: typing_extensions.Literal["bucket_0_value", b"bucket_0_value", "bucket_1_value", b"bucket_1_value"]) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["neg_mean", b"neg_mean", "pos_mean", b"pos_mean"]) -> None: ... class FeatureQuantizationMetadata(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor From d9714b9c033d06ff21b3170c1e3f95e5b62f6cfb Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 16:57:34 +0000 Subject: [PATCH 21/44] Remove nan sentinal --- gigl/common/utils/feature_quantization/torch_ops.py | 6 ++---- gigl/types/graph.py | 8 ++++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py index 488f6764d..872d1dc0a 100644 --- a/gigl/common/utils/feature_quantization/torch_ops.py +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -1,5 +1,3 @@ -import math - import torch from gigl.types.graph import FeatureQuantizationMetadata @@ -28,11 +26,11 @@ def dequantize_torch_tensor( packed_features, dim=q.dequantized_feature_dim, bits=q.bits ).float() if q.bits == 1: - if math.isnan(q.neg_mean) or math.isnan(q.pos_mean): + 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 math.isnan(q.clip_min) or math.isnan(q.clip_max): + 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) diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 91f57545c..8e48700e7 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -121,10 +121,10 @@ class FeatureQuantizationMetadata: packed_feature_dim: int dequantized_feature_dim: int dequantized_feature_keys: tuple[str, ...] = () - clip_min: float = float("nan") - clip_max: float = float("nan") - neg_mean: float = float("nan") - pos_mean: float = float("nan") + clip_min: Optional[float] = None + clip_max: Optional[float] = None + neg_mean: Optional[float] = None + pos_mean: Optional[float] = None def _get_label_edges( From 2819bfdfd890afa63a45c44c2b21bebeba5f7f6a Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 17:46:12 +0000 Subject: [PATCH 22/44] Simplify proto --- .../serialized_graph_metadata_translator.py | 2 +- .../research/gbml/preprocessed_metadata.proto | 4 +-- .../gbml/preprocessed_metadata_pb2.py | 26 +++++++++---------- .../gbml/preprocessed_metadata_pb2.pyi | 11 ++++---- 4 files changed, 21 insertions(+), 22 deletions(-) diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index 121b29d23..755cb4e83 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -38,7 +38,7 @@ def _build_serialized_tfrecord_entity_info( if isinstance(preprocessed_metadata, PreprocessedMetadata.NodeMetadataOutput): if preprocessed_metadata.HasField("quantized_feature_metadata"): quantized_metadata = preprocessed_metadata.quantized_feature_metadata - quantized_feature_keys = list(quantized_metadata.packed_feature_keys) + quantized_feature_keys = [quantized_metadata.quantized_feature_key] quantized_feature_dim = quantized_metadata.packed_feature_dim stored_keys = set(preprocessed_metadata.feature_keys) diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index fef5cc4d8..cd0a64a89 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -20,8 +20,8 @@ message PreprocessedMetadata{ } message FeatureQuantizationMetadata{ - // Packed uint8 fields in output TFRecords. - repeated string packed_feature_keys = 1; + // Field in output TFRecords that stores packed uint8 features. + string quantized_feature_key = 1; // Model-visible feature names produced after dequantizing and appending. repeated string dequantized_feature_keys = 2; // Number of uint8 columns in the packed sidecar tensor. diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.py b/snapchat/research/gbml/preprocessed_metadata_pb2.py index d87fae0a8..98767a69b 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.py +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.py @@ -14,7 +14,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xc2\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xe4\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1b\n\x13packed_feature_keys\x18\x01 \x03(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xc4\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xe6\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1d\n\x15quantized_feature_key\x18\x01 \x01(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') @@ -106,21 +106,21 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._options = None _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_options = b'8\001' _PREPROCESSEDMETADATA._serialized_start=79 - _PREPROCESSEDMETADATA._serialized_end=2193 + _PREPROCESSEDMETADATA._serialized_end=2195 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_start=395 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_end=456 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_start=458 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=521 _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_start=524 - _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=880 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=883 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1277 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1280 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1503 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1506 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1901 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1904 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2047 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2050 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2193 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=882 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=885 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1279 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1282 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1505 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1508 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1903 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1906 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2049 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2052 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2195 # @@protoc_insertion_point(module_scope) diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index 2d21db24d..df328f07a 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -61,16 +61,15 @@ class PreprocessedMetadata(google.protobuf.message.Message): class FeatureQuantizationMetadata(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor - PACKED_FEATURE_KEYS_FIELD_NUMBER: builtins.int + QUANTIZED_FEATURE_KEY_FIELD_NUMBER: builtins.int DEQUANTIZED_FEATURE_KEYS_FIELD_NUMBER: builtins.int PACKED_FEATURE_DIM_FIELD_NUMBER: builtins.int DEQUANTIZED_FEATURE_DIM_FIELD_NUMBER: builtins.int BITS_FIELD_NUMBER: builtins.int LINEAR_FIELD_NUMBER: builtins.int CENTROID_FIELD_NUMBER: builtins.int - @property - def packed_feature_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: - """Packed uint8 fields in output TFRecords.""" + quantized_feature_key: builtins.str + """Field in output TFRecords that stores packed uint8 features.""" @property def dequantized_feature_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """Model-visible feature names produced after dequantizing and appending.""" @@ -87,7 +86,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): def __init__( self, *, - packed_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., + quantized_feature_key: builtins.str = ..., dequantized_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., packed_feature_dim: builtins.int = ..., dequantized_feature_dim: builtins.int = ..., @@ -96,7 +95,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): centroid: global___PreprocessedMetadata.CentroidQuantizationState | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["centroid", b"centroid", "linear", b"linear", "state", b"state"]) -> builtins.bool: ... - def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_dim", b"packed_feature_dim", "packed_feature_keys", b"packed_feature_keys", "state", b"state"]) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_dim", b"packed_feature_dim", "quantized_feature_key", b"quantized_feature_key", "state", b"state"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["state", b"state"]) -> typing_extensions.Literal["linear", "centroid"] | None: ... class NodeMetadataOutput(google.protobuf.message.Message): From 2029e5a88d75b31daae1c6fdd449254e5ca9a33e Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 21:15:11 +0000 Subject: [PATCH 23/44] WIP --- .../data_preprocessor/data_preprocessor.py | 34 +++++++ .../transform/transformed_features_info.py | 6 ++ .../data_preprocessor/lib/transform/utils.py | 93 +++++++++++++++++++ gigl/src/data_preprocessor/lib/types.py | 10 +- 4 files changed, 142 insertions(+), 1 deletion(-) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 543954753..e61eb5b32 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -1,5 +1,6 @@ import argparse import concurrent.futures +import json import sys import threading from collections import defaultdict @@ -286,6 +287,10 @@ def __get_feature_dimension_for_single_data_reference( ) transformed_features_info.features_outputs = preprocessing_spec.features_outputs transformed_features_info.label_outputs = preprocessing_spec.labels_outputs + if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): + transformed_features_info.feature_quantization_output = ( + preprocessing_spec.feature_quantization_output + ) if isinstance(feature_transform_pipeline_result, DataflowPipelineResult): pipeline_state: str = feature_transform_pipeline_result.state @@ -481,6 +486,35 @@ def generate_preprocessed_metadata_pb( feature_dim=feature_dim_output, transform_fn_assets_uri=node_transformed_features_info.transformed_features_transform_fn_assets_path.uri, ) + if node_transformed_features_info.feature_quantization_output is not None: + feature_quantization_output = ( + node_transformed_features_info.feature_quantization_output + ) + bits = feature_quantization_output.bits + dim = len(feature_quantization_output.dequantized_feature_keys) + per_byte = 8 // bits + quantized_feature_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( + quantized_feature_key=feature_quantization_output.quantized_feature_key, + dequantized_feature_keys=feature_quantization_output.dequantized_feature_keys, + packed_feature_dim=(dim + per_byte - 1) // per_byte, + dequantized_feature_dim=dim, + bits=bits, + ) + with tf.io.gfile.GFile( + node_transformed_features_info.feature_quantization_metadata_path.uri + ) as f: + stats = json.loads(f.read()) + if bits == 1: + centroid = quantized_feature_metadata_pb.centroid + centroid.neg_mean = stats["node_quantized_neg_mean"] + centroid.pos_mean = stats["node_quantized_pos_mean"] + else: + linear = quantized_feature_metadata_pb.linear + linear.clip_min = stats["node_quantized_clip_min"] + linear.clip_max = stats["node_quantized_clip_max"] + node_metadata_output_pb.quantized_feature_metadata.CopyFrom( + quantized_feature_metadata_pb + ) preprocessed_metadata_pb.condensed_node_type_to_preprocessed_metadata[ int(condensed_node_type) ].CopyFrom(node_metadata_output_pb) diff --git a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py index c7c0669eb..7171615ff 100644 --- a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py +++ b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py @@ -8,6 +8,7 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.data_preprocessor.lib.types import ( EdgeOutputIdentifier, + FeatureQuantizationOutput, NodeOutputIdentifier, ) @@ -22,6 +23,7 @@ class TransformedFeaturesInfo: raw_data_schema_file_path: GcsUri tft_temp_directory_path: GcsUri transformed_features_file_prefix: GcsUri + feature_quantization_metadata_path: GcsUri transformed_features_schema_path: GcsUri transform_directory_path: GcsUri dataflow_console_uri: Optional[HttpUri] = None @@ -31,6 +33,7 @@ class TransformedFeaturesInfo: features_outputs: Optional[list[str]] = None label_outputs: Optional[list[str]] = None feature_dim_output: Optional[int] = None + feature_quantization_output: Optional[FeatureQuantizationOutput] = None custom_identifier: Optional[str] = None def __init__( @@ -92,6 +95,9 @@ def __init__( custom_identifier=custom_identifier, ) ) + self.feature_quantization_metadata_path = GcsUri.join( + self.transform_directory_path, "feature_quantization_metadata.json" + ) self.transformed_features_schema_path = ( gcs_constants.get_tf_transformed_features_schema_path( diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index 7e9be7081..c6c5dee6c 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -1,3 +1,4 @@ +import json from typing import Any, Callable, Iterable, Optional, Tuple, Union import apache_beam as beam @@ -5,6 +6,7 @@ import tensorflow_data_validation as tfdv import tensorflow_data_validation.utils.display_util import tensorflow_transform +import tensorflow_transform.tf_metadata.dataset_metadata as dataset_metadata import tfx_bsl import tfx_bsl.tfxio.tensor_adapter import tfx_bsl.tfxio.tf_example_record @@ -42,6 +44,49 @@ logger = Logger() +def _extract_quantization_stats( + record_batch: pa.RecordBatch, stat_keys: list[str] +) -> Iterable[str]: + if record_batch.num_rows == 0: + return + key_to_index = {name: i for i, name in enumerate(record_batch.schema.names)} + yield json.dumps( + { + key: float(record_batch.column(key_to_index[key])[0].as_py()) + for key in stat_keys + } + ) + + +def _drop_record_batch_columns( + record_batch: pa.RecordBatch, column_names: list[str] +) -> pa.RecordBatch: + columns_to_drop = set(column_names) + keep_indices = [ + i + for i, name in enumerate(record_batch.schema.names) + if name not in columns_to_drop + ] + return pa.RecordBatch.from_arrays( + [record_batch.column(i) for i in keep_indices], + names=[record_batch.schema.names[i] for i in keep_indices], + ) + + +def _drop_dataset_metadata_features( + metadata: dataset_metadata.DatasetMetadata, feature_names: list[str] +) -> dataset_metadata.DatasetMetadata: + features_to_drop = set(feature_names) + schema = schema_pb2.Schema() + schema.CopyFrom(metadata.schema) + kept_features = [ + feature for feature in schema.feature if feature.name not in features_to_drop + ] + del schema.feature[:] + schema.feature.extend(kept_features) + return dataset_metadata.DatasetMetadata(schema) + + class InstanceDictToTFExample(beam.DoFn): """ Uses a feature spec to process a raw instance dict (read from some tabular data) as a TFExample. These @@ -362,6 +407,54 @@ def get_load_data_and_transform_pipeline_component( if should_use_existing_transform_fn else beam.pvalue.AsSingleton(analyzed_transform_fn[1].deferred_metadata) # type: ignore ) + if ( + isinstance(preprocessing_spec, NodeDataPreprocessingSpec) + and preprocessing_spec.feature_quantization_output is not None + ): + node_quantized_stat_keys = ( + ["node_quantized_neg_mean", "node_quantized_pos_mean"] + if preprocessing_spec.feature_quantization_output.bits == 1 + else ["node_quantized_clip_min", "node_quantized_clip_max"] + ) + _ = ( + transformed_features + | "Extract feature quantization metadata" + >> beam.FlatMap( + _extract_quantization_stats, + stat_keys=node_quantized_stat_keys, + ) + | "Sample feature quantization metadata" + >> beam.combiners.Sample.FixedSizeGlobally(1) + | "Select feature quantization metadata" + >> beam.Map(lambda rows: rows[0]) + | "Write feature quantization metadata" + >> beam.io.WriteToText( + transformed_features_info.feature_quantization_metadata_path.uri, + num_shards=1, + shard_name_template="", + ) + ) + transformed_features = transformed_features | ( + "Drop feature quantization metadata fields" + >> beam.Map( + _drop_record_batch_columns, column_names=node_quantized_stat_keys + ) + ) + if should_use_existing_transform_fn: + resolved_transformed_metadata = _drop_dataset_metadata_features( + transformed_metadata, feature_names=node_quantized_stat_keys + ) + else: + filtered_metadata = analyzed_transform_fn[1].deferred_metadata | ( + "Drop feature quantization schema fields" + >> beam.Map( + _drop_dataset_metadata_features, + feature_names=node_quantized_stat_keys, + ) + ) + resolved_transformed_metadata = beam.pvalue.AsSingleton( + filtered_metadata + ) transformed_features | "Write tf record files" >> BetterWriteToTFRecord( file_path_prefix=transformed_features_info.transformed_features_file_prefix.uri, diff --git a/gigl/src/data_preprocessor/lib/types.py b/gigl/src/data_preprocessor/lib/types.py index 491cc9486..8b6a3b3ac 100644 --- a/gigl/src/data_preprocessor/lib/types.py +++ b/gigl/src/data_preprocessor/lib/types.py @@ -48,6 +48,12 @@ class NodeOutputIdentifier(str): """ +class FeatureQuantizationOutput(NamedTuple): + quantized_feature_key: str + dequantized_feature_keys: list[str] + bits: int + + class EdgeOutputIdentifier(NamedTuple): """ References the TFTransform output fields / column names for src and dst node ids of an edge. @@ -72,6 +78,7 @@ class NodeDataPreprocessingSpec(NamedTuple): pretrained_tft_model_uri: Optional[Uri] = None features_outputs: Optional[list[str]] = None labels_outputs: Optional[list[str]] = None + feature_quantization_output: Optional[FeatureQuantizationOutput] = None def __repr__(self) -> str: return f"""NodeDataPreprocessingSpec( @@ -80,7 +87,8 @@ def __repr__(self) -> str: preprocessing_fn={self.preprocessing_fn}, pretrained_tft_model_uri={self.pretrained_tft_model_uri}, features_outputs={self.features_outputs}, - labels_outputs={self.labels_outputs}) + labels_outputs={self.labels_outputs}, + feature_quantization_output={self.feature_quantization_output}) """ From 0423e5c63640e5fef188bf290801a0e09028085c Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 22:01:12 +0000 Subject: [PATCH 24/44] Migrate quantization to beam --- gigl/common/data/dataloaders.py | 29 ++- .../utils/feature_quantization/__init__.py | 3 +- .../feature_quantization/tensorflow_ops.py | 94 -------- .../data_preprocessor/data_preprocessor.py | 37 ++-- .../lib/transform/feature_quantization.py | 206 ++++++++++++++++++ .../transform/transformed_features_info.py | 4 +- .../data_preprocessor/lib/transform/utils.py | 101 +++------ gigl/src/data_preprocessor/lib/types.py | 10 +- 8 files changed, 273 insertions(+), 211 deletions(-) delete mode 100644 gigl/common/utils/feature_quantization/tensorflow_ops.py create mode 100644 gigl/src/data_preprocessor/lib/transform/feature_quantization.py diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 2026167b2..9ff60338d 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -402,18 +402,11 @@ def load_as_torch_tensors( ) for feature_key in quantized_feature_keys: if feature_key not in feature_spec_dict: - feature_shape = ( - [serialized_tf_record_info.quantized_feature_dim] - if len(quantized_feature_keys) == 1 - else [] - ) logger.info( - f"Injecting quantized feature key {feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape={feature_shape}, dtype=tf.int64)`" + f"Injecting quantized feature key {feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" ) - # TODO(jchmura-sc): Serialize uint8 packed features as raw - # bytes, then decode string as uint8 to avoid int64. feature_spec_dict[feature_key] = tf.io.FixedLenFeature( - shape=feature_shape, dtype=tf.int64 + shape=[], dtype=tf.string ) else: id_concat_axis = 1 @@ -501,14 +494,18 @@ def load_as_torch_tensors( if label_tensor is not None: label_tensors.append(label_tensor) if quantized_feature_keys: - quantized_feature_tensor, _ = _concatenate_features_by_names( - batch, - quantized_feature_keys, - label_keys=[], - dtype=tf.uint8, + if len(quantized_feature_keys) != 1: + raise ValueError( + f"Expected one quantized feature key, got {quantized_feature_keys}." + ) + quantized_feature_tensor = tf.io.decode_raw( + batch[quantized_feature_keys[0]], tf.uint8 + ) + quantized_feature_tensor = tf.reshape( + quantized_feature_tensor, + [-1, serialized_tf_record_info.quantized_feature_dim], ) - if quantized_feature_tensor is not None: - quantized_feature_tensors.append(quantized_feature_tensor) + quantized_feature_tensors.append(quantized_feature_tensor) num_entities_processed += ( id_tensors[-1].shape[0] if entity_type == FeatureTypes.NODE diff --git a/gigl/common/utils/feature_quantization/__init__.py b/gigl/common/utils/feature_quantization/__init__.py index f500f7a93..010ce33e9 100644 --- a/gigl/common/utils/feature_quantization/__init__.py +++ b/gigl/common/utils/feature_quantization/__init__.py @@ -1,4 +1,3 @@ -from gigl.common.utils.feature_quantization.tensorflow_ops import quantize_tft_tensor from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor -__all__ = ["dequantize_torch_tensor", "quantize_tft_tensor"] +__all__ = ["dequantize_torch_tensor"] diff --git a/gigl/common/utils/feature_quantization/tensorflow_ops.py b/gigl/common/utils/feature_quantization/tensorflow_ops.py deleted file mode 100644 index 5608534a0..000000000 --- a/gigl/common/utils/feature_quantization/tensorflow_ops.py +++ /dev/null @@ -1,94 +0,0 @@ -import tensorflow as tf -import tensorflow_transform as tft - -from gigl.common.logger import Logger -from gigl.types.graph import FeatureQuantizationMetadata - -logger = Logger() - - -def quantize_tft_tensor( - x: tft.common_types.ConsistentTensorType, - bits: int, -) -> tuple[tf.Tensor, FeatureQuantizationMetadata]: - """Quantize the final tft tensor dimension into packed uint8 codes.""" - VALID_BITS = (1, 2, 4, 8) - if bits not in VALID_BITS: - raise ValueError(f"bits must be one of {VALID_BITS}, got {bits}") - if isinstance(x, tf.SparseTensor): - raise ValueError("Feature quantization expects a dense numerical tensor.") - if x.shape.rank != 2: - raise ValueError(f"Feature quantization expects a 2D tensor, got {x.shape}.") - - x = tf.cast(x, tf.float32) - q = _build_quantization_metadata(x, bits) - if bits == 1: - # 1-bit quantization keeps only sign; values restore from neg/pos means. - codes = tf.cast(x > 0, tf.uint8) - else: - # Linearly map clipped values into integer buckets. - levels = (1 << bits) - 1 - clipped = tf.clip_by_value(x, q.clip_min, q.clip_max) - scaled = (clipped - q.clip_min) / (q.clip_max - q.clip_min) - codes = tf.cast(tf.round(scaled * levels), tf.uint8) - - return _pack_tft_tensor(codes, bits), q - - -def _build_quantization_metadata( - x: tft.common_types.ConsistentTensorType, - bits: int, -) -> FeatureQuantizationMetadata: - dim = x.shape[-1] - if dim is None: - raise ValueError("Feature quantization expects a known final dimension.") - dim = int(dim) - per_byte = 8 // bits - q = FeatureQuantizationMetadata( - bits=bits, - packed_feature_dim=(dim + per_byte - 1) // per_byte, - dequantized_feature_dim=dim, - ) - if bits == 1: - # Store mean values for negative and positive buckets. - x = tf.reshape(x, [-1]) - neg = tf.cast(x <= 0, tf.float32) - pos = tf.cast(x > 0, tf.float32) - q.neg_mean = tf.math.divide_no_nan(tft.sum(x * neg), tft.sum(neg)) - q.pos_mean = tf.math.divide_no_nan(tft.sum(x * pos), tft.sum(pos)) - tf.debugging.check_numerics(q.neg_mean, "non-finite neg_mean.") - tf.debugging.check_numerics(q.pos_mean, "non-finite pos_mean.") - logger.info( - "Computed 1-bit quantization stats: " - f"neg_mean={q.neg_mean}, pos_mean={q.pos_mean}" - ) - else: - # Store symmetric clip bounds from the 99.5th abs-value percentile. - x_abs = tf.reshape(tf.abs(x), [-1]) - bounds = tft.quantiles(x_abs, num_buckets=1000, epsilon=0.001) - # TFT returns boundaries for k / 1000, k=1..999; 0.995 is index 994. - q.clip_max = tf.maximum(bounds[..., 994], 1e-5) - q.clip_min = -q.clip_max - tf.debugging.check_numerics(q.clip_min, "non-finite clip_min.") - tf.debugging.check_numerics(q.clip_max, "non-finite clip_max.") - logger.info( - f"Computed {bits}-bit quantization stats: " - f"clip_min={q.clip_min}, clip_max={q.clip_max}" - ) - logger.info(f"Built feature quantization metadata: {q}") - return q - - -def _pack_tft_tensor(codes: tf.Tensor, bits: int) -> tf.Tensor: - # Cast to int32 for integer shift/multiply/sum packing math. - codes = tf.cast(codes, tf.int32) - # Pad only the feature dimension of this 2D [row, feature] tensor. - per_byte = 8 // bits - pad = tf.math.floormod(-tf.shape(codes)[-1], per_byte) - codes = tf.pad(codes, [[0, 0], [0, pad]]) - # Group the padded feature dimension into chunks that each form one byte. - codes = tf.reshape(codes, [tf.shape(codes)[0], -1, per_byte]) - # Place the first code in each chunk into the highest bits of the byte. - shifts = bits * tf.range(per_byte - 1, -1, -1, dtype=tf.int32) - weights = tf.bitwise.left_shift(tf.ones_like(shifts), shifts) - return tf.cast(tf.reduce_sum(codes * weights, axis=-1), tf.uint8) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index e61eb5b32..559f4b855 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -288,8 +288,8 @@ def __get_feature_dimension_for_single_data_reference( transformed_features_info.features_outputs = preprocessing_spec.features_outputs transformed_features_info.label_outputs = preprocessing_spec.labels_outputs if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): - transformed_features_info.feature_quantization_output = ( - preprocessing_spec.feature_quantization_output + transformed_features_info.feature_quantization_spec = ( + preprocessing_spec.feature_quantization_spec ) if isinstance(feature_transform_pipeline_result, DataflowPipelineResult): @@ -486,32 +486,27 @@ def generate_preprocessed_metadata_pb( feature_dim=feature_dim_output, transform_fn_assets_uri=node_transformed_features_info.transformed_features_transform_fn_assets_path.uri, ) - if node_transformed_features_info.feature_quantization_output is not None: - feature_quantization_output = ( - node_transformed_features_info.feature_quantization_output - ) - bits = feature_quantization_output.bits - dim = len(feature_quantization_output.dequantized_feature_keys) - per_byte = 8 // bits - quantized_feature_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( - quantized_feature_key=feature_quantization_output.quantized_feature_key, - dequantized_feature_keys=feature_quantization_output.dequantized_feature_keys, - packed_feature_dim=(dim + per_byte - 1) // per_byte, - dequantized_feature_dim=dim, - bits=bits, - ) + if node_transformed_features_info.feature_quantization_spec is not None: with tf.io.gfile.GFile( node_transformed_features_info.feature_quantization_metadata_path.uri ) as f: - stats = json.loads(f.read()) + metadata = json.loads(f.read()) + bits = metadata["bits"] + quantized_feature_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( + quantized_feature_key=metadata["quantized_feature_key"], + dequantized_feature_keys=metadata["dequantized_feature_keys"], + packed_feature_dim=metadata["packed_feature_dim"], + dequantized_feature_dim=metadata["dequantized_feature_dim"], + bits=bits, + ) if bits == 1: centroid = quantized_feature_metadata_pb.centroid - centroid.neg_mean = stats["node_quantized_neg_mean"] - centroid.pos_mean = stats["node_quantized_pos_mean"] + centroid.neg_mean = metadata["neg_mean"] + centroid.pos_mean = metadata["pos_mean"] else: linear = quantized_feature_metadata_pb.linear - linear.clip_min = stats["node_quantized_clip_min"] - linear.clip_max = stats["node_quantized_clip_max"] + linear.clip_min = metadata["clip_min"] + linear.clip_max = metadata["clip_max"] node_metadata_output_pb.quantized_feature_metadata.CopyFrom( quantized_feature_metadata_pb ) diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py new file mode 100644 index 000000000..41e166670 --- /dev/null +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -0,0 +1,206 @@ +import json +from typing import Iterable + +import apache_beam as beam +import numpy as np +import pyarrow as pa +import tensorflow_transform.tf_metadata.dataset_metadata as dataset_metadata +from apache_beam.transforms.stats import ApproximateQuantiles +from tensorflow_metadata.proto.v0 import schema_pb2 + +from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec + + +def build_feature_quantization_stats( + transformed_features: beam.PCollection[pa.RecordBatch], + spec: FeatureQuantizationSpec, +) -> beam.PCollection[dict[str, float]]: + if spec.bits not in (1, 2, 4, 8): + raise ValueError(f"bits must be one of 1, 2, 4, or 8, got {spec.bits}.") + if not spec.feature_keys: + raise ValueError("Feature quantization expects at least one feature key.") + if spec.bits == 1: + return ( + transformed_features + | "Extract centroid quantization values" + >> beam.FlatMap(_iter_feature_values, feature_keys=spec.feature_keys) + | "Compute centroid quantization stats" + >> beam.CombineGlobally(_CentroidStatsFn()) + ) + return ( + transformed_features + | "Extract linear quantization abs values" + >> beam.FlatMap( + _iter_feature_values, feature_keys=spec.feature_keys, use_abs=True + ) + | "Compute linear quantization quantiles" + >> ApproximateQuantiles.Globally(num_quantiles=1000) + | "Build linear quantization stats" >> beam.Map(_linear_stats_from_quantiles) + ) + + +def quantize_record_batch( + record_batch: pa.RecordBatch, + spec: FeatureQuantizationSpec, + stats: dict[str, float], +) -> pa.RecordBatch: + features = _build_feature_matrix(record_batch, spec.feature_keys) + if spec.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 << spec.bits) - 1 + clip_min = stats["clip_min"] + clip_max = stats["clip_max"] + clipped = np.clip(features, clip_min, clip_max) + scaled = (clipped - clip_min) / (clip_max - clip_min) + codes = np.rint(scaled * levels).astype(np.uint8) + + packed = _pack_codes(codes, spec.bits) + keep_indices = [ + i + for i, name in enumerate(record_batch.schema.names) + if name not in spec.feature_keys and name != spec.quantized_feature_key + ] + arrays = [record_batch.column(i) for i in keep_indices] + names = [record_batch.schema.names[i] for i in keep_indices] + arrays.append(pa.array([row.tobytes() for row in packed], type=pa.binary())) + names.append(spec.quantized_feature_key) + return pa.RecordBatch.from_arrays(arrays, names=names) + + +def feature_quantization_metadata_json( + stats: dict[str, float], + spec: FeatureQuantizationSpec, +) -> str: + per_byte = 8 // spec.bits + dim = len(spec.feature_keys) + metadata = { + "quantized_feature_key": spec.quantized_feature_key, + "dequantized_feature_keys": list(spec.feature_keys), + "packed_feature_dim": (dim + per_byte - 1) // per_byte, + "dequantized_feature_dim": dim, + "bits": spec.bits, + } + metadata.update(stats) + return json.dumps(metadata) + + +def apply_feature_quantization_schema( + metadata: dataset_metadata.DatasetMetadata, + spec: FeatureQuantizationSpec, +) -> dataset_metadata.DatasetMetadata: + features_to_drop = set(spec.feature_keys) + schema = schema_pb2.Schema() + schema.CopyFrom(metadata.schema) + kept_features = [ + feature + for feature in schema.feature + if feature.name not in features_to_drop + and feature.name != spec.quantized_feature_key + ] + del schema.feature[:] + schema.feature.extend(kept_features) + quantized_feature = schema.feature.add() + quantized_feature.name = spec.quantized_feature_key + quantized_feature.type = schema_pb2.BYTES + quantized_feature.value_count.min = 1 + quantized_feature.value_count.max = 1 + return dataset_metadata.DatasetMetadata(schema) + + +def _iter_feature_values( + record_batch: pa.RecordBatch, + feature_keys: list[str], + use_abs: bool = False, +) -> Iterable[float]: + values = _build_feature_matrix(record_batch, feature_keys).reshape(-1) + if use_abs: + values = np.abs(values) + for value in values: + if np.isfinite(value): + yield float(value) + + +def _build_feature_matrix( + record_batch: pa.RecordBatch, feature_keys: list[str] +) -> np.ndarray: + key_to_index = {name: i for i, name in enumerate(record_batch.schema.names)} + columns: list[np.ndarray] = [] + for feature_key in feature_keys: + if feature_key not in key_to_index: + raise ValueError(f"Feature key {feature_key} not found in RecordBatch.") + values = np.asarray( + record_batch.column(key_to_index[feature_key]).to_numpy( + zero_copy_only=False + ), + dtype=np.float32, + ) + if values.ndim != 1: + raise ValueError( + f"Feature quantization currently expects scalar features; " + f"got {feature_key} with shape {values.shape}." + ) + columns.append(values) + if not columns: + return np.empty((record_batch.num_rows, 0), dtype=np.float32) + return np.stack(columns, axis=1) + + +def _pack_codes(codes: np.ndarray, bits: int) -> np.ndarray: + # Cast to uint16 for integer shift/multiply/sum packing math. + 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) + # Place the first code in each chunk into the highest bits of the byte. + 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) + + +def _linear_stats_from_quantiles(quantiles: list[float]) -> dict[str, float]: + if not quantiles: + raise ValueError("Cannot compute quantization stats from no values.") + # Store symmetric clip bounds from the approximate 99.5th abs-value percentile. + index = round(0.995 * (len(quantiles) - 1)) + clip_max = max(float(quantiles[index]), 1e-5) + return {"clip_min": -clip_max, "clip_max": clip_max} + + +class _CentroidStatsFn(beam.CombineFn): + def create_accumulator(self) -> tuple[float, int, float, int]: + return 0.0, 0, 0.0, 0 + + def add_input( + self, accumulator: tuple[float, int, float, int], value: float + ) -> tuple[float, int, float, int]: + neg_sum, neg_count, pos_sum, pos_count = accumulator + if value <= 0: + return neg_sum + value, neg_count + 1, pos_sum, pos_count + return neg_sum, neg_count, pos_sum + value, pos_count + 1 + + def merge_accumulators( + self, accumulators: Iterable[tuple[float, int, float, int]] + ) -> tuple[float, int, float, int]: + neg_sum = neg_count = pos_sum = pos_count = 0 + for acc_neg_sum, acc_neg_count, acc_pos_sum, acc_pos_count in accumulators: + neg_sum += acc_neg_sum + neg_count += acc_neg_count + pos_sum += acc_pos_sum + pos_count += acc_pos_count + return neg_sum, neg_count, pos_sum, pos_count + + def extract_output( + self, accumulator: tuple[float, int, float, int] + ) -> dict[str, float]: + neg_sum, neg_count, pos_sum, pos_count = accumulator + # Store mean values for negative and positive buckets. + return { + "neg_mean": neg_sum / neg_count if neg_count else 0.0, + "pos_mean": pos_sum / pos_count if pos_count else 0.0, + } diff --git a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py index 7171615ff..38044c07b 100644 --- a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py +++ b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py @@ -8,7 +8,7 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.data_preprocessor.lib.types import ( EdgeOutputIdentifier, - FeatureQuantizationOutput, + FeatureQuantizationSpec, NodeOutputIdentifier, ) @@ -33,7 +33,7 @@ class TransformedFeaturesInfo: features_outputs: Optional[list[str]] = None label_outputs: Optional[list[str]] = None feature_dim_output: Optional[int] = None - feature_quantization_output: Optional[FeatureQuantizationOutput] = None + feature_quantization_spec: Optional[FeatureQuantizationSpec] = None custom_identifier: Optional[str] = None def __init__( diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index c6c5dee6c..7fa681a4d 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -1,4 +1,3 @@ -import json from typing import Any, Callable, Iterable, Optional, Tuple, Union import apache_beam as beam @@ -6,7 +5,6 @@ import tensorflow_data_validation as tfdv import tensorflow_data_validation.utils.display_util import tensorflow_transform -import tensorflow_transform.tf_metadata.dataset_metadata as dataset_metadata import tfx_bsl import tfx_bsl.tfxio.tensor_adapter import tfx_bsl.tfxio.tf_example_record @@ -29,6 +27,12 @@ EdgeDataReference, NodeDataReference, ) +from gigl.src.data_preprocessor.lib.transform.feature_quantization import ( + apply_feature_quantization_schema, + build_feature_quantization_stats, + feature_quantization_metadata_json, + quantize_record_batch, +) from gigl.src.data_preprocessor.lib.transform.tf_value_encoder import TFValueEncoder from gigl.src.data_preprocessor.lib.transform.transformed_features_info import ( TransformedFeaturesInfo, @@ -44,49 +48,6 @@ logger = Logger() -def _extract_quantization_stats( - record_batch: pa.RecordBatch, stat_keys: list[str] -) -> Iterable[str]: - if record_batch.num_rows == 0: - return - key_to_index = {name: i for i, name in enumerate(record_batch.schema.names)} - yield json.dumps( - { - key: float(record_batch.column(key_to_index[key])[0].as_py()) - for key in stat_keys - } - ) - - -def _drop_record_batch_columns( - record_batch: pa.RecordBatch, column_names: list[str] -) -> pa.RecordBatch: - columns_to_drop = set(column_names) - keep_indices = [ - i - for i, name in enumerate(record_batch.schema.names) - if name not in columns_to_drop - ] - return pa.RecordBatch.from_arrays( - [record_batch.column(i) for i in keep_indices], - names=[record_batch.schema.names[i] for i in keep_indices], - ) - - -def _drop_dataset_metadata_features( - metadata: dataset_metadata.DatasetMetadata, feature_names: list[str] -) -> dataset_metadata.DatasetMetadata: - features_to_drop = set(feature_names) - schema = schema_pb2.Schema() - schema.CopyFrom(metadata.schema) - kept_features = [ - feature for feature in schema.feature if feature.name not in features_to_drop - ] - del schema.feature[:] - schema.feature.extend(kept_features) - return dataset_metadata.DatasetMetadata(schema) - - class InstanceDictToTFExample(beam.DoFn): """ Uses a feature spec to process a raw instance dict (read from some tabular data) as a TFExample. These @@ -407,26 +368,22 @@ def get_load_data_and_transform_pipeline_component( if should_use_existing_transform_fn else beam.pvalue.AsSingleton(analyzed_transform_fn[1].deferred_metadata) # type: ignore ) - if ( - isinstance(preprocessing_spec, NodeDataPreprocessingSpec) - and preprocessing_spec.feature_quantization_output is not None - ): - node_quantized_stat_keys = ( - ["node_quantized_neg_mean", "node_quantized_pos_mean"] - if preprocessing_spec.feature_quantization_output.bits == 1 - else ["node_quantized_clip_min", "node_quantized_clip_max"] + if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): + feature_quantization_spec = preprocessing_spec.feature_quantization_spec + else: + feature_quantization_spec = None + if feature_quantization_spec is not None: + quantization_stats = build_feature_quantization_stats( + transformed_features=transformed_features, + spec=feature_quantization_spec, ) _ = ( - transformed_features - | "Extract feature quantization metadata" - >> beam.FlatMap( - _extract_quantization_stats, - stat_keys=node_quantized_stat_keys, + quantization_stats + | "Build feature quantization metadata JSON" + >> beam.Map( + feature_quantization_metadata_json, + spec=feature_quantization_spec, ) - | "Sample feature quantization metadata" - >> beam.combiners.Sample.FixedSizeGlobally(1) - | "Select feature quantization metadata" - >> beam.Map(lambda rows: rows[0]) | "Write feature quantization metadata" >> beam.io.WriteToText( transformed_features_info.feature_quantization_metadata_path.uri, @@ -435,25 +392,27 @@ def get_load_data_and_transform_pipeline_component( ) ) transformed_features = transformed_features | ( - "Drop feature quantization metadata fields" + "Quantize transformed feature RecordBatches" >> beam.Map( - _drop_record_batch_columns, column_names=node_quantized_stat_keys + quantize_record_batch, + spec=feature_quantization_spec, + stats=beam.pvalue.AsSingleton(quantization_stats), ) ) if should_use_existing_transform_fn: - resolved_transformed_metadata = _drop_dataset_metadata_features( - transformed_metadata, feature_names=node_quantized_stat_keys + resolved_transformed_metadata = apply_feature_quantization_schema( + transformed_metadata, spec=feature_quantization_spec ) else: - filtered_metadata = analyzed_transform_fn[1].deferred_metadata | ( - "Drop feature quantization schema fields" + quantized_metadata = analyzed_transform_fn[1].deferred_metadata | ( + "Apply feature quantization schema" >> beam.Map( - _drop_dataset_metadata_features, - feature_names=node_quantized_stat_keys, + apply_feature_quantization_schema, + spec=feature_quantization_spec, ) ) resolved_transformed_metadata = beam.pvalue.AsSingleton( - filtered_metadata + quantized_metadata ) transformed_features | "Write tf record files" >> BetterWriteToTFRecord( diff --git a/gigl/src/data_preprocessor/lib/types.py b/gigl/src/data_preprocessor/lib/types.py index 8b6a3b3ac..44ee7989a 100644 --- a/gigl/src/data_preprocessor/lib/types.py +++ b/gigl/src/data_preprocessor/lib/types.py @@ -48,10 +48,10 @@ class NodeOutputIdentifier(str): """ -class FeatureQuantizationOutput(NamedTuple): - quantized_feature_key: str - dequantized_feature_keys: list[str] +class FeatureQuantizationSpec(NamedTuple): + feature_keys: list[str] bits: int + quantized_feature_key: str class EdgeOutputIdentifier(NamedTuple): @@ -78,7 +78,7 @@ class NodeDataPreprocessingSpec(NamedTuple): pretrained_tft_model_uri: Optional[Uri] = None features_outputs: Optional[list[str]] = None labels_outputs: Optional[list[str]] = None - feature_quantization_output: Optional[FeatureQuantizationOutput] = None + feature_quantization_spec: Optional[FeatureQuantizationSpec] = None def __repr__(self) -> str: return f"""NodeDataPreprocessingSpec( @@ -88,7 +88,7 @@ def __repr__(self) -> str: pretrained_tft_model_uri={self.pretrained_tft_model_uri}, features_outputs={self.features_outputs}, labels_outputs={self.labels_outputs}, - feature_quantization_output={self.feature_quantization_output}) + feature_quantization_spec={self.feature_quantization_spec}) """ From c7ed95076b0027e5e1f23b4f8b93107dc8b43809 Mon Sep 17 00:00:00 2001 From: jchmura Date: Thu, 2 Jul 2026 22:39:05 +0000 Subject: [PATCH 25/44] Try to simplify --- .../lib/transform/feature_quantization.py | 178 +++++++++--------- .../data_preprocessor/lib/transform/utils.py | 27 ++- gigl/src/data_preprocessor/lib/types.py | 1 - 3 files changed, 102 insertions(+), 104 deletions(-) diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index 41e166670..affda5ffd 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -4,69 +4,72 @@ import apache_beam as beam import numpy as np import pyarrow as pa -import tensorflow_transform.tf_metadata.dataset_metadata as dataset_metadata from apache_beam.transforms.stats import ApproximateQuantiles from tensorflow_metadata.proto.v0 import schema_pb2 +from tensorflow_transform.tf_metadata.dataset_metadata import DatasetMetadata +from gigl.common.logger import Logger from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec +logger = Logger() +_NODE_QUANTIZED_FEATURE_KEY = "node_quantized_features" +_CentroidAcc = tuple[float, int, float, int] + def build_feature_quantization_stats( - transformed_features: beam.PCollection[pa.RecordBatch], + record_batches: beam.PCollection[pa.RecordBatch], spec: FeatureQuantizationSpec, ) -> beam.PCollection[dict[str, float]]: if spec.bits not in (1, 2, 4, 8): raise ValueError(f"bits must be one of 1, 2, 4, or 8, got {spec.bits}.") if not spec.feature_keys: raise ValueError("Feature quantization expects at least one feature key.") + logger.info( + f"Building Beam feature quantization stats for {len(spec.feature_keys)} " + f"features with bits={spec.bits}: {spec.feature_keys}" + ) if spec.bits == 1: return ( - transformed_features - | "Extract centroid quantization values" - >> beam.FlatMap(_iter_feature_values, feature_keys=spec.feature_keys) + record_batches | "Compute centroid quantization stats" - >> beam.CombineGlobally(_CentroidStatsFn()) + >> beam.CombineGlobally(_CentroidStatsFn(spec.feature_keys)) ) return ( - transformed_features - | "Extract linear quantization abs values" - >> beam.FlatMap( - _iter_feature_values, feature_keys=spec.feature_keys, use_abs=True - ) + record_batches + | "Build linear quantization value batches" + >> beam.Map(_build_abs_feature_values, feature_keys=spec.feature_keys) | "Compute linear quantization quantiles" - >> ApproximateQuantiles.Globally(num_quantiles=1000) + >> ApproximateQuantiles.Globally(num_quantiles=1000, input_batched=True) | "Build linear quantization stats" >> beam.Map(_linear_stats_from_quantiles) ) def quantize_record_batch( - record_batch: pa.RecordBatch, + batch: pa.RecordBatch, spec: FeatureQuantizationSpec, stats: dict[str, float], ) -> pa.RecordBatch: - features = _build_feature_matrix(record_batch, spec.feature_keys) + features = _build_feature_matrix(batch, spec.feature_keys) if spec.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 << spec.bits) - 1 - clip_min = stats["clip_min"] - clip_max = stats["clip_max"] - clipped = np.clip(features, clip_min, clip_max) - scaled = (clipped - clip_min) / (clip_max - clip_min) + 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) packed = _pack_codes(codes, spec.bits) + drop_keys = set(spec.feature_keys) | {_NODE_QUANTIZED_FEATURE_KEY} keep_indices = [ - i - for i, name in enumerate(record_batch.schema.names) - if name not in spec.feature_keys and name != spec.quantized_feature_key + i for i, name in enumerate(batch.schema.names) if name not in drop_keys ] - arrays = [record_batch.column(i) for i in keep_indices] - names = [record_batch.schema.names[i] for i in keep_indices] + arrays = [batch.column(i) for i in keep_indices] + names = [batch.schema.names[i] for i in keep_indices] arrays.append(pa.array([row.tobytes() for row in packed], type=pa.binary())) - names.append(spec.quantized_feature_key) + names.append(_NODE_QUANTIZED_FEATURE_KEY) return pa.RecordBatch.from_arrays(arrays, names=names) @@ -77,75 +80,66 @@ def feature_quantization_metadata_json( per_byte = 8 // spec.bits dim = len(spec.feature_keys) metadata = { - "quantized_feature_key": spec.quantized_feature_key, + "quantized_feature_key": _NODE_QUANTIZED_FEATURE_KEY, "dequantized_feature_keys": list(spec.feature_keys), "packed_feature_dim": (dim + per_byte - 1) // per_byte, "dequantized_feature_dim": dim, "bits": spec.bits, } metadata.update(stats) + logger.info(f"Writing feature quantization metadata: {metadata}") return json.dumps(metadata) def apply_feature_quantization_schema( - metadata: dataset_metadata.DatasetMetadata, + metadata: DatasetMetadata, spec: FeatureQuantizationSpec, -) -> dataset_metadata.DatasetMetadata: - features_to_drop = set(spec.feature_keys) +) -> DatasetMetadata: + drop_keys = set(spec.feature_keys) | {_NODE_QUANTIZED_FEATURE_KEY} schema = schema_pb2.Schema() schema.CopyFrom(metadata.schema) kept_features = [ - feature - for feature in schema.feature - if feature.name not in features_to_drop - and feature.name != spec.quantized_feature_key + feature for feature in schema.feature if feature.name not in drop_keys ] del schema.feature[:] schema.feature.extend(kept_features) quantized_feature = schema.feature.add() - quantized_feature.name = spec.quantized_feature_key + quantized_feature.name = _NODE_QUANTIZED_FEATURE_KEY quantized_feature.type = schema_pb2.BYTES quantized_feature.value_count.min = 1 quantized_feature.value_count.max = 1 - return dataset_metadata.DatasetMetadata(schema) - - -def _iter_feature_values( - record_batch: pa.RecordBatch, - feature_keys: list[str], - use_abs: bool = False, -) -> Iterable[float]: - values = _build_feature_matrix(record_batch, feature_keys).reshape(-1) - if use_abs: - values = np.abs(values) - for value in values: - if np.isfinite(value): - yield float(value) - - -def _build_feature_matrix( - record_batch: pa.RecordBatch, feature_keys: list[str] -) -> np.ndarray: - key_to_index = {name: i for i, name in enumerate(record_batch.schema.names)} - columns: list[np.ndarray] = [] - for feature_key in feature_keys: - if feature_key not in key_to_index: - raise ValueError(f"Feature key {feature_key} not found in RecordBatch.") - values = np.asarray( - record_batch.column(key_to_index[feature_key]).to_numpy( - zero_copy_only=False - ), - dtype=np.float32, - ) + logger.info( + f"Updated transformed schema for feature quantization: dropped " + f"{len(spec.feature_keys)} features and added bytes feature " + f"{_NODE_QUANTIZED_FEATURE_KEY}." + ) + return DatasetMetadata(schema) + + +def _build_abs_feature_values( + batch: pa.RecordBatch, feature_keys: list[str] +) -> list[float]: + values = np.abs(_build_feature_matrix(batch, feature_keys).reshape(-1)) + return values[np.isfinite(values)].astype(float).tolist() + + +def _build_feature_matrix(batch: pa.RecordBatch, feature_keys: list[str]) -> np.ndarray: + key_to_idx = {name: i for i, name in enumerate(batch.schema.names)} + cols: list[np.ndarray] = [] + for key in feature_keys: + if key not in key_to_idx: + raise ValueError(f"Feature key {key} not found in RecordBatch.") + col = batch.column(key_to_idx[key]) + values = np.asarray(col.to_numpy(zero_copy_only=False), dtype=np.float32) if values.ndim != 1: raise ValueError( f"Feature quantization currently expects scalar features; " - f"got {feature_key} with shape {values.shape}." + f"got {key} with shape {values.shape}." ) - columns.append(values) - if not columns: - return np.empty((record_batch.num_rows, 0), dtype=np.float32) - return np.stack(columns, axis=1) + cols.append(values) + if not cols: + return np.empty((batch.num_rows, 0), dtype=np.float32) + return np.stack(cols, axis=1) def _pack_codes(codes: np.ndarray, bits: int) -> np.ndarray: @@ -169,38 +163,48 @@ def _linear_stats_from_quantiles(quantiles: list[float]) -> dict[str, float]: # Store symmetric clip bounds from the approximate 99.5th abs-value percentile. index = round(0.995 * (len(quantiles) - 1)) clip_max = max(float(quantiles[index]), 1e-5) - return {"clip_min": -clip_max, "clip_max": clip_max} + stats = {"clip_min": -clip_max, "clip_max": clip_max} + logger.info(f"Computed Beam feature quantization stats: {stats}") + return stats class _CentroidStatsFn(beam.CombineFn): - def create_accumulator(self) -> tuple[float, int, float, int]: + def __init__(self, feature_keys: list[str]): + self._feature_keys = feature_keys + + def create_accumulator(self) -> _CentroidAcc: return 0.0, 0, 0.0, 0 def add_input( - self, accumulator: tuple[float, int, float, int], value: float - ) -> tuple[float, int, float, int]: + self, accumulator: _CentroidAcc, batch: pa.RecordBatch + ) -> _CentroidAcc: neg_sum, neg_count, pos_sum, pos_count = accumulator - if value <= 0: - return neg_sum + value, neg_count + 1, pos_sum, pos_count - return neg_sum, neg_count, pos_sum + value, pos_count + 1 + values = _build_feature_matrix(batch, self._feature_keys).reshape(-1) + values = values[np.isfinite(values)] + neg = values <= 0 + pos = values > 0 + return ( + neg_sum + float(values[neg].sum()), + neg_count + int(neg.sum()), + pos_sum + float(values[pos].sum()), + pos_count + int(pos.sum()), + ) - def merge_accumulators( - self, accumulators: Iterable[tuple[float, int, float, int]] - ) -> tuple[float, int, float, int]: + def merge_accumulators(self, accumulators: Iterable[_CentroidAcc]) -> _CentroidAcc: neg_sum = neg_count = pos_sum = pos_count = 0 - for acc_neg_sum, acc_neg_count, acc_pos_sum, acc_pos_count in accumulators: - neg_sum += acc_neg_sum - neg_count += acc_neg_count - pos_sum += acc_pos_sum - pos_count += acc_pos_count + for n_sum, n_count, p_sum, p_count in accumulators: + neg_sum += n_sum + neg_count += n_count + pos_sum += p_sum + pos_count += p_count return neg_sum, neg_count, pos_sum, pos_count - def extract_output( - self, accumulator: tuple[float, int, float, int] - ) -> dict[str, float]: + def extract_output(self, accumulator: _CentroidAcc) -> dict[str, float]: neg_sum, neg_count, pos_sum, pos_count = accumulator # Store mean values for negative and positive buckets. - return { + stats = { "neg_mean": neg_sum / neg_count if neg_count else 0.0, "pos_mean": pos_sum / pos_count if pos_count else 0.0, } + logger.info(f"Computed Beam feature quantization stats: {stats}") + return stats diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index 7fa681a4d..b93b3e9bd 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -369,21 +369,16 @@ def get_load_data_and_transform_pipeline_component( else beam.pvalue.AsSingleton(analyzed_transform_fn[1].deferred_metadata) # type: ignore ) if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): - feature_quantization_spec = preprocessing_spec.feature_quantization_spec + q_spec = preprocessing_spec.feature_quantization_spec else: - feature_quantization_spec = None - if feature_quantization_spec is not None: - quantization_stats = build_feature_quantization_stats( - transformed_features=transformed_features, - spec=feature_quantization_spec, - ) + q_spec = None + if q_spec is not None: + logger.info(f"Applying Beam feature quantization with spec: {q_spec}") + stats = build_feature_quantization_stats(transformed_features, q_spec) _ = ( - quantization_stats + stats | "Build feature quantization metadata JSON" - >> beam.Map( - feature_quantization_metadata_json, - spec=feature_quantization_spec, - ) + >> beam.Map(feature_quantization_metadata_json, spec=q_spec) | "Write feature quantization metadata" >> beam.io.WriteToText( transformed_features_info.feature_quantization_metadata_path.uri, @@ -395,20 +390,20 @@ def get_load_data_and_transform_pipeline_component( "Quantize transformed feature RecordBatches" >> beam.Map( quantize_record_batch, - spec=feature_quantization_spec, - stats=beam.pvalue.AsSingleton(quantization_stats), + spec=q_spec, + stats=beam.pvalue.AsSingleton(stats), ) ) if should_use_existing_transform_fn: resolved_transformed_metadata = apply_feature_quantization_schema( - transformed_metadata, spec=feature_quantization_spec + transformed_metadata, spec=q_spec ) else: quantized_metadata = analyzed_transform_fn[1].deferred_metadata | ( "Apply feature quantization schema" >> beam.Map( apply_feature_quantization_schema, - spec=feature_quantization_spec, + spec=q_spec, ) ) resolved_transformed_metadata = beam.pvalue.AsSingleton( diff --git a/gigl/src/data_preprocessor/lib/types.py b/gigl/src/data_preprocessor/lib/types.py index 44ee7989a..a3308b00d 100644 --- a/gigl/src/data_preprocessor/lib/types.py +++ b/gigl/src/data_preprocessor/lib/types.py @@ -51,7 +51,6 @@ class NodeOutputIdentifier(str): class FeatureQuantizationSpec(NamedTuple): feature_keys: list[str] bits: int - quantized_feature_key: str class EdgeOutputIdentifier(NamedTuple): From f25d8889da6988c439bb93026070e2e00ae29e51 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 15:57:51 +0000 Subject: [PATCH 26/44] Fix quantization spec copy --- gigl/common/data/load_torch_tensors.py | 10 +++++++++- gigl/src/data_preprocessor/data_preprocessor.py | 6 ++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index e8ae33a4d..e035076f7 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -191,6 +191,12 @@ def _data_loading_process( graph_type, serialized_entity_tf_record_info, ) in serialized_tf_record_info.items(): + logger.info( + f"Rank {rank} loading {entity_type} graph type {graph_type} with " + f"{len(serialized_entity_tf_record_info.feature_keys)} feature keys, " + f"quantized_feature_keys={serialized_entity_tf_record_info.quantized_feature_keys}, " + f"quantized_feature_dim={serialized_entity_tf_record_info.quantized_feature_dim}" + ) # We currently do not support training with labels for edge entities if ( serialized_entity_tf_record_info.label_keys @@ -235,7 +241,9 @@ def _data_loading_process( ) 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}" + 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}. " + f"Serialized quantized_feature_keys={serialized_entity_tf_record_info.quantized_feature_keys}, " + f"quantized_feature_dim={serialized_entity_tf_record_info.quantized_feature_dim}" ) if entity_labels is not None: diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 559f4b855..65f1880e5 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -727,7 +727,13 @@ def inner() -> FeatureSpecDict: pretrained_tft_model_uri=input_node_preprocessing_spec.pretrained_tft_model_uri, features_outputs=input_node_preprocessing_spec.features_outputs, labels_outputs=input_node_preprocessing_spec.labels_outputs, + feature_quantization_spec=input_node_preprocessing_spec.feature_quantization_spec, ) + if enumerated_node_data_preprocessing_spec.feature_quantization_spec: + logger.info( + f"Carrying node feature quantization spec through enumeration: " + f"{enumerated_node_data_preprocessing_spec.feature_quantization_spec}" + ) enumerated_node_refs_to_preprocessing_specs[ enumerated_node_metadata.enumerated_node_data_reference ] = enumerated_node_data_preprocessing_spec From 0c54a4aaa073481719913b9aef4ad1f0b48496d7 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 17:07:55 +0000 Subject: [PATCH 27/44] TFrecord io fix serialize struct array --- gigl/common/data/dataloaders.py | 6 ++---- .../data_preprocessor/lib/transform/feature_quantization.py | 4 +++- .../test_assets/distributed/run_distributed_partitioner.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 9ff60338d..7e99cf749 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -139,7 +139,6 @@ def _concatenate_features_by_names( feature_key_to_tf_tensor: dict[str, tf.Tensor], feature_keys: Sequence[str], label_keys: Sequence[str], - dtype: tf.dtypes.DType = tf.float32, ) -> tuple[Optional[tf.Tensor], Optional[tf.Tensor]]: """ Concatenates feature tensors in the order specified by feature names. @@ -151,7 +150,6 @@ def _concatenate_features_by_names( feature_key_to_tf_tensor (dict[str, tf.Tensor]): A dictionary mapping feature names to their corresponding tf tensors. feature_keys (Sequence[str]): A list of feature names specifying the order in which tensors should be concatenated. label_keys (Sequence[str]): Name of the label columns for the current entity. - dtype (tf.dtypes.DType): Type to cast concatenated tensors to. Returns: Tuple[ @@ -176,8 +174,8 @@ def _concatenate_features_by_names( # it back to int. Note that this is ok for small int values (less than 2^24, or ~16 million). # For large int values, we will need to round it when converting back # from float, as otherwise there will be precision loss. - if tensor.dtype != dtype: - tensor = tf.cast(tensor, dtype) + if tensor.dtype != tf.float32: + tensor = tf.cast(tensor, tf.float32) # Reshape 1D tensor to column vector if len(tensor.shape) == 1: diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index affda5ffd..94d3b1c26 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -68,7 +68,9 @@ def quantize_record_batch( ] arrays = [batch.column(i) for i in keep_indices] names = [batch.schema.names[i] for i in keep_indices] - arrays.append(pa.array([row.tobytes() for row in packed], type=pa.binary())) + arrays.append( + pa.array([[row.tobytes()] for row in packed], type=pa.list_(pa.binary())) + ) names.append(_NODE_QUANTIZED_FEATURE_KEY) return pa.RecordBatch.from_arrays(arrays, names=names) diff --git a/tests/test_assets/distributed/run_distributed_partitioner.py b/tests/test_assets/distributed/run_distributed_partitioner.py index 07dce41d6..f557e283b 100644 --- a/tests/test_assets/distributed/run_distributed_partitioner.py +++ b/tests/test_assets/distributed/run_distributed_partitioner.py @@ -120,7 +120,7 @@ def run_distributed_partitioner( del node_features ( output_node_features, - output_quantized_node_features, # TODO(jchmura-sc): Update tests once API is decided on + _, output_node_labels, ) = dist_partitioner.partition_node_features_and_labels( node_partition_book=output_node_partition_book From 44677320723826751282a5e94b2c146a25faef9b Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 17:50:10 +0000 Subject: [PATCH 28/44] Derive packed feature dim from in FeatureQuantizationMetadata --- .../utils/feature_quantization/torch_ops.py | 9 ++----- .../serialized_graph_metadata_translator.py | 25 +++++++----------- .../data_preprocessor/data_preprocessor.py | 1 - .../lib/transform/feature_quantization.py | 2 -- gigl/types/graph.py | 16 +++++++++++- .../research/gbml/preprocessed_metadata.proto | 10 +++---- .../gbml/preprocessed_metadata_pb2.py | 26 +++++++++---------- .../gbml/preprocessed_metadata_pb2.pyi | 6 +---- 8 files changed, 45 insertions(+), 50 deletions(-) diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py index 872d1dc0a..b3a492d00 100644 --- a/gigl/common/utils/feature_quantization/torch_ops.py +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -10,14 +10,9 @@ def dequantize_torch_tensor( """Reconstruct approximate float features from packed uint8 codes.""" q = metadata - VALID_BITS = (1, 2, 4, 8) - if q.bits not in VALID_BITS: - raise ValueError(f"bits must be one of {VALID_BITS}, got {q.bits}") - per_byte = 8 // q.bits - expected_packed_dim = (q.dequantized_feature_dim + per_byte - 1) // per_byte - if packed_features.size(-1) != expected_packed_dim: + if packed_features.size(-1) != q.packed_feature_dim: raise ValueError( - f"Expected packed feature dim {expected_packed_dim} for " + 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)}." ) diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index 755cb4e83..ead3f03ee 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -38,8 +38,9 @@ def _build_serialized_tfrecord_entity_info( if isinstance(preprocessed_metadata, PreprocessedMetadata.NodeMetadataOutput): if preprocessed_metadata.HasField("quantized_feature_metadata"): quantized_metadata = preprocessed_metadata.quantized_feature_metadata + metadata = _build_feature_quantization_metadata(quantized_metadata) quantized_feature_keys = [quantized_metadata.quantized_feature_key] - quantized_feature_dim = quantized_metadata.packed_feature_dim + quantized_feature_dim = metadata.packed_feature_dim stored_keys = set(preprocessed_metadata.feature_keys) stored_keys.update(preprocessed_metadata.label_keys) @@ -71,25 +72,19 @@ def _build_serialized_tfrecord_entity_info( def _build_feature_quantization_metadata( quantized_metadata: PreprocessedMetadata.FeatureQuantizationMetadata, ) -> FeatureQuantizationMetadata: - bits = quantized_metadata.bits + metadata = FeatureQuantizationMetadata( + bits=quantized_metadata.bits, + dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, + dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), + ) state = quantized_metadata.WhichOneof("state") - if bits not in {1, 2, 4, 8}: - raise ValueError( - f"Expected quantized feature bits to be one of 1, 2, 4, or 8, got {bits}." - ) - expected_state = "centroid" if bits == 1 else "linear" + expected_state = "centroid" if metadata.bits == 1 else "linear" if state != expected_state: raise ValueError( - f"Expected {expected_state} quantization state for {bits}-bit features." + f"Expected {expected_state} quantization state for {metadata.bits}-bit features." ) - metadata = FeatureQuantizationMetadata( - bits=bits, - packed_feature_dim=quantized_metadata.packed_feature_dim, - dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, - dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), - ) - if bits == 1: + if metadata.bits == 1: metadata.neg_mean = quantized_metadata.centroid.neg_mean metadata.pos_mean = quantized_metadata.centroid.pos_mean else: diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 65f1880e5..63a5996b7 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -495,7 +495,6 @@ def generate_preprocessed_metadata_pb( quantized_feature_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( quantized_feature_key=metadata["quantized_feature_key"], dequantized_feature_keys=metadata["dequantized_feature_keys"], - packed_feature_dim=metadata["packed_feature_dim"], dequantized_feature_dim=metadata["dequantized_feature_dim"], bits=bits, ) diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index 94d3b1c26..756e65f1c 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -79,12 +79,10 @@ def feature_quantization_metadata_json( stats: dict[str, float], spec: FeatureQuantizationSpec, ) -> str: - per_byte = 8 // spec.bits dim = len(spec.feature_keys) metadata = { "quantized_feature_key": _NODE_QUANTIZED_FEATURE_KEY, "dequantized_feature_keys": list(spec.feature_keys), - "packed_feature_dim": (dim + per_byte - 1) // per_byte, "dequantized_feature_dim": dim, "bits": spec.bits, } diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 8e48700e7..2a3569783 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -118,7 +118,6 @@ class FeatureQuantizationMetadata: """Metadata needed to unpack/dequantize append-only packed features.""" bits: int - packed_feature_dim: int dequantized_feature_dim: int dequantized_feature_keys: tuple[str, ...] = () clip_min: Optional[float] = None @@ -126,6 +125,21 @@ class FeatureQuantizationMetadata: neg_mean: Optional[float] = None pos_mean: Optional[float] = None + def __post_init__(self) -> None: + valid_bits = (1, 2, 4, 8) + if self.bits not in valid_bits: + raise ValueError(f"bits must be one of {valid_bits}, got {self.bits}") + if self.dequantized_feature_dim < 0: + raise ValueError( + "dequantized_feature_dim must be non-negative, got " + f"{self.dequantized_feature_dim}" + ) + + @property + def packed_feature_dim(self) -> int: + per_byte = 8 // self.bits + return (self.dequantized_feature_dim + per_byte - 1) // per_byte + def _get_label_edges( labeled_edge_index: torch.Tensor, diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index cd0a64a89..c1ec3f738 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -24,16 +24,14 @@ message PreprocessedMetadata{ string quantized_feature_key = 1; // Model-visible feature names produced after dequantizing and appending. repeated string dequantized_feature_keys = 2; - // Number of uint8 columns in the packed sidecar tensor. - uint32 packed_feature_dim = 3; // Number of fp32 columns produced after unpacking/dequantizing. - uint32 dequantized_feature_dim = 4; + uint32 dequantized_feature_dim = 3; // 1 means centroid; 2, 4, and 8 mean linear buckets. - uint32 bits = 5; + uint32 bits = 4; oneof state { - LinearQuantizationState linear = 6; - CentroidQuantizationState centroid = 7; + LinearQuantizationState linear = 5; + CentroidQuantizationState centroid = 6; } } diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.py b/snapchat/research/gbml/preprocessed_metadata_pb2.py index 98767a69b..72067fd31 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.py +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.py @@ -14,7 +14,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xc4\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xe6\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1d\n\x15quantized_feature_key\x18\x01 \x01(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1a\n\x12packed_feature_dim\x18\x03 \x01(\r\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x04 \x01(\r\x12\x0c\n\x04\x62its\x18\x05 \x01(\r\x12V\n\x06linear\x18\x06 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x07 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xa8\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xca\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1d\n\x15quantized_feature_key\x18\x01 \x01(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x03 \x01(\r\x12\x0c\n\x04\x62its\x18\x04 \x01(\r\x12V\n\x06linear\x18\x05 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x06 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') @@ -106,21 +106,21 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._options = None _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_options = b'8\001' _PREPROCESSEDMETADATA._serialized_start=79 - _PREPROCESSEDMETADATA._serialized_end=2195 + _PREPROCESSEDMETADATA._serialized_end=2167 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_start=395 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_end=456 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_start=458 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=521 _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_start=524 - _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=882 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=885 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1279 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1282 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1505 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1508 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1903 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1906 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2049 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2052 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2195 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=854 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=857 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1251 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1254 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1477 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1480 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1875 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1878 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2021 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2024 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2167 # @@protoc_insertion_point(module_scope) diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index df328f07a..b0ae2422a 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -63,7 +63,6 @@ class PreprocessedMetadata(google.protobuf.message.Message): QUANTIZED_FEATURE_KEY_FIELD_NUMBER: builtins.int DEQUANTIZED_FEATURE_KEYS_FIELD_NUMBER: builtins.int - PACKED_FEATURE_DIM_FIELD_NUMBER: builtins.int DEQUANTIZED_FEATURE_DIM_FIELD_NUMBER: builtins.int BITS_FIELD_NUMBER: builtins.int LINEAR_FIELD_NUMBER: builtins.int @@ -73,8 +72,6 @@ class PreprocessedMetadata(google.protobuf.message.Message): @property def dequantized_feature_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """Model-visible feature names produced after dequantizing and appending.""" - packed_feature_dim: builtins.int - """Number of uint8 columns in the packed sidecar tensor.""" dequantized_feature_dim: builtins.int """Number of fp32 columns produced after unpacking/dequantizing.""" bits: builtins.int @@ -88,14 +85,13 @@ class PreprocessedMetadata(google.protobuf.message.Message): *, quantized_feature_key: builtins.str = ..., dequantized_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., - packed_feature_dim: builtins.int = ..., dequantized_feature_dim: builtins.int = ..., bits: builtins.int = ..., linear: global___PreprocessedMetadata.LinearQuantizationState | None = ..., centroid: global___PreprocessedMetadata.CentroidQuantizationState | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["centroid", b"centroid", "linear", b"linear", "state", b"state"]) -> builtins.bool: ... - def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_dim", b"packed_feature_dim", "quantized_feature_key", b"quantized_feature_key", "state", b"state"]) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "quantized_feature_key", b"quantized_feature_key", "state", b"state"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["state", b"state"]) -> typing_extensions.Literal["linear", "centroid"] | None: ... class NodeMetadataOutput(google.protobuf.message.Message): From 543ad0281de042c465d8c391805bca9ea81cdefc Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 18:40:42 +0000 Subject: [PATCH 29/44] WIP Fix raw schema write --- .../data_preprocessor/data_preprocessor.py | 35 +++++++++++++++---- .../lib/transform/feature_quantization.py | 20 +++++------ .../data_preprocessor/lib/transform/utils.py | 11 ++++-- .../research/gbml/preprocessed_metadata.proto | 8 ++--- .../gbml/preprocessed_metadata_pb2.pyi | 8 ++--- 5 files changed, 51 insertions(+), 31 deletions(-) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 63a5996b7..cceec78bb 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -11,6 +11,7 @@ import tensorflow_data_validation as tfdv from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult from apache_beam.runners.runner import PipelineState +from google.protobuf import text_format from tensorflow_transform.tf_metadata import schema_utils import gigl.common.utils.dataflow @@ -58,6 +59,9 @@ EdgeDataReference, NodeDataReference, ) +from gigl.src.data_preprocessor.lib.transform.feature_quantization import ( + apply_feature_quantization_schema, +) from gigl.src.data_preprocessor.lib.transform.transformed_features_info import ( TransformedFeaturesInfo, ) @@ -274,23 +278,42 @@ def __get_feature_dimension_for_single_data_reference( feature_dimension += feature_shape[0] return feature_dimension + feature_outputs = preprocessing_spec.features_outputs + q_spec = ( + preprocessing_spec.feature_quantization_spec + if isinstance(preprocessing_spec, NodeDataPreprocessingSpec) + else None + ) + if q_spec is not None: + # Feature quantization runs after TFT, so rewrite the trainer-facing + # schema to match the final TFRecords rather than the pure TFT output. + # TODO(jchmura-sc): Consider writing a separate final schema artifact. + schema_path = transformed_features_info.transformed_features_schema_path + schema = tfdv.load_schema_text(schema_path.uri) + schema = apply_feature_quantization_schema(schema, q_spec) + with tf.io.gfile.GFile(schema_path.uri, "w") as f: + f.write(text_format.MessageToString(schema)) + if feature_outputs is not None: + q_keys = set(q_spec.feature_keys) + feature_outputs = [ + feature for feature in feature_outputs if feature not in q_keys + ] + # Find and save the feature dimension if there is any - if preprocessing_spec.features_outputs is not None: + if feature_outputs is not None: transformed_features_info.feature_dim_output = __get_feature_dimension_for_single_data_reference( schema_path=transformed_features_info.transformed_features_schema_path, - feature_outputs=preprocessing_spec.features_outputs, + feature_outputs=feature_outputs, ) # Carry forward the identifier, features and label outputs from the preprocessing spec. transformed_features_info.identifier_output = ( preprocessing_spec.identifier_output ) - transformed_features_info.features_outputs = preprocessing_spec.features_outputs + transformed_features_info.features_outputs = feature_outputs transformed_features_info.label_outputs = preprocessing_spec.labels_outputs if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): - transformed_features_info.feature_quantization_spec = ( - preprocessing_spec.feature_quantization_spec - ) + transformed_features_info.feature_quantization_spec = q_spec if isinstance(feature_transform_pipeline_result, DataflowPipelineResult): pipeline_state: str = feature_transform_pipeline_result.state diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index 756e65f1c..60556c310 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -6,7 +6,6 @@ import pyarrow as pa from apache_beam.transforms.stats import ApproximateQuantiles from tensorflow_metadata.proto.v0 import schema_pb2 -from tensorflow_transform.tf_metadata.dataset_metadata import DatasetMetadata from gigl.common.logger import Logger from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec @@ -92,18 +91,17 @@ def feature_quantization_metadata_json( def apply_feature_quantization_schema( - metadata: DatasetMetadata, - spec: FeatureQuantizationSpec, -) -> DatasetMetadata: + schema: schema_pb2.Schema, spec: FeatureQuantizationSpec +) -> schema_pb2.Schema: drop_keys = set(spec.feature_keys) | {_NODE_QUANTIZED_FEATURE_KEY} - schema = schema_pb2.Schema() - schema.CopyFrom(metadata.schema) + quantized_schema = schema_pb2.Schema() + quantized_schema.CopyFrom(schema) kept_features = [ - feature for feature in schema.feature if feature.name not in drop_keys + feature for feature in quantized_schema.feature if feature.name not in drop_keys ] - del schema.feature[:] - schema.feature.extend(kept_features) - quantized_feature = schema.feature.add() + del quantized_schema.feature[:] + quantized_schema.feature.extend(kept_features) + quantized_feature = quantized_schema.feature.add() quantized_feature.name = _NODE_QUANTIZED_FEATURE_KEY quantized_feature.type = schema_pb2.BYTES quantized_feature.value_count.min = 1 @@ -113,7 +111,7 @@ def apply_feature_quantization_schema( f"{len(spec.feature_keys)} features and added bytes feature " f"{_NODE_QUANTIZED_FEATURE_KEY}." ) - return DatasetMetadata(schema) + return quantized_schema def _build_abs_feature_values( diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index b93b3e9bd..1633e5b8b 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -12,6 +12,7 @@ from tensorflow_metadata.proto.v0 import schema_pb2, statistics_pb2 from tensorflow_transform import beam as tft_beam from tensorflow_transform.tf_metadata import schema_utils +from tensorflow_transform.tf_metadata.dataset_metadata import DatasetMetadata from tfx_bsl.tfxio.record_based_tfxio import RecordBasedTFXIO from gigl.common import GcsUri, LocalUri, Uri @@ -395,14 +396,18 @@ def get_load_data_and_transform_pipeline_component( ) ) if should_use_existing_transform_fn: - resolved_transformed_metadata = apply_feature_quantization_schema( - transformed_metadata, spec=q_spec + resolved_transformed_metadata = DatasetMetadata( + apply_feature_quantization_schema( + transformed_metadata.schema, spec=q_spec + ) ) else: quantized_metadata = analyzed_transform_fn[1].deferred_metadata | ( "Apply feature quantization schema" >> beam.Map( - apply_feature_quantization_schema, + lambda metadata, spec: DatasetMetadata( + apply_feature_quantization_schema(metadata.schema, spec) + ), spec=q_spec, ) ) diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index c1ec3f738..f69c459ce 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -4,11 +4,9 @@ package snapchat.research.gbml; message PreprocessedMetadata{ message LinearQuantizationState{ - // Value produced by quantized code 0. Inputs are clipped to this value - // before linear quantization. + // Lower clipping bound; dequantized value for linear code 0. float clip_min = 1; - // Value produced by the largest quantized code. Inputs are clipped to this - // value before linear quantization. + // Upper clipping bound; dequantized value for max linear code. float clip_max = 2; } @@ -28,7 +26,7 @@ message PreprocessedMetadata{ uint32 dequantized_feature_dim = 3; // 1 means centroid; 2, 4, and 8 mean linear buckets. uint32 bits = 4; - + // Stats required to dequantize features for the selected bit width. oneof state { LinearQuantizationState linear = 5; CentroidQuantizationState centroid = 6; diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index b0ae2422a..c1bd01a7b 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -26,13 +26,9 @@ class PreprocessedMetadata(google.protobuf.message.Message): CLIP_MIN_FIELD_NUMBER: builtins.int CLIP_MAX_FIELD_NUMBER: builtins.int clip_min: builtins.float - """Value produced by quantized code 0. Inputs are clipped to this value - before linear quantization. - """ + """Lower clipping bound; dequantized value for linear code 0.""" clip_max: builtins.float - """Value produced by the largest quantized code. Inputs are clipped to this - value before linear quantization. - """ + """Upper clipping bound; dequantized value for max linear code.""" def __init__( self, *, From daaab317585d8987a7849e5ea175af16af6b24dd Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 20:15:41 +0000 Subject: [PATCH 30/44] Move numpy ops out of preprocessor --- .../utils/feature_quantization/__init__.py | 3 -- .../utils/feature_quantization/numpy_ops.py | 42 +++++++++++++++++++ gigl/distributed/utils/neighborloader.py | 2 +- .../lib/transform/feature_quantization.py | 29 +------------ 4 files changed, 45 insertions(+), 31 deletions(-) create mode 100644 gigl/common/utils/feature_quantization/numpy_ops.py diff --git a/gigl/common/utils/feature_quantization/__init__.py b/gigl/common/utils/feature_quantization/__init__.py index 010ce33e9..e69de29bb 100644 --- a/gigl/common/utils/feature_quantization/__init__.py +++ b/gigl/common/utils/feature_quantization/__init__.py @@ -1,3 +0,0 @@ -from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor - -__all__ = ["dequantize_torch_tensor"] diff --git a/gigl/common/utils/feature_quantization/numpy_ops.py b/gigl/common/utils/feature_quantization/numpy_ops.py new file mode 100644 index 000000000..e161406c1 --- /dev/null +++ b/gigl/common/utils/feature_quantization/numpy_ops.py @@ -0,0 +1,42 @@ +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.""" + if bits not in (1, 2, 4, 8): + raise ValueError(f"bits must be one of 1, 2, 4, or 8, got {bits}.") + if codes.ndim != 2: + raise ValueError(f"Expected a 2D code array, got shape {codes.shape}.") + 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) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index c848c868f..83b2301c6 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -13,7 +13,7 @@ from torch_geometric.typing import EdgeType, NodeType from gigl.common.logger import Logger -from gigl.common.utils.feature_quantization import dequantize_torch_tensor +from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor from gigl.distributed.sampler import NODE_QUANTIZED_FEATURES_METADATA_KEY from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_NODE_TYPE, diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index 60556c310..7e3ceceac 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -8,6 +8,7 @@ from tensorflow_metadata.proto.v0 import schema_pb2 from gigl.common.logger import Logger +from gigl.common.utils.feature_quantization.numpy_ops import quantize_ndarray from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec logger = Logger() @@ -49,18 +50,7 @@ def quantize_record_batch( stats: dict[str, float], ) -> pa.RecordBatch: features = _build_feature_matrix(batch, spec.feature_keys) - if spec.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 << spec.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) - - packed = _pack_codes(codes, spec.bits) + packed = quantize_ndarray(features, bits=spec.bits, stats=stats) drop_keys = set(spec.feature_keys) | {_NODE_QUANTIZED_FEATURE_KEY} keep_indices = [ i for i, name in enumerate(batch.schema.names) if name not in drop_keys @@ -140,21 +130,6 @@ def _build_feature_matrix(batch: pa.RecordBatch, feature_keys: list[str]) -> np. return np.stack(cols, axis=1) -def _pack_codes(codes: np.ndarray, bits: int) -> np.ndarray: - # Cast to uint16 for integer shift/multiply/sum packing math. - 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) - # Place the first code in each chunk into the highest bits of the byte. - 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) - - def _linear_stats_from_quantiles(quantiles: list[float]) -> dict[str, float]: if not quantiles: raise ValueError("Cannot compute quantization stats from no values.") From 7e565efb01b316e6027c61600c5cd4a5e9edfd7b Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 21:03:58 +0000 Subject: [PATCH 31/44] wip --- gigl/common/data/dataloaders.py | 38 +++++++++---------- gigl/common/data/load_torch_tensors.py | 22 ++++------- .../serialized_graph_metadata_translator.py | 9 +++-- 3 files changed, 31 insertions(+), 38 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 7e99cf749..3e2f12d75 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -1,6 +1,6 @@ import time from copy import deepcopy -from dataclasses import dataclass, field +from dataclasses import dataclass from functools import partial from typing import Callable, NamedTuple, Optional, Sequence, Tuple, Union @@ -45,12 +45,12 @@ 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 names to load for the current node entity. - quantized_feature_keys: Sequence[str] = field(default_factory=list) + # Packed uint8 feature name to load for the current node entity. + quantized_feature_key: Optional[str] = None # Number of packed uint8 columns for the current node entity. quantized_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) + label_keys: Sequence[str] = () # The regex pattern to match the TFRecord files at the specified prefix tfrecord_uri_pattern: str = ".*-of-.*\.tfrecord(\.gz)?$" @@ -376,7 +376,7 @@ def load_as_torch_tensors( """ entity_key = serialized_tf_record_info.entity_key feature_keys = serialized_tf_record_info.feature_keys - quantized_feature_keys = serialized_tf_record_info.quantized_feature_keys + quantized_feature_key = serialized_tf_record_info.quantized_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 @@ -398,14 +398,16 @@ def load_as_torch_tensors( feature_spec_dict[entity_key] = tf.io.FixedLenFeature( shape=[], dtype=tf.int64 ) - for feature_key in quantized_feature_keys: - if feature_key not in feature_spec_dict: - logger.info( - f"Injecting quantized feature key {feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" - ) - feature_spec_dict[feature_key] = tf.io.FixedLenFeature( - shape=[], dtype=tf.string - ) + if ( + quantized_feature_key is not None + and quantized_feature_key not in feature_spec_dict + ): + logger.info( + f"Injecting quantized feature key {quantized_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" + ) + feature_spec_dict[quantized_feature_key] = tf.io.FixedLenFeature( + shape=[], dtype=tf.string + ) else: id_concat_axis = 1 proccess_id_tensor = lambda t: tf.stack( @@ -449,7 +451,7 @@ def load_as_torch_tensors( else: empty_feature = None - if quantized_feature_keys: + if quantized_feature_key is not None: empty_quantized_feature = torch.empty( (0, serialized_tf_record_info.quantized_feature_dim), dtype=torch.uint8, @@ -491,13 +493,9 @@ def load_as_torch_tensors( feature_tensors.append(feature_tensor) if label_tensor is not None: label_tensors.append(label_tensor) - if quantized_feature_keys: - if len(quantized_feature_keys) != 1: - raise ValueError( - f"Expected one quantized feature key, got {quantized_feature_keys}." - ) + if quantized_feature_key is not None: quantized_feature_tensor = tf.io.decode_raw( - batch[quantized_feature_keys[0]], tf.uint8 + batch[quantized_feature_key], tf.uint8 ) quantized_feature_tensor = tf.reshape( quantized_feature_tensor, diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index e035076f7..3e44af0a4 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -27,10 +27,10 @@ _ID_FMT = "{entity}_ids" _FEATURE_FMT = "{entity}_features" +_QUANTIZED_FEATURE_FMT = "{entity}_quantized_features" _LABEL_FMT = "{entity}_labels" _EDGE_WEIGHTS_KEY = "edge_weights" _NODE_KEY = "node" -_NODE_QUANTIZED_KEY = "node_quantized_features" def _extract_weight_col( @@ -191,12 +191,6 @@ def _data_loading_process( graph_type, serialized_entity_tf_record_info, ) in serialized_tf_record_info.items(): - logger.info( - f"Rank {rank} loading {entity_type} graph type {graph_type} with " - f"{len(serialized_entity_tf_record_info.feature_keys)} feature keys, " - f"quantized_feature_keys={serialized_entity_tf_record_info.quantized_feature_keys}, " - f"quantized_feature_dim={serialized_entity_tf_record_info.quantized_feature_dim}" - ) # We currently do not support training with labels for edge entities if ( serialized_entity_tf_record_info.label_keys @@ -206,7 +200,7 @@ def _data_loading_process( "Label keys are not supported for edge entities" ) if ( - serialized_entity_tf_record_info.quantized_feature_keys + serialized_entity_tf_record_info.quantized_feature_key is not None and not serialized_entity_tf_record_info.is_node_entity ): raise NotImplementedError( @@ -241,9 +235,7 @@ def _data_loading_process( ) 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}. " - f"Serialized quantized_feature_keys={serialized_entity_tf_record_info.quantized_feature_keys}, " - f"quantized_feature_dim={serialized_entity_tf_record_info.quantized_feature_dim}" + 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: @@ -349,8 +341,8 @@ 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 and entity_type == _NODE_KEY: - output_dict[_NODE_QUANTIZED_KEY] = ( + if quantized_features: + output_dict[_QUANTIZED_FEATURE_FMT.format(entity=entity_type)] = ( list(quantized_features.values())[0] if is_input_homogeneous else quantized_features @@ -525,7 +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(_NODE_QUANTIZED_KEY, None) + node_quantized_features = node_output_dict.get( + _QUANTIZED_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)] diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index ead3f03ee..e7ef99c4f 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -33,18 +33,19 @@ def _build_serialized_tfrecord_entity_info( Returns: SerializedTFRecordInfo: Stored metadata for current entity """ - quantized_feature_keys: list[str] = [] + quantized_feature_key = None quantized_feature_dim = 0 if isinstance(preprocessed_metadata, PreprocessedMetadata.NodeMetadataOutput): if preprocessed_metadata.HasField("quantized_feature_metadata"): quantized_metadata = preprocessed_metadata.quantized_feature_metadata metadata = _build_feature_quantization_metadata(quantized_metadata) - quantized_feature_keys = [quantized_metadata.quantized_feature_key] + quantized_feature_key = quantized_metadata.quantized_feature_key quantized_feature_dim = metadata.packed_feature_dim stored_keys = set(preprocessed_metadata.feature_keys) stored_keys.update(preprocessed_metadata.label_keys) - stored_keys.update(quantized_feature_keys) + if quantized_feature_key is not None: + stored_keys.add(quantized_feature_key) # PreprocessedMetadataPbWrapper may add synthetic dequantized feature keys # to the logical feature schema. Only keep fields physically stored in # TFRecords here; dequantized features are reconstructed later from uint8. @@ -62,7 +63,7 @@ def _build_serialized_tfrecord_entity_info( feature_spec=filtered_feature_spec_dict, feature_dim=preprocessed_metadata.feature_dim, entity_key=entity_key, - quantized_feature_keys=quantized_feature_keys, + quantized_feature_key=quantized_feature_key, quantized_feature_dim=quantized_feature_dim, label_keys=list(preprocessed_metadata.label_keys), tfrecord_uri_pattern=tfrecord_uri_pattern, From fb44c1bcb8ff9580735c41e4034d1a44d63cf845 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 21:05:04 +0000 Subject: [PATCH 32/44] wip --- gigl/common/data/dataloaders.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 3e2f12d75..9186a0f23 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -1,6 +1,6 @@ import time from copy import deepcopy -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import partial from typing import Callable, NamedTuple, Optional, Sequence, Tuple, Union @@ -50,7 +50,7 @@ class SerializedTFRecordInfo: # Number of packed uint8 columns for the current node entity. quantized_feature_dim: int = 0 # Name of the label columns for the current entity, defaults to an empty list. - label_keys: Sequence[str] = () + label_keys: Sequence[str] = field(default_factory=list) # The regex pattern to match the TFRecord files at the specified prefix tfrecord_uri_pattern: str = ".*-of-.*\.tfrecord(\.gz)?$" From 6201e026d888fc18a11d70c38e9c2f6bec4d56e7 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 21:27:58 +0000 Subject: [PATCH 33/44] WIP --- gigl/common/utils/feature_quantization/numpy_ops.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/gigl/common/utils/feature_quantization/numpy_ops.py b/gigl/common/utils/feature_quantization/numpy_ops.py index e161406c1..bb8e30f28 100644 --- a/gigl/common/utils/feature_quantization/numpy_ops.py +++ b/gigl/common/utils/feature_quantization/numpy_ops.py @@ -26,10 +26,6 @@ def quantize_ndarray( def pack_codes(codes: np.ndarray, bits: int) -> np.ndarray: """Pack low-bit feature codes high-bits-first along the final dimension.""" - if bits not in (1, 2, 4, 8): - raise ValueError(f"bits must be one of 1, 2, 4, or 8, got {bits}.") - if codes.ndim != 2: - raise ValueError(f"Expected a 2D code array, got shape {codes.shape}.") per_byte = 8 // bits pad = (-codes.shape[-1]) % per_byte if pad: From 681d5e9042c04704e90f71cfd1e65e4b81eee6a0 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 21:38:53 +0000 Subject: [PATCH 34/44] Rename to packed_feature --- gigl/common/data/dataloaders.py | 24 ++++++++--------- gigl/common/data/load_torch_tensors.py | 10 +++---- gigl/distributed/base_sampler.py | 8 +++--- gigl/distributed/sampler.py | 2 +- gigl/distributed/utils/neighborloader.py | 8 +++--- .../serialized_graph_metadata_translator.py | 16 ++++++------ .../data_preprocessor/data_preprocessor.py | 2 +- .../lib/transform/feature_quantization.py | 22 ++++++++-------- .../research/gbml/preprocessed_metadata.proto | 2 +- .../gbml/preprocessed_metadata_pb2.py | 26 +++++++++---------- .../gbml/preprocessed_metadata_pb2.pyi | 8 +++--- 11 files changed, 64 insertions(+), 64 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 9186a0f23..498309ce3 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -46,9 +46,9 @@ class SerializedTFRecordInfo: # 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. - quantized_feature_key: Optional[str] = None + packed_feature_key: Optional[str] = None # Number of packed uint8 columns for the current node entity. - quantized_feature_dim: int = 0 + 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 @@ -376,7 +376,7 @@ def load_as_torch_tensors( """ entity_key = serialized_tf_record_info.entity_key feature_keys = serialized_tf_record_info.feature_keys - quantized_feature_key = serialized_tf_record_info.quantized_feature_key + 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 @@ -399,13 +399,13 @@ def load_as_torch_tensors( shape=[], dtype=tf.int64 ) if ( - quantized_feature_key is not None - and quantized_feature_key not in feature_spec_dict + packed_feature_key is not None + and packed_feature_key not in feature_spec_dict ): logger.info( - f"Injecting quantized feature key {quantized_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" + f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" ) - feature_spec_dict[quantized_feature_key] = tf.io.FixedLenFeature( + feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature( shape=[], dtype=tf.string ) else: @@ -451,9 +451,9 @@ def load_as_torch_tensors( else: empty_feature = None - if quantized_feature_key is not None: + if packed_feature_key is not None: empty_quantized_feature = torch.empty( - (0, serialized_tf_record_info.quantized_feature_dim), + (0, serialized_tf_record_info.packed_feature_dim), dtype=torch.uint8, ) else: @@ -493,13 +493,13 @@ def load_as_torch_tensors( feature_tensors.append(feature_tensor) if label_tensor is not None: label_tensors.append(label_tensor) - if quantized_feature_key is not None: + if packed_feature_key is not None: quantized_feature_tensor = tf.io.decode_raw( - batch[quantized_feature_key], tf.uint8 + batch[packed_feature_key], tf.uint8 ) quantized_feature_tensor = tf.reshape( quantized_feature_tensor, - [-1, serialized_tf_record_info.quantized_feature_dim], + [-1, serialized_tf_record_info.packed_feature_dim], ) quantized_feature_tensors.append(quantized_feature_tensor) num_entities_processed += ( diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index 3e44af0a4..7b3545e33 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -27,7 +27,7 @@ _ID_FMT = "{entity}_ids" _FEATURE_FMT = "{entity}_features" -_QUANTIZED_FEATURE_FMT = "{entity}_quantized_features" +_PACKED_FEATURE_FMT = "{entity}_packed_features" _LABEL_FMT = "{entity}_labels" _EDGE_WEIGHTS_KEY = "edge_weights" _NODE_KEY = "node" @@ -200,11 +200,11 @@ def _data_loading_process( "Label keys are not supported for edge entities" ) if ( - serialized_entity_tf_record_info.quantized_feature_key is not None + serialized_entity_tf_record_info.packed_feature_key is not None and not serialized_entity_tf_record_info.is_node_entity ): raise NotImplementedError( - "Quantized feature keys are not supported for edge entities" + "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, @@ -342,7 +342,7 @@ def _data_loading_process( list(features.values())[0] if is_input_homogeneous else features ) if quantized_features: - output_dict[_QUANTIZED_FEATURE_FMT.format(entity=entity_type)] = ( + output_dict[_PACKED_FEATURE_FMT.format(entity=entity_type)] = ( list(quantized_features.values())[0] if is_input_homogeneous else quantized_features @@ -518,7 +518,7 @@ 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( - _QUANTIZED_FEATURE_FMT.format(entity=_NODE_KEY), None + _PACKED_FEATURE_FMT.format(entity=_NODE_KEY), None ) node_labels = node_output_dict.get(_LABEL_FMT.format(entity=_NODE_KEY), None) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 91aec64c7..d162620bd 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -17,7 +17,7 @@ from gigl.distributed.sampler import ( NEGATIVE_LABEL_METADATA_KEY, - NODE_QUANTIZED_FEATURES_METADATA_KEY, + NODE_PACKED_FEATURES_METADATA_KEY, POSITIVE_LABEL_METADATA_KEY, ABLPNodeSamplerInput, ) @@ -354,7 +354,7 @@ async def _collate_fn( ) for ntype, quantized_nfeats in quantized_nfeat_dict.items(): result_map[ - f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{as_str(ntype)}" + f"#META.{NODE_PACKED_FEATURES_METADATA_KEY}.{as_str(ntype)}" ] = quantized_nfeats else: quantized_nfeat_fut_dict = {} @@ -366,7 +366,7 @@ async def _collate_fn( for ntype, fut in quantized_nfeat_fut_dict.items(): quantized_nfeats = await wrap_torch_future(fut) result_map[ - f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{as_str(ntype)}" + f"#META.{NODE_PACKED_FEATURES_METADATA_KEY}.{as_str(ntype)}" ] = quantized_nfeats if self.dist_edge_feature is not None and self.with_edge: efeat_fut_dict = {} @@ -431,7 +431,7 @@ async def _collate_fn( if self.dist_node_quantized_feature is not None: fut = self.dist_node_quantized_feature.async_get(output.node) quantized_nfeats = await wrap_torch_future(fut) - result_map[f"#META.{NODE_QUANTIZED_FEATURES_METADATA_KEY}"] = ( + result_map[f"#META.{NODE_PACKED_FEATURES_METADATA_KEY}"] = ( quantized_nfeats ) if self.dist_edge_feature is not None: diff --git a/gigl/distributed/sampler.py b/gigl/distributed/sampler.py index a216fa41c..7789c6731 100644 --- a/gigl/distributed/sampler.py +++ b/gigl/distributed/sampler.py @@ -8,7 +8,7 @@ POSITIVE_LABEL_METADATA_KEY: Final[str] = "gigl_positive_labels." NEGATIVE_LABEL_METADATA_KEY: Final[str] = "gigl_negative_labels." -NODE_QUANTIZED_FEATURES_METADATA_KEY: Final[str] = "node_quantized_features" +NODE_PACKED_FEATURES_METADATA_KEY: Final[str] = "node_packed_features" class ABLPNodeSamplerInput(NodeSamplerInput): diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 83b2301c6..2527564a1 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -14,7 +14,7 @@ from gigl.common.logger import Logger from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor -from gigl.distributed.sampler import NODE_QUANTIZED_FEATURES_METADATA_KEY +from gigl.distributed.sampler import NODE_PACKED_FEATURES_METADATA_KEY from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_NODE_TYPE, FeatureInfo, @@ -375,12 +375,12 @@ def materialize_node_store( DEFAULT_HOMOGENEOUS_NODE_TYPE ] metadata_key = ( - f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}." + f"{NODE_PACKED_FEATURES_METADATA_KEY}." f"{DEFAULT_HOMOGENEOUS_NODE_TYPE}" ) else: quantization_metadata = node_quantization_metadata - metadata_key = NODE_QUANTIZED_FEATURES_METADATA_KEY + metadata_key = NODE_PACKED_FEATURES_METADATA_KEY packed_features = metadata.pop(metadata_key, None) if packed_features is None: raise ValueError( @@ -396,7 +396,7 @@ def materialize_node_store( dict[NodeType, FeatureQuantizationMetadata], node_quantization_metadata ) for node_type, quantization_metadata in node_quantization_metadata.items(): - metadata_key = f"{NODE_QUANTIZED_FEATURES_METADATA_KEY}.{node_type}" + metadata_key = f"{NODE_PACKED_FEATURES_METADATA_KEY}.{node_type}" packed_features = metadata.pop(metadata_key, None) if packed_features is None: continue diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index e7ef99c4f..be5dcf909 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -33,19 +33,19 @@ def _build_serialized_tfrecord_entity_info( Returns: SerializedTFRecordInfo: Stored metadata for current entity """ - quantized_feature_key = None - quantized_feature_dim = 0 + packed_feature_key = None + packed_feature_dim = 0 if isinstance(preprocessed_metadata, PreprocessedMetadata.NodeMetadataOutput): if preprocessed_metadata.HasField("quantized_feature_metadata"): quantized_metadata = preprocessed_metadata.quantized_feature_metadata metadata = _build_feature_quantization_metadata(quantized_metadata) - quantized_feature_key = quantized_metadata.quantized_feature_key - quantized_feature_dim = metadata.packed_feature_dim + packed_feature_key = quantized_metadata.packed_feature_key + packed_feature_dim = metadata.packed_feature_dim stored_keys = set(preprocessed_metadata.feature_keys) stored_keys.update(preprocessed_metadata.label_keys) - if quantized_feature_key is not None: - stored_keys.add(quantized_feature_key) + if packed_feature_key is not None: + stored_keys.add(packed_feature_key) # PreprocessedMetadataPbWrapper may add synthetic dequantized feature keys # to the logical feature schema. Only keep fields physically stored in # TFRecords here; dequantized features are reconstructed later from uint8. @@ -63,8 +63,8 @@ def _build_serialized_tfrecord_entity_info( feature_spec=filtered_feature_spec_dict, feature_dim=preprocessed_metadata.feature_dim, entity_key=entity_key, - quantized_feature_key=quantized_feature_key, - quantized_feature_dim=quantized_feature_dim, + packed_feature_key=packed_feature_key, + packed_feature_dim=packed_feature_dim, label_keys=list(preprocessed_metadata.label_keys), tfrecord_uri_pattern=tfrecord_uri_pattern, ) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index cceec78bb..9e739863d 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -516,7 +516,7 @@ def generate_preprocessed_metadata_pb( metadata = json.loads(f.read()) bits = metadata["bits"] quantized_feature_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( - quantized_feature_key=metadata["quantized_feature_key"], + packed_feature_key=metadata["packed_feature_key"], dequantized_feature_keys=metadata["dequantized_feature_keys"], dequantized_feature_dim=metadata["dequantized_feature_dim"], bits=bits, diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index 7e3ceceac..af47f3d8f 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -12,7 +12,7 @@ from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec logger = Logger() -_NODE_QUANTIZED_FEATURE_KEY = "node_quantized_features" +_NODE_PACKED_FEATURE_KEY = "node_packed_features" _CentroidAcc = tuple[float, int, float, int] @@ -51,7 +51,7 @@ def quantize_record_batch( ) -> pa.RecordBatch: features = _build_feature_matrix(batch, spec.feature_keys) packed = quantize_ndarray(features, bits=spec.bits, stats=stats) - drop_keys = set(spec.feature_keys) | {_NODE_QUANTIZED_FEATURE_KEY} + drop_keys = set(spec.feature_keys) | {_NODE_PACKED_FEATURE_KEY} keep_indices = [ i for i, name in enumerate(batch.schema.names) if name not in drop_keys ] @@ -60,7 +60,7 @@ def quantize_record_batch( arrays.append( pa.array([[row.tobytes()] for row in packed], type=pa.list_(pa.binary())) ) - names.append(_NODE_QUANTIZED_FEATURE_KEY) + names.append(_NODE_PACKED_FEATURE_KEY) return pa.RecordBatch.from_arrays(arrays, names=names) @@ -70,7 +70,7 @@ def feature_quantization_metadata_json( ) -> str: dim = len(spec.feature_keys) metadata = { - "quantized_feature_key": _NODE_QUANTIZED_FEATURE_KEY, + "packed_feature_key": _NODE_PACKED_FEATURE_KEY, "dequantized_feature_keys": list(spec.feature_keys), "dequantized_feature_dim": dim, "bits": spec.bits, @@ -83,7 +83,7 @@ def feature_quantization_metadata_json( def apply_feature_quantization_schema( schema: schema_pb2.Schema, spec: FeatureQuantizationSpec ) -> schema_pb2.Schema: - drop_keys = set(spec.feature_keys) | {_NODE_QUANTIZED_FEATURE_KEY} + drop_keys = set(spec.feature_keys) | {_NODE_PACKED_FEATURE_KEY} quantized_schema = schema_pb2.Schema() quantized_schema.CopyFrom(schema) kept_features = [ @@ -91,15 +91,15 @@ def apply_feature_quantization_schema( ] del quantized_schema.feature[:] quantized_schema.feature.extend(kept_features) - quantized_feature = quantized_schema.feature.add() - quantized_feature.name = _NODE_QUANTIZED_FEATURE_KEY - quantized_feature.type = schema_pb2.BYTES - quantized_feature.value_count.min = 1 - quantized_feature.value_count.max = 1 + packed_feature = quantized_schema.feature.add() + packed_feature.name = _NODE_PACKED_FEATURE_KEY + packed_feature.type = schema_pb2.BYTES + packed_feature.value_count.min = 1 + packed_feature.value_count.max = 1 logger.info( f"Updated transformed schema for feature quantization: dropped " f"{len(spec.feature_keys)} features and added bytes feature " - f"{_NODE_QUANTIZED_FEATURE_KEY}." + f"{_NODE_PACKED_FEATURE_KEY}." ) return quantized_schema diff --git a/proto/snapchat/research/gbml/preprocessed_metadata.proto b/proto/snapchat/research/gbml/preprocessed_metadata.proto index f69c459ce..c45be4967 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -19,7 +19,7 @@ message PreprocessedMetadata{ message FeatureQuantizationMetadata{ // Field in output TFRecords that stores packed uint8 features. - string quantized_feature_key = 1; + string packed_feature_key = 1; // Model-visible feature names produced after dequantizing and appending. repeated string dequantized_feature_keys = 2; // Number of fp32 columns produced after unpacking/dequantizing. diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.py b/snapchat/research/gbml/preprocessed_metadata_pb2.py index 72067fd31..9f377e2b1 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.py +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.py @@ -14,7 +14,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xa8\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xca\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1d\n\x15quantized_feature_key\x18\x01 \x01(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x03 \x01(\r\x12\x0c\n\x04\x62its\x18\x04 \x01(\r\x12V\n\x06linear\x18\x05 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x06 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2snapchat/research/gbml/preprocessed_metadata.proto\x12\x16snapchat.research.gbml\"\xa5\x10\n\x14PreprocessedMetadata\x12\x8f\x01\n,condensed_node_type_to_preprocessed_metadata\x18\x01 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedNodeTypeToPreprocessedMetadataEntry\x12\x8f\x01\n,condensed_edge_type_to_preprocessed_metadata\x18\x02 \x03(\x0b\x32Y.snapchat.research.gbml.PreprocessedMetadata.CondensedEdgeTypeToPreprocessedMetadataEntry\x1a=\n\x17LinearQuantizationState\x12\x10\n\x08\x63lip_min\x18\x01 \x01(\x02\x12\x10\n\x08\x63lip_max\x18\x02 \x01(\x02\x1a?\n\x19\x43\x65ntroidQuantizationState\x12\x10\n\x08neg_mean\x18\x01 \x01(\x02\x12\x10\n\x08pos_mean\x18\x02 \x01(\x02\x1a\xc7\x02\n\x1b\x46\x65\x61tureQuantizationMetadata\x12\x1a\n\x12packed_feature_key\x18\x01 \x01(\t\x12 \n\x18\x64\x65quantized_feature_keys\x18\x02 \x03(\t\x12\x1f\n\x17\x64\x65quantized_feature_dim\x18\x03 \x01(\r\x12\x0c\n\x04\x62its\x18\x04 \x01(\r\x12V\n\x06linear\x18\x05 \x01(\x0b\x32\x44.snapchat.research.gbml.PreprocessedMetadata.LinearQuantizationStateH\x00\x12Z\n\x08\x63\x65ntroid\x18\x06 \x01(\x0b\x32\x46.snapchat.research.gbml.PreprocessedMetadata.CentroidQuantizationStateH\x00\x42\x07\n\x05state\x1a\x8a\x03\n\x12NodeMetadataOutput\x12\x13\n\x0bnode_id_key\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x02 \x03(\t\x12\x12\n\nlabel_keys\x18\x03 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x04 \x01(\t\x12\x12\n\nschema_uri\x18\x05 \x01(\t\x12$\n\x1c\x65numerated_node_ids_bq_table\x18\x06 \x01(\t\x12%\n\x1d\x65numerated_node_data_bq_table\x18\x07 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x08 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\t \x01(\t\x12l\n\x1aquantized_feature_metadata\x18\n \x01(\x0b\x32H.snapchat.research.gbml.PreprocessedMetadata.FeatureQuantizationMetadataB\x0e\n\x0c_feature_dim\x1a\xdf\x01\n\x10\x45\x64geMetadataInfo\x12\x14\n\x0c\x66\x65\x61ture_keys\x18\x01 \x03(\t\x12\x12\n\nlabel_keys\x18\x02 \x03(\t\x12\x1b\n\x13tfrecord_uri_prefix\x18\x03 \x01(\t\x12\x12\n\nschema_uri\x18\x04 \x01(\t\x12%\n\x1d\x65numerated_edge_data_bq_table\x18\x05 \x01(\t\x12\x18\n\x0b\x66\x65\x61ture_dim\x18\x06 \x01(\rH\x00\x88\x01\x01\x12\x1f\n\x17transform_fn_assets_uri\x18\x07 \x01(\tB\x0e\n\x0c_feature_dim\x1a\x8b\x03\n\x12\x45\x64geMetadataOutput\x12\x17\n\x0fsrc_node_id_key\x18\x01 \x01(\t\x12\x17\n\x0f\x64st_node_id_key\x18\x02 \x01(\t\x12U\n\x0emain_edge_info\x18\x03 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfo\x12^\n\x12positive_edge_info\x18\x04 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x00\x88\x01\x01\x12^\n\x12negative_edge_info\x18\x05 \x01(\x0b\x32=.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataInfoH\x01\x88\x01\x01\x42\x15\n\x13_positive_edge_infoB\x15\n\x13_negative_edge_info\x1a\x8f\x01\n,CondensedNodeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.NodeMetadataOutput:\x02\x38\x01\x1a\x8f\x01\n,CondensedEdgeTypeToPreprocessedMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12N\n\x05value\x18\x02 \x01(\x0b\x32?.snapchat.research.gbml.PreprocessedMetadata.EdgeMetadataOutput:\x02\x38\x01\x62\x06proto3') @@ -106,21 +106,21 @@ _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._options = None _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_options = b'8\001' _PREPROCESSEDMETADATA._serialized_start=79 - _PREPROCESSEDMETADATA._serialized_end=2167 + _PREPROCESSEDMETADATA._serialized_end=2164 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_start=395 _PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE._serialized_end=456 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_start=458 _PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE._serialized_end=521 _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_start=524 - _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=854 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=857 - _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1251 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1254 - _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1477 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1480 - _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1875 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1878 - _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2021 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2024 - _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2167 + _PREPROCESSEDMETADATA_FEATUREQUANTIZATIONMETADATA._serialized_end=851 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_start=854 + _PREPROCESSEDMETADATA_NODEMETADATAOUTPUT._serialized_end=1248 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_start=1251 + _PREPROCESSEDMETADATA_EDGEMETADATAINFO._serialized_end=1474 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_start=1477 + _PREPROCESSEDMETADATA_EDGEMETADATAOUTPUT._serialized_end=1872 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=1875 + _PREPROCESSEDMETADATA_CONDENSEDNODETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2018 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_start=2021 + _PREPROCESSEDMETADATA_CONDENSEDEDGETYPETOPREPROCESSEDMETADATAENTRY._serialized_end=2164 # @@protoc_insertion_point(module_scope) diff --git a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi index c1bd01a7b..5f35dc8b4 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -57,13 +57,13 @@ class PreprocessedMetadata(google.protobuf.message.Message): class FeatureQuantizationMetadata(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor - QUANTIZED_FEATURE_KEY_FIELD_NUMBER: builtins.int + PACKED_FEATURE_KEY_FIELD_NUMBER: builtins.int DEQUANTIZED_FEATURE_KEYS_FIELD_NUMBER: builtins.int DEQUANTIZED_FEATURE_DIM_FIELD_NUMBER: builtins.int BITS_FIELD_NUMBER: builtins.int LINEAR_FIELD_NUMBER: builtins.int CENTROID_FIELD_NUMBER: builtins.int - quantized_feature_key: builtins.str + packed_feature_key: builtins.str """Field in output TFRecords that stores packed uint8 features.""" @property def dequantized_feature_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: @@ -79,7 +79,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): def __init__( self, *, - quantized_feature_key: builtins.str = ..., + packed_feature_key: builtins.str = ..., dequantized_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., dequantized_feature_dim: builtins.int = ..., bits: builtins.int = ..., @@ -87,7 +87,7 @@ class PreprocessedMetadata(google.protobuf.message.Message): centroid: global___PreprocessedMetadata.CentroidQuantizationState | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["centroid", b"centroid", "linear", b"linear", "state", b"state"]) -> builtins.bool: ... - def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "quantized_feature_key", b"quantized_feature_key", "state", b"state"]) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["bits", b"bits", "centroid", b"centroid", "dequantized_feature_dim", b"dequantized_feature_dim", "dequantized_feature_keys", b"dequantized_feature_keys", "linear", b"linear", "packed_feature_key", b"packed_feature_key", "state", b"state"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["state", b"state"]) -> typing_extensions.Literal["linear", "centroid"] | None: ... class NodeMetadataOutput(google.protobuf.message.Message): From 496fa9458cf2ba4f135087c73d784f898de98862 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 22:17:03 +0000 Subject: [PATCH 35/44] Dont need spec in transformed feature info --- gigl/src/data_preprocessor/data_preprocessor.py | 12 ++++++------ .../lib/transform/transformed_features_info.py | 2 -- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 9e739863d..c394108f6 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -312,8 +312,6 @@ def __get_feature_dimension_for_single_data_reference( ) transformed_features_info.features_outputs = feature_outputs transformed_features_info.label_outputs = preprocessing_spec.labels_outputs - if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): - transformed_features_info.feature_quantization_spec = q_spec if isinstance(feature_transform_pipeline_result, DataflowPipelineResult): pipeline_state: str = feature_transform_pipeline_result.state @@ -509,10 +507,12 @@ def generate_preprocessed_metadata_pb( feature_dim=feature_dim_output, transform_fn_assets_uri=node_transformed_features_info.transformed_features_transform_fn_assets_path.uri, ) - if node_transformed_features_info.feature_quantization_spec is not None: - with tf.io.gfile.GFile( - node_transformed_features_info.feature_quantization_metadata_path.uri - ) as f: + metadata_path = ( + node_transformed_features_info.feature_quantization_metadata_path.uri + ) + if tf.io.gfile.exists(metadata_path): + logger.info(f"Adding feature quantization metadata from {metadata_path}") + with tf.io.gfile.GFile(metadata_path) as f: metadata = json.loads(f.read()) bits = metadata["bits"] quantized_feature_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( diff --git a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py index 38044c07b..903a02057 100644 --- a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py +++ b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py @@ -8,7 +8,6 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.src.data_preprocessor.lib.types import ( EdgeOutputIdentifier, - FeatureQuantizationSpec, NodeOutputIdentifier, ) @@ -33,7 +32,6 @@ class TransformedFeaturesInfo: features_outputs: Optional[list[str]] = None label_outputs: Optional[list[str]] = None feature_dim_output: Optional[int] = None - feature_quantization_spec: Optional[FeatureQuantizationSpec] = None custom_identifier: Optional[str] = None def __init__( From a64951ef4f4e383e8efa695656847efea0998cbb Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 23:09:00 +0000 Subject: [PATCH 36/44] Keep non-feature quant diff minimal prepreprocess --- .../data_preprocessor/data_preprocessor.py | 41 ++++----- .../lib/transform/feature_quantization.py | 90 +++++++++++++++---- .../data_preprocessor/lib/transform/utils.py | 58 +++--------- 3 files changed, 104 insertions(+), 85 deletions(-) diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index c394108f6..d6c25e844 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -11,7 +11,6 @@ import tensorflow_data_validation as tfdv from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult from apache_beam.runners.runner import PipelineState -from google.protobuf import text_format from tensorflow_transform.tf_metadata import schema_utils import gigl.common.utils.dataflow @@ -19,7 +18,7 @@ import gigl.src.common.constants.local_fs as local_fs_constants import gigl.src.data_preprocessor.lib.transform.utils as transform_utils from gigl.analytics.graph_validation import BQGraphValidator -from gigl.common import Uri, UriFactory +from gigl.common import GcsUri, Uri, UriFactory from gigl.common.logger import Logger from gigl.common.metrics.decorators import flushes_metrics, profileit from gigl.common.utils import os_utils @@ -59,9 +58,6 @@ EdgeDataReference, NodeDataReference, ) -from gigl.src.data_preprocessor.lib.transform.feature_quantization import ( - apply_feature_quantization_schema, -) from gigl.src.data_preprocessor.lib.transform.transformed_features_info import ( TransformedFeaturesInfo, ) @@ -226,6 +222,17 @@ def __preprocess_single_data_reference( entity_type=entity_type, custom_identifier=custom_identifier, ) + q_spec = None + if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): + q_spec = preprocessing_spec.feature_quantization_spec + if q_spec is not None: + # Quantization changes the final TFRecord schema after TFT runs, so + # write the trainer-facing schema as its own Beam output. + transformed_features_info.transformed_features_schema_path = GcsUri.join( + transformed_features_info.transform_directory_path, + "final_metadata", + "schema.pbtxt", + ) def __get_feature_preprocessing_job_msgs( is_start: bool, @@ -279,25 +286,9 @@ def __get_feature_dimension_for_single_data_reference( return feature_dimension feature_outputs = preprocessing_spec.features_outputs - q_spec = ( - preprocessing_spec.feature_quantization_spec - if isinstance(preprocessing_spec, NodeDataPreprocessingSpec) - else None - ) - if q_spec is not None: - # Feature quantization runs after TFT, so rewrite the trainer-facing - # schema to match the final TFRecords rather than the pure TFT output. - # TODO(jchmura-sc): Consider writing a separate final schema artifact. - schema_path = transformed_features_info.transformed_features_schema_path - schema = tfdv.load_schema_text(schema_path.uri) - schema = apply_feature_quantization_schema(schema, q_spec) - with tf.io.gfile.GFile(schema_path.uri, "w") as f: - f.write(text_format.MessageToString(schema)) - if feature_outputs is not None: - q_keys = set(q_spec.feature_keys) - feature_outputs = [ - feature for feature in feature_outputs if feature not in q_keys - ] + if q_spec is not None and feature_outputs is not None: + q_keys = set(q_spec.feature_keys) + feature_outputs = [f for f in feature_outputs if f not in q_keys] # Find and save the feature dimension if there is any if feature_outputs is not None: @@ -511,7 +502,7 @@ def generate_preprocessed_metadata_pb( node_transformed_features_info.feature_quantization_metadata_path.uri ) if tf.io.gfile.exists(metadata_path): - logger.info(f"Adding feature quantization metadata from {metadata_path}") + logger.info(f"Adding feature quantization metadata: {metadata_path}") with tf.io.gfile.GFile(metadata_path) as f: metadata = json.loads(f.read()) bits = metadata["bits"] diff --git a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py index af47f3d8f..68dfcf5de 100644 --- a/gigl/src/data_preprocessor/lib/transform/feature_quantization.py +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -5,7 +5,9 @@ import numpy as np import pyarrow as pa from apache_beam.transforms.stats import ApproximateQuantiles +from google.protobuf import text_format from tensorflow_metadata.proto.v0 import schema_pb2 +from tensorflow_transform.tf_metadata.dataset_metadata import DatasetMetadata from gigl.common.logger import Logger from gigl.common.utils.feature_quantization.numpy_ops import quantize_ndarray @@ -16,7 +18,69 @@ _CentroidAcc = tuple[float, int, float, int] -def build_feature_quantization_stats( +def apply_feature_quantization_transform( + transformed_features: beam.PCollection[pa.RecordBatch], + transformed_metadata: DatasetMetadata, + analyzed_metadata: beam.PCollection[DatasetMetadata] | None, + spec: FeatureQuantizationSpec, + metadata_path: str, + schema_path: str, +): + logger.info(f"Applying Beam feature quantization with spec: {spec}") + stats = _build_feature_quantization_stats(transformed_features, spec) + _ = ( + stats + | "Build feature quantization metadata JSON" + >> beam.Map(_feature_quantization_metadata_json, spec=spec) + | "Write feature quantization metadata" + >> beam.io.WriteToText( + metadata_path, + num_shards=1, + shard_name_template="", + ) + ) + transformed_features = transformed_features | ( + "Quantize transformed feature RecordBatches" + >> beam.Map( + _quantize_record_batch, + spec=spec, + stats=beam.pvalue.AsSingleton(stats), + ) + ) + if analyzed_metadata is None: + transformed_metadata = DatasetMetadata( + _apply_feature_quantization_schema(transformed_metadata.schema, spec) + ) + schema_text = transformed_features.pipeline | ( + "Create feature quantization schema text" + >> beam.Create([text_format.MessageToString(transformed_metadata.schema)]) + ) + else: + if analyzed_metadata is None: + raise ValueError("Expected analyzed metadata for fresh TFTransform fn.") + quantized_metadata = analyzed_metadata | ( + "Apply feature quantization schema" + >> beam.Map( + lambda metadata, spec: DatasetMetadata( + _apply_feature_quantization_schema(metadata.schema, spec) + ), + spec=spec, + ) + ) + schema_text = quantized_metadata | ( + "Serialize feature quantization schema" + >> beam.Map(lambda metadata: text_format.MessageToString(metadata.schema)) + ) + transformed_metadata = beam.pvalue.AsSingleton(quantized_metadata) + _ = schema_text | "Write feature quantization schema" >> beam.io.WriteToText( + schema_path, + num_shards=1, + shard_name_template="", + ) + return transformed_features, transformed_metadata + + +def _build_feature_quantization_stats( record_batches: beam.PCollection[pa.RecordBatch], spec: FeatureQuantizationSpec, ) -> beam.PCollection[dict[str, float]]: @@ -44,7 +108,7 @@ def build_feature_quantization_stats( ) -def quantize_record_batch( +def _quantize_record_batch( batch: pa.RecordBatch, spec: FeatureQuantizationSpec, stats: dict[str, float], @@ -52,11 +116,10 @@ def quantize_record_batch( features = _build_feature_matrix(batch, spec.feature_keys) packed = quantize_ndarray(features, bits=spec.bits, stats=stats) drop_keys = set(spec.feature_keys) | {_NODE_PACKED_FEATURE_KEY} - keep_indices = [ - i for i, name in enumerate(batch.schema.names) if name not in drop_keys - ] + schema_names = batch.schema.names + keep_indices = [i for i, name in enumerate(schema_names) if name not in drop_keys] arrays = [batch.column(i) for i in keep_indices] - names = [batch.schema.names[i] for i in keep_indices] + names = [schema_names[i] for i in keep_indices] arrays.append( pa.array([[row.tobytes()] for row in packed], type=pa.list_(pa.binary())) ) @@ -64,23 +127,22 @@ def quantize_record_batch( return pa.RecordBatch.from_arrays(arrays, names=names) -def feature_quantization_metadata_json( +def _feature_quantization_metadata_json( stats: dict[str, float], spec: FeatureQuantizationSpec, ) -> str: - dim = len(spec.feature_keys) metadata = { "packed_feature_key": _NODE_PACKED_FEATURE_KEY, "dequantized_feature_keys": list(spec.feature_keys), - "dequantized_feature_dim": dim, + "dequantized_feature_dim": len(spec.feature_keys), "bits": spec.bits, + **stats, } - metadata.update(stats) logger.info(f"Writing feature quantization metadata: {metadata}") return json.dumps(metadata) -def apply_feature_quantization_schema( +def _apply_feature_quantization_schema( schema: schema_pb2.Schema, spec: FeatureQuantizationSpec ) -> schema_pb2.Schema: drop_keys = set(spec.feature_keys) | {_NODE_PACKED_FEATURE_KEY} @@ -125,8 +187,6 @@ def _build_feature_matrix(batch: pa.RecordBatch, feature_keys: list[str]) -> np. f"got {key} with shape {values.shape}." ) cols.append(values) - if not cols: - return np.empty((batch.num_rows, 0), dtype=np.float32) return np.stack(cols, axis=1) @@ -134,8 +194,7 @@ def _linear_stats_from_quantiles(quantiles: list[float]) -> dict[str, float]: if not quantiles: raise ValueError("Cannot compute quantization stats from no values.") # Store symmetric clip bounds from the approximate 99.5th abs-value percentile. - index = round(0.995 * (len(quantiles) - 1)) - clip_max = max(float(quantiles[index]), 1e-5) + clip_max = max(float(quantiles[round(0.995 * (len(quantiles) - 1))]), 1e-5) stats = {"clip_min": -clip_max, "clip_max": clip_max} logger.info(f"Computed Beam feature quantization stats: {stats}") return stats @@ -174,7 +233,6 @@ def merge_accumulators(self, accumulators: Iterable[_CentroidAcc]) -> _CentroidA def extract_output(self, accumulator: _CentroidAcc) -> dict[str, float]: neg_sum, neg_count, pos_sum, pos_count = accumulator - # Store mean values for negative and positive buckets. stats = { "neg_mean": neg_sum / neg_count if neg_count else 0.0, "pos_mean": pos_sum / pos_count if pos_count else 0.0, diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index 1633e5b8b..8e478c779 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -12,7 +12,6 @@ from tensorflow_metadata.proto.v0 import schema_pb2, statistics_pb2 from tensorflow_transform import beam as tft_beam from tensorflow_transform.tf_metadata import schema_utils -from tensorflow_transform.tf_metadata.dataset_metadata import DatasetMetadata from tfx_bsl.tfxio.record_based_tfxio import RecordBasedTFXIO from gigl.common import GcsUri, LocalUri, Uri @@ -29,10 +28,7 @@ NodeDataReference, ) from gigl.src.data_preprocessor.lib.transform.feature_quantization import ( - apply_feature_quantization_schema, - build_feature_quantization_stats, - feature_quantization_metadata_json, - quantize_record_batch, + apply_feature_quantization_transform, ) from gigl.src.data_preprocessor.lib.transform.tf_value_encoder import TFValueEncoder from gigl.src.data_preprocessor.lib.transform.transformed_features_info import ( @@ -369,51 +365,25 @@ def get_load_data_and_transform_pipeline_component( if should_use_existing_transform_fn else beam.pvalue.AsSingleton(analyzed_transform_fn[1].deferred_metadata) # type: ignore ) + q_spec = None if isinstance(preprocessing_spec, NodeDataPreprocessingSpec): q_spec = preprocessing_spec.feature_quantization_spec - else: - q_spec = None if q_spec is not None: - logger.info(f"Applying Beam feature quantization with spec: {q_spec}") - stats = build_feature_quantization_stats(transformed_features, q_spec) - _ = ( - stats - | "Build feature quantization metadata JSON" - >> beam.Map(feature_quantization_metadata_json, spec=q_spec) - | "Write feature quantization metadata" - >> beam.io.WriteToText( - transformed_features_info.feature_quantization_metadata_path.uri, - num_shards=1, - shard_name_template="", - ) - ) - transformed_features = transformed_features | ( - "Quantize transformed feature RecordBatches" - >> beam.Map( - quantize_record_batch, - spec=q_spec, - stats=beam.pvalue.AsSingleton(stats), - ) - ) if should_use_existing_transform_fn: - resolved_transformed_metadata = DatasetMetadata( - apply_feature_quantization_schema( - transformed_metadata.schema, spec=q_spec - ) - ) + analyzed_metadata = None else: - quantized_metadata = analyzed_transform_fn[1].deferred_metadata | ( - "Apply feature quantization schema" - >> beam.Map( - lambda metadata, spec: DatasetMetadata( - apply_feature_quantization_schema(metadata.schema, spec) - ), - spec=q_spec, - ) - ) - resolved_transformed_metadata = beam.pvalue.AsSingleton( - quantized_metadata + analyzed_metadata = analyzed_transform_fn[1].deferred_metadata + + transformed_features, resolved_transformed_metadata = ( + apply_feature_quantization_transform( + transformed_features=transformed_features, + transformed_metadata=transformed_metadata, + analyzed_metadata=analyzed_metadata, + spec=q_spec, + metadata_path=transformed_features_info.feature_quantization_metadata_path.uri, + schema_path=transformed_features_info.transformed_features_schema_path.uri, ) + ) transformed_features | "Write tf record files" >> BetterWriteToTFRecord( file_path_prefix=transformed_features_info.transformed_features_file_prefix.uri, From 5ef9cf476f893112a1618d6fc2d9030037f30dd6 Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 23:25:55 +0000 Subject: [PATCH 37/44] wip --- gigl/distributed/utils/neighborloader.py | 3 +-- gigl/src/data_preprocessor/data_preprocessor.py | 5 ----- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 2527564a1..cd92101c2 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -375,8 +375,7 @@ def materialize_node_store( DEFAULT_HOMOGENEOUS_NODE_TYPE ] metadata_key = ( - f"{NODE_PACKED_FEATURES_METADATA_KEY}." - f"{DEFAULT_HOMOGENEOUS_NODE_TYPE}" + f"{NODE_PACKED_FEATURES_METADATA_KEY}.{DEFAULT_HOMOGENEOUS_NODE_TYPE}" ) else: quantization_metadata = node_quantization_metadata diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index d6c25e844..045f897dc 100644 --- a/gigl/src/data_preprocessor/data_preprocessor.py +++ b/gigl/src/data_preprocessor/data_preprocessor.py @@ -742,11 +742,6 @@ def inner() -> FeatureSpecDict: labels_outputs=input_node_preprocessing_spec.labels_outputs, feature_quantization_spec=input_node_preprocessing_spec.feature_quantization_spec, ) - if enumerated_node_data_preprocessing_spec.feature_quantization_spec: - logger.info( - f"Carrying node feature quantization spec through enumeration: " - f"{enumerated_node_data_preprocessing_spec.feature_quantization_spec}" - ) enumerated_node_refs_to_preprocessing_specs[ enumerated_node_metadata.enumerated_node_data_reference ] = enumerated_node_data_preprocessing_spec From 9b26000729176b6f60e5476b7fb7670b1ae5f65b Mon Sep 17 00:00:00 2001 From: jchmura Date: Mon, 6 Jul 2026 23:28:54 +0000 Subject: [PATCH 38/44] wip --- .../pb_wrappers/preprocessed_metadata.py | 43 ++++++++----------- 1 file changed, 18 insertions(+), 25 deletions(-) diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index 6e13a9422..8b6b74193 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -2,7 +2,7 @@ from typing import Callable, cast import tensorflow as tf -from tensorflow_metadata.proto.v0.schema_pb2 import Feature, Schema +from tensorflow_metadata.proto.v0.schema_pb2 import Schema import gigl.common.utils.local_fs as LocalFsUtils from gigl.common import GcsUri, LocalUri, Uri, UriFactory @@ -224,35 +224,23 @@ def __get_feature_spec_for_feature_keys( self, feature_spec: FeatureSpecDict, feature_keys: list[str], - synthetic_feature_keys: set[str], ) -> FeatureSpecDict: """ Return feature spec for the given feature keys """ - return { - feature: tf.io.FixedLenFeature(shape=[], dtype=tf.float32) - if feature in synthetic_feature_keys - else feature_spec[feature] - for feature in feature_keys - } + return {feature: feature_spec[feature] for feature in feature_keys} def __get_schema_for_feature_keys( self, feature_schema: Schema, feature_spec: FeatureSpecDict, feature_keys: list[str], - synthetic_feature_keys: set[str], ) -> FeatureSchemaDict: """ Return feature schema for the given feature keys """ - all_features_in_feature_spec = list(feature_spec.keys()) - return { - feature: Feature(name=feature) - if feature in synthetic_feature_keys - else feature_schema.feature[all_features_in_feature_spec.index(feature)] - for feature in feature_keys - } + feature_to_schema = dict(zip(feature_spec.keys(), feature_schema.feature)) + return {feature: feature_to_schema[feature] for feature in feature_keys} def __build_feature_schema( self, @@ -276,20 +264,26 @@ def __build_feature_schema( else (None, {}) ) - for synthetic_feature_key in synthetic_feature_keys: - if synthetic_feature_key in raw_feature_spec: - raise ValueError( - f"Synthetic feature key {synthetic_feature_key} already exists " - "in raw schema; dequantized keys must not overlap normal keys." - ) - # Dequantized features are materialized later from packed uint8 tensors, # so they need logical fp32 schema entries even though TFT did not write them. + if synthetic_feature_keys: + if raw_feature_schema is None: + raise ValueError("Synthetic feature keys require a transformed schema.") + for synthetic_feature_key in synthetic_feature_keys: + if synthetic_feature_key in raw_feature_spec: + raise ValueError( + f"Synthetic feature key {synthetic_feature_key} already exists " + "in raw schema; dequantized keys must not overlap normal keys." + ) + raw_feature_spec[synthetic_feature_key] = tf.io.FixedLenFeature( + shape=[], dtype=tf.float32 + ) + raw_feature_schema.feature.add().name = synthetic_feature_key + feature_spec = ( self.__get_feature_spec_for_feature_keys( feature_spec=raw_feature_spec, feature_keys=feature_keys, - synthetic_feature_keys=synthetic_feature_keys, ) if feature_keys and raw_feature_schema else {} @@ -299,7 +293,6 @@ def __build_feature_schema( feature_schema=raw_feature_schema, feature_spec=raw_feature_spec, feature_keys=feature_keys, - synthetic_feature_keys=synthetic_feature_keys, ) if feature_keys and raw_feature_schema else {} From 641ff3b4f43cf98788d883390aaf9301114a5827 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 7 Jul 2026 15:37:20 +0000 Subject: [PATCH 39/44] Simplify pb wrapper --- .../pb_wrappers/preprocessed_metadata.py | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index 8b6b74193..0caf7ab81 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -221,9 +221,7 @@ def __post_init__(self): ) def __get_feature_spec_for_feature_keys( - self, - feature_spec: FeatureSpecDict, - feature_keys: list[str], + self, feature_spec: FeatureSpecDict, feature_keys: list[str] ) -> FeatureSpecDict: """ Return feature spec for the given feature keys @@ -239,8 +237,11 @@ def __get_schema_for_feature_keys( """ Return feature schema for the given feature keys """ - feature_to_schema = dict(zip(feature_spec.keys(), feature_schema.feature)) - return {feature: feature_to_schema[feature] for feature in feature_keys} + all_features_in_feature_spec = list(feature_spec.keys()) + return { + feature: feature_schema.feature[all_features_in_feature_spec.index(feature)] + for feature in feature_keys + } def __build_feature_schema( self, @@ -269,21 +270,21 @@ def __build_feature_schema( if synthetic_feature_keys: if raw_feature_schema is None: raise ValueError("Synthetic feature keys require a transformed schema.") - for synthetic_feature_key in synthetic_feature_keys: - if synthetic_feature_key in raw_feature_spec: - raise ValueError( - f"Synthetic feature key {synthetic_feature_key} already exists " - "in raw schema; dequantized keys must not overlap normal keys." - ) - raw_feature_spec[synthetic_feature_key] = tf.io.FixedLenFeature( + overlapping_keys = synthetic_feature_keys.intersection(raw_feature_spec) + if overlapping_keys: + raise ValueError( + "Generated dequantized feature keys must be disjoint from " + f"normal features: {sorted(overlapping_keys)}." + ) + for key in synthetic_feature_keys: + raw_feature_spec[key] = tf.io.FixedLenFeature( shape=[], dtype=tf.float32 ) - raw_feature_schema.feature.add().name = synthetic_feature_key + raw_feature_schema.feature.add().name = key feature_spec = ( self.__get_feature_spec_for_feature_keys( - feature_spec=raw_feature_spec, - feature_keys=feature_keys, + feature_spec=raw_feature_spec, feature_keys=feature_keys ) if feature_keys and raw_feature_schema else {} From 3e3d680fc4a3a5eccc56c4775cc3c5a8284da49c Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 7 Jul 2026 15:56:55 +0000 Subject: [PATCH 40/44] Clearn up dist datasett --- gigl/distributed/base_sampler.py | 2 +- gigl/distributed/dist_dataset.py | 69 ++++++++----------- .../pb_wrappers/preprocessed_metadata.py | 3 + 3 files changed, 33 insertions(+), 41 deletions(-) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index d162620bd..ff0eba570 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -107,7 +107,7 @@ def __init__(self, *args, **kwargs): data.num_partitions, data.partition_idx, data.node_quantized_features, - data.node_quantized_feature_pb, + data.node_pb, local_only=False, rpc_router=self.rpc_router, device=self.device, diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index 066aa3caf..93ddf2ced 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -159,7 +159,6 @@ def __init__( self._node_ids: Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]] = ( node_ids ) - self._node_quantized_features = node_quantized_feature_partition self._num_train = num_train self._num_val = num_val @@ -167,9 +166,11 @@ def __init__( # These fields are added so we can extract the node and edge feature dimensions and data type in the dataloader without having to lazily initialize the features. self._node_feature_info = node_feature_info + self._edge_feature_info = edge_feature_info + self._node_quantized_feature_info = node_quantized_feature_info + self._node_quantized_features = node_quantized_feature_partition self._node_quantization_metadata = node_quantization_metadata - self._edge_feature_info = edge_feature_info self._degree_tensor: Optional[ Union[torch.Tensor, dict[NodeType, torch.Tensor]] @@ -235,11 +236,12 @@ def node_quantized_features( ) -> Optional[Union[Feature, dict[NodeType, Feature]]]: return self._node_quantized_features - @property - def node_quantized_feature_pb( + @node_quantized_features.setter + def node_quantized_features( self, - ) -> Optional[Union[PartitionBook, dict[NodeType, PartitionBook]]]: - return self.node_pb if self._node_quantized_features is not None else None + new_node_quantized_features: Optional[Union[Feature, dict[NodeType, Feature]]], + ): + self._node_quantized_features = new_node_quantized_features @property def edge_features(self) -> Optional[Union[Feature, dict[EdgeType, Feature]]]: @@ -339,9 +341,6 @@ def node_feature_info( def node_quantized_feature_info( self, ) -> Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: - """ - Contains dimension and dtype for packed uint8 node features. - """ return self._node_quantized_feature_info @property @@ -797,13 +796,22 @@ def _initialize_node_quantized_features( logger.info("Found no node quantized features to initialize") return - self._node_quantized_features = _build_feature_store( - feature_data=node_quantized_features, - id2idx=node_quantized_feature_id_to_index, - dtype=torch.uint8, - ) + # GLT only exposes init_node_features for the standard node feature + # store. Calling it here would overwrite self._node_features, so build + # this sidecar Feature store directly, mirroring GLT's construction: + # https://github.com/alibaba/graphlearn-for-pytorch/blob/88ff111ac0d9e45c6c9d2d18cfc5883dca07e9f9/graphlearn_torch/python/data/dataset.py#L236 if isinstance(node_quantized_features, Mapping): + assert isinstance(node_quantized_feature_id_to_index, Mapping) + self._node_quantized_features = { + node_type: Feature( + feature_tensor=features_per_node_type, + id2index=node_quantized_feature_id_to_index[node_type], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + with_gpu=False, + dtype=torch.uint8, + ) + for node_type, features_per_node_type in node_quantized_features.items() + } self._node_quantized_feature_info = {} for node_type, features_per_node_type in node_quantized_features.items(): assert not isinstance(node_type, EdgeType) @@ -815,6 +823,13 @@ def _initialize_node_quantized_features( f"Initialized node quantized features for heterogeneous graph to dataset with node types: {node_quantized_features.keys()}" ) else: + assert not isinstance(node_quantized_feature_id_to_index, Mapping) + self._node_quantized_features = Feature( + feature_tensor=node_quantized_features, + id2index=node_quantized_feature_id_to_index, + with_gpu=False, + dtype=torch.uint8, + ) self._node_quantized_feature_info = FeatureInfo( dim=node_quantized_features.size(1), dtype=node_quantized_features.dtype, @@ -1334,32 +1349,6 @@ def _prepare_feature_data( return None, None -def _build_feature_store( - feature_data: Union[torch.Tensor, dict[_EntityType, torch.Tensor]], - id2idx: Union[TensorDataType, dict[_EntityType, TensorDataType]], - dtype: torch.dtype, -) -> Union[Feature, dict[_EntityType, Feature]]: - """Build a GLT Feature store without registering it as standard node features.""" - if isinstance(feature_data, Mapping): - assert isinstance(id2idx, Mapping) - return { - entity_key: Feature( - feature_tensor=features, - id2index=id2idx[entity_key], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. - with_gpu=False, - dtype=dtype, - ) - for entity_key, features in feature_data.items() - } - assert not isinstance(id2idx, Mapping) - return Feature( - feature_tensor=feature_data, - id2index=id2idx, - with_gpu=False, - dtype=dtype, - ) - - ## Pickling Registration # The serialization function (share_ipc) first pushes all member variable tensors # to the shared memory, and then packages all references to the tensors in one ipc diff --git a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py index 0caf7ab81..6f122f9c4 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -276,6 +276,9 @@ def __build_feature_schema( "Generated dequantized feature keys must be disjoint from " f"normal features: {sorted(overlapping_keys)}." ) + logger.info( + f"Injecting synthetic feature schema keys: {synthetic_feature_keys}" + ) for key in synthetic_feature_keys: raw_feature_spec[key] = tf.io.FixedLenFeature( shape=[], dtype=tf.float32 From e2494b41ca929700db173e3b353895d0785f84b1 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 7 Jul 2026 16:07:44 +0000 Subject: [PATCH 41/44] Remove quant feature info as its not needed --- gigl/distributed/base_dist_loader.py | 1 - gigl/distributed/dist_ablp_neighborloader.py | 2 -- gigl/distributed/distributed_neighborloader.py | 2 -- gigl/distributed/utils/neighborloader.py | 4 ---- 4 files changed, 9 deletions(-) diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index b9d2b059b..62df7bb08 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -241,7 +241,6 @@ def __init__( dataset_schema.is_homogeneous_with_labeled_edge_type ) self._node_feature_info = dataset_schema.node_feature_info - self._node_quantized_feature_info = dataset_schema.node_quantized_feature_info self._node_quantization_metadata = dataset_schema.node_quantization_metadata self._edge_feature_info = dataset_schema.edge_feature_info diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index c4ca79669..05492b3fd 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -577,7 +577,6 @@ def _setup_for_colocated( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=dataset.node_feature_info, - node_quantized_feature_info=dataset.node_quantized_feature_info, node_quantization_metadata=dataset.node_quantization_metadata, edge_feature_info=dataset.edge_feature_info, edge_dir=dataset.edge_dir, @@ -743,7 +742,6 @@ def _setup_for_graph_store( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=node_feature_info, - node_quantized_feature_info=None, node_quantization_metadata=None, edge_feature_info=edge_feature_info, edge_dir=dataset.fetch_edge_dir(), diff --git a/gigl/distributed/distributed_neighborloader.py b/gigl/distributed/distributed_neighborloader.py index a2486af17..110e76b03 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -409,7 +409,6 @@ def _setup_for_graph_store( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=node_feature_info, - node_quantized_feature_info=None, node_quantization_metadata=None, edge_feature_info=edge_feature_info, edge_dir=dataset.fetch_edge_dir(), @@ -528,7 +527,6 @@ def _setup_for_colocated( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=dataset.node_feature_info, - node_quantized_feature_info=dataset.node_quantized_feature_info, node_quantization_metadata=dataset.node_quantization_metadata, edge_feature_info=dataset.edge_feature_info, edge_dir=dataset.edge_dir, diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index cd92101c2..afe661d99 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -49,10 +49,6 @@ class DatasetSchema: edge_types: Optional[list[EdgeType]] # Node feature info. node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]] - # Packed uint8 node feature info. - node_quantized_feature_info: Optional[ - Union[FeatureInfo, dict[NodeType, FeatureInfo]] - ] # Quantization metadata for append-only packed node features. node_quantization_metadata: Optional[ Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] From 7db1896cf8aeaf1921c95252e3595b713d8d7362 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 7 Jul 2026 16:46:43 +0000 Subject: [PATCH 42/44] WIP --- gigl/distributed/base_dist_loader.py | 2 +- gigl/distributed/dist_ablp_neighborloader.py | 4 ++-- gigl/distributed/distributed_neighborloader.py | 4 ++-- gigl/distributed/utils/neighborloader.py | 4 ++-- .../utils/serialized_graph_metadata_translator.py | 14 ++++++-------- 5 files changed, 13 insertions(+), 15 deletions(-) diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index 62df7bb08..280bd010f 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -241,8 +241,8 @@ def __init__( dataset_schema.is_homogeneous_with_labeled_edge_type ) self._node_feature_info = dataset_schema.node_feature_info - self._node_quantization_metadata = dataset_schema.node_quantization_metadata 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 diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 05492b3fd..98f18f515 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -577,8 +577,8 @@ def _setup_for_colocated( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=dataset.node_feature_info, - node_quantization_metadata=dataset.node_quantization_metadata, edge_feature_info=dataset.edge_feature_info, + node_quantization_metadata=dataset.node_quantization_metadata, edge_dir=dataset.edge_dir, ), ) @@ -742,8 +742,8 @@ def _setup_for_graph_store( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=node_feature_info, - node_quantization_metadata=None, edge_feature_info=edge_feature_info, + node_quantization_metadata=None, edge_dir=dataset.fetch_edge_dir(), ), backend_key, diff --git a/gigl/distributed/distributed_neighborloader.py b/gigl/distributed/distributed_neighborloader.py index 110e76b03..82cdd32ea 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -409,8 +409,8 @@ def _setup_for_graph_store( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=node_feature_info, - node_quantization_metadata=None, edge_feature_info=edge_feature_info, + node_quantization_metadata=None, edge_dir=dataset.fetch_edge_dir(), ), backend_key, @@ -527,8 +527,8 @@ def _setup_for_colocated( is_homogeneous_with_labeled_edge_type=is_homogeneous_with_labeled_edge_type, edge_types=edge_types, node_feature_info=dataset.node_feature_info, - node_quantization_metadata=dataset.node_quantization_metadata, edge_feature_info=dataset.edge_feature_info, + node_quantization_metadata=dataset.node_quantization_metadata, edge_dir=dataset.edge_dir, ), ) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index afe661d99..a93d23347 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -49,12 +49,12 @@ class DatasetSchema: edge_types: Optional[list[EdgeType]] # Node feature info. node_feature_info: Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]] + # Edge feature info. + edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]] # Quantization metadata for append-only packed node features. node_quantization_metadata: Optional[ Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ] - # Edge feature info. - edge_feature_info: Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]] # Edge direction. edge_dir: Union[str, Literal["in", "out"]] diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index be5dcf909..a4decea2a 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -42,17 +42,15 @@ def _build_serialized_tfrecord_entity_info( packed_feature_key = quantized_metadata.packed_feature_key packed_feature_dim = metadata.packed_feature_dim - stored_keys = set(preprocessed_metadata.feature_keys) - stored_keys.update(preprocessed_metadata.label_keys) + physical_keys = set(preprocessed_metadata.feature_keys) + physical_keys.update(preprocessed_metadata.label_keys) if packed_feature_key is not None: - stored_keys.add(packed_feature_key) + physical_keys.add(packed_feature_key) # PreprocessedMetadataPbWrapper may add synthetic dequantized feature keys # to the logical feature schema. Only keep fields physically stored in # TFRecords here; dequantized features are reconstructed later from uint8. - filtered_feature_spec_dict = { - feature_key: feature_spec - for feature_key, feature_spec in feature_spec_dict.items() - if feature_key in stored_keys + feature_spec_dict = { + key: spec for key, spec in feature_spec_dict.items() if key in physical_keys } return SerializedTFRecordInfo( @@ -60,7 +58,7 @@ def _build_serialized_tfrecord_entity_info( preprocessed_metadata.tfrecord_uri_prefix ), feature_keys=list(preprocessed_metadata.feature_keys), - feature_spec=filtered_feature_spec_dict, + feature_spec=feature_spec_dict, feature_dim=preprocessed_metadata.feature_dim, entity_key=entity_key, packed_feature_key=packed_feature_key, From 4eaffe80cc64c5b6e5f97c8a0a2dd379bc838114 Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 7 Jul 2026 16:59:07 +0000 Subject: [PATCH 43/44] WIP --- gigl/distributed/utils/neighborloader.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index a93d23347..db8f86012 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -346,20 +346,15 @@ def materialize_quantized_node_features( if node_quantization_metadata is None: return data, metadata - def materialize_node_store( + def materialize( node_store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata ) -> None: dequantized = dequantize_torch_tensor(packed_features, metadata=q) x = getattr(node_store, "x", None) if x is None: node_store.x = dequantized - return - if x.size(0) != dequantized.size(0): - raise ValueError( - "Cannot materialize quantized features with " - f"{dequantized.size(0)} rows into existing x with {x.size(0)} rows." - ) - node_store.x = torch.cat([x, dequantized], dim=1) + else: + node_store.x = torch.cat([x, dequantized], dim=1) if isinstance(data, Data): if isinstance(node_quantization_metadata, dict): @@ -381,7 +376,7 @@ def materialize_node_store( raise ValueError( f"Missing packed quantized node features in metadata key {metadata_key}." ) - materialize_node_store(data, packed_features, quantization_metadata) + materialize(data, packed_features, quantization_metadata) return data, metadata assert isinstance(node_quantization_metadata, dict), ( @@ -395,7 +390,7 @@ def materialize_node_store( packed_features = metadata.pop(metadata_key, None) if packed_features is None: continue - materialize_node_store(data[node_type], packed_features, quantization_metadata) + materialize(data[node_type], packed_features, quantization_metadata) return data, metadata From 75dfd0d2365680353a1e9131a7eb89adf9089cea Mon Sep 17 00:00:00 2001 From: jchmura Date: Tue, 7 Jul 2026 17:04:09 +0000 Subject: [PATCH 44/44] ty check --- gigl/distributed/utils/neighborloader.py | 8 ++++---- gigl/src/data_preprocessor/lib/transform/utils.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index db8f86012..13ed280a8 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -347,14 +347,14 @@ def materialize_quantized_node_features( return data, metadata def materialize( - node_store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata + store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata ) -> None: dequantized = dequantize_torch_tensor(packed_features, metadata=q) - x = getattr(node_store, "x", None) + x = getattr(store, "x", None) if x is None: - node_store.x = dequantized + store.x = dequantized else: - node_store.x = torch.cat([x, dequantized], dim=1) + store.x = torch.cat([x, dequantized], dim=1) if isinstance(data, Data): if isinstance(node_quantization_metadata, dict): diff --git a/gigl/src/data_preprocessor/lib/transform/utils.py b/gigl/src/data_preprocessor/lib/transform/utils.py index 8e478c779..3be7bc94d 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -372,7 +372,7 @@ def get_load_data_and_transform_pipeline_component( if should_use_existing_transform_fn: analyzed_metadata = None else: - analyzed_metadata = analyzed_transform_fn[1].deferred_metadata + analyzed_metadata = analyzed_transform_fn[1].deferred_metadata # type: ignore transformed_features, resolved_transformed_metadata = ( apply_feature_quantization_transform(