diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 0ac83f0e7..498309ce3 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 name to load for the current node entity. + packed_feature_key: Optional[str] = None + # Number of packed uint8 columns for the current node entity. + packed_feature_dim: int = 0 # Name of the label columns for the current entity, defaults to an empty list. label_keys: Sequence[str] = field(default_factory=list) # The regex pattern to match the TFRecord files at the specified prefix @@ -367,10 +372,11 @@ def load_as_torch_tensors( serialized_tf_record_info (SerializedTFRecordInfo): Information for how TFRecord files are serialized on disk. tf_dataset_options (TFDatasetOptions): The options to use when building the dataset. Returns: - LoadedEntityTensors: The (id_tensor, feature_tensor, label_tensor) for the loaded entities. + LoadedEntityTensors: The (id_tensor, feature_tensor, quantized_feature_tensor, label_tensor) for the loaded entities. """ entity_key = serialized_tf_record_info.entity_key feature_keys = serialized_tf_record_info.feature_keys + packed_feature_key = serialized_tf_record_info.packed_feature_key label_keys = serialized_tf_record_info.label_keys # We make a deep copy of the feature spec dict so that future modifications don't redirect to the input @@ -392,6 +398,16 @@ def load_as_torch_tensors( feature_spec_dict[entity_key] = tf.io.FixedLenFeature( shape=[], dtype=tf.int64 ) + if ( + packed_feature_key is not None + and packed_feature_key not in feature_spec_dict + ): + logger.info( + f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" + ) + feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature( + shape=[], dtype=tf.string + ) else: id_concat_axis = 1 proccess_id_tensor = lambda t: tf.stack( @@ -435,13 +451,24 @@ def load_as_torch_tensors( else: empty_feature = None + if packed_feature_key is not None: + empty_quantized_feature = torch.empty( + (0, serialized_tf_record_info.packed_feature_dim), + dtype=torch.uint8, + ) + else: + empty_quantized_feature = None + if label_keys: empty_label = torch.empty(0, len(label_keys)) else: empty_label = None return LoadedEntityTensors( - ids=empty_entity, features=empty_feature, labels=empty_label + ids=empty_entity, + features=empty_feature, + quantized_features=empty_quantized_feature, + labels=empty_label, ) dataset = TFRecordDataLoader._build_dataset_for_uris( @@ -454,6 +481,7 @@ def load_as_torch_tensors( num_entities_processed = 0 id_tensors: list[torch.Tensor] = [] feature_tensors: list[torch.Tensor] = [] + quantized_feature_tensors: list[torch.Tensor] = [] label_tensors: list[torch.Tensor] = [] for idx, batch in enumerate(dataset): id_tensors.append(proccess_id_tensor(batch)) @@ -465,6 +493,15 @@ def load_as_torch_tensors( feature_tensors.append(feature_tensor) if label_tensor is not None: label_tensors.append(label_tensor) + if packed_feature_key is not None: + quantized_feature_tensor = tf.io.decode_raw( + batch[packed_feature_key], tf.uint8 + ) + quantized_feature_tensor = tf.reshape( + quantized_feature_tensor, + [-1, serialized_tf_record_info.packed_feature_dim], + ) + quantized_feature_tensors.append(quantized_feature_tensor) num_entities_processed += ( id_tensors[-1].shape[0] if entity_type == FeatureTypes.NODE @@ -483,11 +520,16 @@ def load_as_torch_tensors( tf.concat(id_tensors, axis=id_concat_axis) ) output_feature_tensor: Optional[torch.Tensor] = None + output_quantized_feature_tensor: Optional[torch.Tensor] = None output_label_tensor: Optional[torch.Tensor] = None if feature_tensors: output_feature_tensor = _tf_tensor_to_torch_tensor( tf.concat(feature_tensors, axis=0) ) + if quantized_feature_tensors: + output_quantized_feature_tensor = _tf_tensor_to_torch_tensor( + tf.concat(quantized_feature_tensors, axis=0) + ).to(torch.uint8) if label_tensors: output_label_tensor = _tf_tensor_to_torch_tensor( tf.concat(label_tensors, axis=0) @@ -503,5 +545,8 @@ def load_as_torch_tensors( f"Converted {num_entities_processed:,} {entity_type.name} to torch tensors in {end - start:.2f} seconds" ) return LoadedEntityTensors( - ids=id_tensor, features=output_feature_tensor, labels=output_label_tensor + ids=id_tensor, + features=output_feature_tensor, + quantized_features=output_quantized_feature_tensor, + labels=output_label_tensor, ) diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index 2ba43ce10..7b3545e33 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -18,6 +18,7 @@ from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, + FeatureQuantizationMetadata, LoadedGraphTensors, ) from gigl.utils.share_memory import share_memory @@ -26,6 +27,7 @@ _ID_FMT = "{entity}_ids" _FEATURE_FMT = "{entity}_features" +_PACKED_FEATURE_FMT = "{entity}_packed_features" _LABEL_FMT = "{entity}_labels" _EDGE_WEIGHTS_KEY = "edge_weights" _NODE_KEY = "node" @@ -113,6 +115,10 @@ class SerializedGraphMetadata: negative_label_entity_info: Optional[ Union[SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]] ] = None + # Optional node quantization metadata. + node_quantization_metadata: Optional[ + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] + ] = None def _data_loading_process( @@ -178,6 +184,7 @@ def _data_loading_process( ids: dict[Union[NodeType, EdgeType], torch.Tensor] = {} features: dict[Union[NodeType, EdgeType], torch.Tensor] = {} + quantized_features: dict[Union[NodeType, EdgeType], torch.Tensor] = {} labels: dict[Union[NodeType, EdgeType], torch.Tensor] = {} weights: dict[Union[NodeType, EdgeType], torch.Tensor] = {} for ( @@ -192,12 +199,20 @@ def _data_loading_process( raise NotImplementedError( "Label keys are not supported for edge entities" ) + if ( + serialized_entity_tf_record_info.packed_feature_key is not None + and not serialized_entity_tf_record_info.is_node_entity + ): + raise NotImplementedError( + "Packed feature keys are not supported for edge entities" + ) loaded_entity = tf_record_dataloader.load_as_torch_tensors( serialized_tf_record_info=serialized_entity_tf_record_info, tf_dataset_options=tf_dataset_options, ) entity_ids = loaded_entity.ids entity_features = loaded_entity.features + entity_quantized_features = loaded_entity.quantized_features entity_labels = loaded_entity.labels ids[graph_type] = entity_ids logger.info( @@ -213,6 +228,16 @@ def _data_loading_process( f"Rank {rank} did not detect {entity_type} features for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}" ) + if entity_quantized_features is not None: + quantized_features[graph_type] = entity_quantized_features + logger.info( + f"Rank {rank} finished loading {entity_type} quantized features of shape {entity_quantized_features.shape} for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}" + ) + else: + logger.info( + f"Rank {rank} did not detect {entity_type} quantized features for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}" + ) + if entity_labels is not None: labels[graph_type] = entity_labels logger.info( @@ -291,6 +316,12 @@ def _data_loading_process( share_memory(features) # We convert the features back to homogeneous from the default heterogeneous setup if our provided input was homogeneous + if quantized_features: + logger.info( + f"Rank {rank} is attempting to share {entity_type} quantized feature memory for tfrecord directories: {all_tf_record_uris}" + ) + share_memory(quantized_features) + if labels: logger.info( f"Rank {rank} is attempting to share {entity_type} label memory for tfrecord directories: {all_tf_record_uris}" @@ -310,6 +341,12 @@ def _data_loading_process( output_dict[_FEATURE_FMT.format(entity=entity_type)] = ( list(features.values())[0] if is_input_homogeneous else features ) + if quantized_features: + output_dict[_PACKED_FEATURE_FMT.format(entity=entity_type)] = ( + list(quantized_features.values())[0] + if is_input_homogeneous + else quantized_features + ) if labels: output_dict[_LABEL_FMT.format(entity=entity_type)] = ( list(labels.values())[0] if is_input_homogeneous else labels @@ -480,6 +517,9 @@ def load_torch_tensors_from_tf_record( node_ids = node_output_dict[_ID_FMT.format(entity=_NODE_KEY)] node_features = node_output_dict.get(_FEATURE_FMT.format(entity=_NODE_KEY), None) + node_quantized_features = node_output_dict.get( + _PACKED_FEATURE_FMT.format(entity=_NODE_KEY), None + ) node_labels = node_output_dict.get(_LABEL_FMT.format(entity=_NODE_KEY), None) edge_index = edge_output_dict[_ID_FMT.format(entity=_EDGE_KEY)] @@ -507,6 +547,7 @@ def load_torch_tensors_from_tf_record( return LoadedGraphTensors( node_ids=node_ids, node_features=node_features, + node_quantized_features=node_quantized_features, node_labels=node_labels, edge_index=edge_index, edge_features=edge_features, diff --git a/gigl/common/utils/feature_quantization/__init__.py b/gigl/common/utils/feature_quantization/__init__.py new file mode 100644 index 000000000..e69de29bb 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..bb8e30f28 --- /dev/null +++ b/gigl/common/utils/feature_quantization/numpy_ops.py @@ -0,0 +1,38 @@ +from collections.abc import Mapping + +import numpy as np + + +def quantize_ndarray( + features: np.ndarray, *, bits: int, stats: Mapping[str, float] +) -> np.ndarray: + """Quantize a 2D float array into packed uint8 codes.""" + if bits not in (1, 2, 4, 8): + raise ValueError(f"bits must be one of 1, 2, 4, or 8, got {bits}.") + if features.ndim != 2: + raise ValueError(f"Expected a 2D feature array, got shape {features.shape}.") + if bits == 1: + # 1-bit quantization keeps only sign; values restore from neg/pos means. + codes = (features > 0).astype(np.uint8) + else: + # Linearly map clipped values into integer buckets. + levels = (1 << bits) - 1 + lo, hi = stats["clip_min"], stats["clip_max"] + clipped = np.clip(features, lo, hi) + scaled = (clipped - lo) / (hi - lo) + codes = np.rint(scaled * levels).astype(np.uint8) + return pack_codes(codes, bits) + + +def pack_codes(codes: np.ndarray, bits: int) -> np.ndarray: + """Pack low-bit feature codes high-bits-first along the final dimension.""" + per_byte = 8 // bits + pad = (-codes.shape[-1]) % per_byte + if pad: + # Pad only the feature dimension of this 2D [row, feature] array. + codes = np.pad(codes, ((0, 0), (0, pad)), constant_values=0) + # Group the padded feature dimension into chunks that each form one byte. + codes = codes.reshape(codes.shape[0], -1, per_byte).astype(np.uint16) + shifts = bits * np.arange(per_byte - 1, -1, -1, dtype=np.uint16) + weights = (1 << shifts).astype(np.uint16) + return np.sum(codes * weights, axis=-1).astype(np.uint8) diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py new file mode 100644 index 000000000..b3a492d00 --- /dev/null +++ b/gigl/common/utils/feature_quantization/torch_ops.py @@ -0,0 +1,42 @@ +import torch + +from gigl.types.graph import FeatureQuantizationMetadata + + +def dequantize_torch_tensor( + packed_features: torch.Tensor, + metadata: FeatureQuantizationMetadata, +) -> torch.Tensor: + """Reconstruct approximate float features from packed uint8 codes.""" + q = metadata + + if packed_features.size(-1) != q.packed_feature_dim: + raise ValueError( + f"Expected packed feature dim {q.packed_feature_dim} for " + f"{q.dequantized_feature_dim} {q.bits}-bit features, got " + f"{packed_features.size(-1)}." + ) + + codes = _unpack_torch_tensor( + packed_features, dim=q.dequantized_feature_dim, bits=q.bits + ).float() + if q.bits == 1: + if q.neg_mean is None or q.pos_mean is None: + raise ValueError("1-bit dequantization requires pos_mean/neg_mean") + return torch.where(codes.bool(), q.pos_mean, q.neg_mean) + else: + if q.clip_min is None or q.clip_max is None: + raise ValueError(f"{q.bits}-bit dequantization requires clip_min/clip_max") + levels = (1 << q.bits) - 1 + return q.clip_min + (codes / levels) * (q.clip_max - q.clip_min) + + +def _unpack_torch_tensor( + packed_features: torch.Tensor, *, dim: int, bits: int +) -> torch.Tensor: + per_byte = 8 // bits + mask = (1 << bits) - 1 + # Extract high-bits-first codes from each packed byte. + shifts = bits * torch.arange(per_byte - 1, -1, -1, device=packed_features.device) + codes = (packed_features.unsqueeze(-1).to(torch.int16) >> shifts).bitwise_and(mask) + return codes.reshape(*packed_features.shape[:-1], -1)[..., :dim].to(torch.uint8) diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index 02d7c69e1..280bd010f 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -242,6 +242,7 @@ def __init__( ) self._node_feature_info = dataset_schema.node_feature_info self._edge_feature_info = dataset_schema.edge_feature_info + self._node_quantization_metadata = dataset_schema.node_quantization_metadata self._sampler_options = sampler_options self._non_blocking_transfers = non_blocking_transfers diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 986ba5d58..ff0eba570 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -17,6 +17,7 @@ from gigl.distributed.sampler import ( NEGATIVE_LABEL_METADATA_KEY, + NODE_PACKED_FEATURES_METADATA_KEY, POSITIVE_LABEL_METADATA_KEY, ABLPNodeSamplerInput, ) @@ -91,6 +92,27 @@ 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 + ): + # 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, + data.node_quantized_features, + data.node_pb, + local_only=False, + rpc_router=self.rpc_router, + device=self.device, + ) + def _prepare_sample_loop_inputs( self, inputs: NodeSamplerInput, @@ -322,6 +344,30 @@ 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_PACKED_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_PACKED_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_PACKED_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..1a5f859b5 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, @@ -217,7 +222,12 @@ 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, diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 50f42f5a9..98f18f515 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -33,6 +33,7 @@ extract_edge_type_metadata, extract_metadata, labeled_to_homogeneous, + materialize_quantized_node_features, set_missing_features, shard_nodes_by_process, strip_label_edges, @@ -577,6 +578,7 @@ def _setup_for_colocated( edge_types=edge_types, node_feature_info=dataset.node_feature_info, edge_feature_info=dataset.edge_feature_info, + node_quantization_metadata=dataset.node_quantization_metadata, edge_dir=dataset.edge_dir, ), ) @@ -741,6 +743,7 @@ def _setup_for_graph_store( edge_types=edge_types, node_feature_info=node_feature_info, edge_feature_info=edge_feature_info, + node_quantization_metadata=None, edge_dir=dataset.fetch_edge_dir(), ), backend_key, @@ -876,6 +879,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 = materialize_quantized_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..93ddf2ced 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -23,6 +23,7 @@ from gigl.types.graph import ( FeatureInfo, FeaturePartitionData, + FeatureQuantizationMetadata, GraphPartitionData, PartitionOutput, ) @@ -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,15 @@ 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[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] + ] = None, edge_feature_info: Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] ] = None, @@ -94,9 +107,12 @@ 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 + 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 +125,7 @@ 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. 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. @@ -151,6 +168,10 @@ def __init__( 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._degree_tensor: Optional[ Union[torch.Tensor, dict[NodeType, torch.Tensor]] ] = degree_tensor @@ -209,6 +230,19 @@ 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 + + @node_quantized_features.setter + def node_quantized_features( + self, + 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]]]: """ @@ -303,6 +337,20 @@ def node_feature_info( """ return self._node_feature_info + @property + def node_quantized_feature_info( + self, + ) -> Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: + return self._node_quantized_feature_info + + @property + def node_quantization_metadata( + self, + ) -> Optional[ + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] + ]: + return self._node_quantization_metadata + @property def edge_feature_info( self, @@ -723,6 +771,73 @@ 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 + + # 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) + self._node_quantized_feature_info[node_type] = FeatureInfo( + 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()}" + ) + 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, + ) + logger.info( + "Initialized node quantized features for homogeneous graph to dataset" + ) + def _initialize_node_labels( self, node_partition_book: Union[PartitionBook, dict[NodeType, PartitionBook]], @@ -897,6 +1012,13 @@ 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, + ) + 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 +1057,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 +1069,13 @@ 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[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] + ], Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]], Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]], Optional[int], @@ -959,6 +1089,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 +1101,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 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 +1123,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 +1135,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 @@ -1234,6 +1370,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 +1396,15 @@ 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[ + FeatureQuantizationMetadata, + dict[NodeType, FeatureQuantizationMetadata], + ] + ], # 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..04de8ce72 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,29 @@ 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 +1166,14 @@ 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 +1190,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 +1560,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 +1578,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 +1612,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 +1646,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 +1663,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 +1673,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 +1894,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 +1934,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..82cdd32ea 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -29,6 +29,7 @@ SamplingClusterSetup, extract_metadata, labeled_to_homogeneous, + materialize_quantized_node_features, set_missing_features, shard_nodes_by_process, strip_label_edges, @@ -409,6 +410,7 @@ def _setup_for_graph_store( edge_types=edge_types, node_feature_info=node_feature_info, edge_feature_info=edge_feature_info, + node_quantization_metadata=None, edge_dir=dataset.fetch_edge_dir(), ), backend_key, @@ -526,6 +528,7 @@ def _setup_for_colocated( edge_types=edge_types, node_feature_info=dataset.node_feature_info, edge_feature_info=dataset.edge_feature_info, + node_quantization_metadata=dataset.node_quantization_metadata, edge_dir=dataset.edge_dir, ), ) @@ -550,6 +553,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 = materialize_quantized_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/sampler.py b/gigl/distributed/sampler.py index e99dd65dc..7789c6731 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_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 876582208..13ed280a8 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 @@ -13,7 +13,14 @@ 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.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor +from gigl.distributed.sampler import NODE_PACKED_FEATURES_METADATA_KEY +from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_NODE_TYPE, + FeatureInfo, + FeatureQuantizationMetadata, + is_label_edge_type, +) logger = Logger() @@ -44,6 +51,10 @@ class DatasetSchema: 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 direction. edge_dir: Union[str, Literal["in", "out"]] @@ -324,6 +335,65 @@ def set_missing_features( return data +def materialize_quantized_node_features( + data: _GraphType, + metadata: dict[str, torch.Tensor], + node_quantization_metadata: Optional[ + Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] + ], +) -> tuple[_GraphType, dict[str, torch.Tensor]]: + """Materialize packed quantized node features into PyG node feature tensors.""" + if node_quantization_metadata is None: + return data, metadata + + def materialize( + store, packed_features: torch.Tensor, q: FeatureQuantizationMetadata + ) -> None: + dequantized = dequantize_torch_tensor(packed_features, metadata=q) + x = getattr(store, "x", None) + if x is None: + store.x = dequantized + else: + store.x = torch.cat([x, dequantized], dim=1) + + if isinstance(data, Data): + if isinstance(node_quantization_metadata, dict): + homogeneous_quantization_metadata = cast( + dict[NodeType, FeatureQuantizationMetadata], + node_quantization_metadata, + ) + quantization_metadata = homogeneous_quantization_metadata[ + DEFAULT_HOMOGENEOUS_NODE_TYPE + ] + metadata_key = ( + f"{NODE_PACKED_FEATURES_METADATA_KEY}.{DEFAULT_HOMOGENEOUS_NODE_TYPE}" + ) + else: + quantization_metadata = node_quantization_metadata + metadata_key = NODE_PACKED_FEATURES_METADATA_KEY + packed_features = metadata.pop(metadata_key, None) + if packed_features is None: + raise ValueError( + f"Missing packed quantized node features in metadata key {metadata_key}." + ) + materialize(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_PACKED_FEATURES_METADATA_KEY}.{node_type}" + packed_features = metadata.pop(metadata_key, None) + if packed_features is None: + continue + materialize(data[node_type], packed_features, quantization_metadata) + 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..d9cdaabe8 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -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 FeatureQuantizationMetadata, to_homogeneous from snapchat.research.gbml.preprocessed_metadata_pb2 import PreprocessedMetadata @@ -33,6 +33,26 @@ def _build_serialized_tfrecord_entity_info( Returns: SerializedTFRecordInfo: Stored metadata for current entity """ + 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) + packed_feature_key = quantized_metadata.packed_feature_key + packed_feature_dim = metadata.packed_feature_dim + + physical_keys = set(preprocessed_metadata.feature_keys) + physical_keys.update(preprocessed_metadata.label_keys) + if packed_feature_key is not None: + 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. + feature_spec_dict = { + key: spec for key, spec in feature_spec_dict.items() if key in physical_keys + } + return SerializedTFRecordInfo( tfrecord_uri_prefix=UriFactory.create_uri( preprocessed_metadata.tfrecord_uri_prefix @@ -41,11 +61,41 @@ def _build_serialized_tfrecord_entity_info( feature_spec=feature_spec_dict, feature_dim=preprocessed_metadata.feature_dim, entity_key=entity_key, + 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, ) +def _build_feature_quantization_metadata( + quantized_metadata: PreprocessedMetadata.FeatureQuantizationMetadata, +) -> FeatureQuantizationMetadata: + state = quantized_metadata.WhichOneof("state") + expected_state = "centroid" if quantized_metadata.bits == 1 else "linear" + if state != expected_state: + raise ValueError( + f"Expected {expected_state} quantization state for {quantized_metadata.bits}-bit features." + ) + + neg_mean = pos_mean = clip_min = clip_max = None + if quantized_metadata.bits == 1: + neg_mean = quantized_metadata.centroid.neg_mean + pos_mean = quantized_metadata.centroid.pos_mean + else: + clip_min = quantized_metadata.linear.clip_min + clip_max = quantized_metadata.linear.clip_max + return FeatureQuantizationMetadata( + bits=quantized_metadata.bits, + dequantized_feature_dim=quantized_metadata.dequantized_feature_dim, + dequantized_feature_keys=tuple(quantized_metadata.dequantized_feature_keys), + clip_min=clip_min, + clip_max=clip_max, + neg_mean=neg_mean, + pos_mean=pos_mean, + ) + + def convert_pb_to_serialized_graph_metadata( preprocessed_metadata_pb_wrapper: PreprocessedMetadataPbWrapper, graph_metadata_pb_wrapper: GraphMetadataPbWrapper, @@ -65,6 +115,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, FeatureQuantizationMetadata] = {} preprocessed_metadata_pb = preprocessed_metadata_pb_wrapper.preprocessed_metadata_pb @@ -92,6 +143,12 @@ def convert_pb_to_serialized_graph_metadata( entity_key=node_key, tfrecord_uri_pattern=tfrecord_uri_pattern, ) + 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 = ( @@ -159,6 +216,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 +230,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..6f122f9c4 100644 --- a/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py +++ b/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py @@ -1,6 +1,7 @@ from dataclasses import dataclass, field from typing import Callable, cast +import tensorflow as tf from tensorflow_metadata.proto.v0.schema_pb2 import Schema import gigl.common.utils.local_fs as LocalFsUtils @@ -65,9 +66,24 @@ 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 @@ -80,7 +96,8 @@ 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), ) condensed_node_type_to_feature_schema_map[ CondensedNodeType(condensed_node_type) @@ -231,6 +248,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 @@ -240,12 +258,33 @@ 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, {}) ) + # 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.") + 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)}." + ) + 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 + ) + 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 diff --git a/gigl/src/data_preprocessor/data_preprocessor.py b/gigl/src/data_preprocessor/data_preprocessor.py index 543954753..045f897dc 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 @@ -17,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 @@ -221,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, @@ -273,18 +285,23 @@ def __get_feature_dimension_for_single_data_reference( feature_dimension += feature_shape[0] return feature_dimension + feature_outputs = preprocessing_spec.features_outputs + 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 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(feature_transform_pipeline_result, DataflowPipelineResult): @@ -481,6 +498,31 @@ 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, ) + 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: {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( + packed_feature_key=metadata["packed_feature_key"], + dequantized_feature_keys=metadata["dequantized_feature_keys"], + dequantized_feature_dim=metadata["dequantized_feature_dim"], + bits=bits, + ) + if bits == 1: + centroid = quantized_feature_metadata_pb.centroid + centroid.neg_mean = metadata["neg_mean"] + centroid.pos_mean = metadata["pos_mean"] + else: + linear = quantized_feature_metadata_pb.linear + linear.clip_min = metadata["clip_min"] + linear.clip_max = metadata["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) @@ -698,6 +740,7 @@ 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, ) enumerated_node_refs_to_preprocessing_specs[ enumerated_node_metadata.enumerated_node_data_reference 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..68dfcf5de --- /dev/null +++ b/gigl/src/data_preprocessor/lib/transform/feature_quantization.py @@ -0,0 +1,241 @@ +import json +from typing import Iterable + +import apache_beam as beam +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 +from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec + +logger = Logger() +_NODE_PACKED_FEATURE_KEY = "node_packed_features" +_CentroidAcc = tuple[float, int, float, int] + + +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]]: + 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 ( + record_batches + | "Compute centroid quantization stats" + >> beam.CombineGlobally(_CentroidStatsFn(spec.feature_keys)) + ) + return ( + 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, input_batched=True) + | "Build linear quantization stats" >> beam.Map(_linear_stats_from_quantiles) + ) + + +def _quantize_record_batch( + batch: pa.RecordBatch, + spec: FeatureQuantizationSpec, + stats: dict[str, float], +) -> 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_PACKED_FEATURE_KEY} + 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 = [schema_names[i] for i in keep_indices] + arrays.append( + pa.array([[row.tobytes()] for row in packed], type=pa.list_(pa.binary())) + ) + names.append(_NODE_PACKED_FEATURE_KEY) + return pa.RecordBatch.from_arrays(arrays, names=names) + + +def _feature_quantization_metadata_json( + stats: dict[str, float], + spec: FeatureQuantizationSpec, +) -> str: + metadata = { + "packed_feature_key": _NODE_PACKED_FEATURE_KEY, + "dequantized_feature_keys": list(spec.feature_keys), + "dequantized_feature_dim": len(spec.feature_keys), + "bits": spec.bits, + **stats, + } + logger.info(f"Writing feature quantization metadata: {metadata}") + return json.dumps(metadata) + + +def _apply_feature_quantization_schema( + schema: schema_pb2.Schema, spec: FeatureQuantizationSpec +) -> schema_pb2.Schema: + drop_keys = set(spec.feature_keys) | {_NODE_PACKED_FEATURE_KEY} + quantized_schema = schema_pb2.Schema() + quantized_schema.CopyFrom(schema) + kept_features = [ + feature for feature in quantized_schema.feature if feature.name not in drop_keys + ] + del quantized_schema.feature[:] + quantized_schema.feature.extend(kept_features) + 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_PACKED_FEATURE_KEY}." + ) + return quantized_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 {key} with shape {values.shape}." + ) + cols.append(values) + return np.stack(cols, axis=1) + + +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. + 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 + + +class _CentroidStatsFn(beam.CombineFn): + 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: _CentroidAcc, batch: pa.RecordBatch + ) -> _CentroidAcc: + neg_sum, neg_count, pos_sum, pos_count = accumulator + 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[_CentroidAcc]) -> _CentroidAcc: + neg_sum = neg_count = pos_sum = pos_count = 0 + 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: _CentroidAcc) -> dict[str, float]: + neg_sum, neg_count, pos_sum, pos_count = accumulator + 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/transformed_features_info.py b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py index c7c0669eb..903a02057 100644 --- a/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py +++ b/gigl/src/data_preprocessor/lib/transform/transformed_features_info.py @@ -22,6 +22,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 @@ -92,6 +93,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..3be7bc94d 100644 --- a/gigl/src/data_preprocessor/lib/transform/utils.py +++ b/gigl/src/data_preprocessor/lib/transform/utils.py @@ -27,6 +27,9 @@ EdgeDataReference, NodeDataReference, ) +from gigl.src.data_preprocessor.lib.transform.feature_quantization import ( + 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 ( TransformedFeaturesInfo, @@ -362,6 +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 + if q_spec is not None: + if should_use_existing_transform_fn: + analyzed_metadata = None + else: + analyzed_metadata = analyzed_transform_fn[1].deferred_metadata # type: ignore + + 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, diff --git a/gigl/src/data_preprocessor/lib/types.py b/gigl/src/data_preprocessor/lib/types.py index 491cc9486..a3308b00d 100644 --- a/gigl/src/data_preprocessor/lib/types.py +++ b/gigl/src/data_preprocessor/lib/types.py @@ -48,6 +48,11 @@ class NodeOutputIdentifier(str): """ +class FeatureQuantizationSpec(NamedTuple): + 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 +77,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_spec: Optional[FeatureQuantizationSpec] = None def __repr__(self) -> str: return f"""NodeDataPreprocessingSpec( @@ -80,7 +86,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_spec={self.feature_quantization_spec}) """ diff --git a/gigl/types/graph.py b/gigl/types/graph.py index e7323693e..8dce27c69 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,34 @@ class FeatureInfo: dtype: torch.dtype +@dataclass(frozen=True) +class FeatureQuantizationMetadata: + """Metadata needed to unpack/dequantize append-only packed features.""" + + bits: int + dequantized_feature_dim: int + dequantized_feature_keys: tuple[str, ...] = () + clip_min: Optional[float] = None + clip_max: Optional[float] = None + 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, edge_dir: Literal["in", "out"], @@ -151,6 +185,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..c45be4967 100644 --- a/proto/snapchat/research/gbml/preprocessed_metadata.proto +++ b/proto/snapchat/research/gbml/preprocessed_metadata.proto @@ -3,6 +3,36 @@ syntax = "proto3"; package snapchat.research.gbml; message PreprocessedMetadata{ + message LinearQuantizationState{ + // Lower clipping bound; dequantized value for linear code 0. + float clip_min = 1; + // Upper clipping bound; dequantized value for max linear code. + float clip_max = 2; + } + + message CentroidQuantizationState{ + // 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{ + // Field in output TFRecords that stores packed uint8 features. + 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. + 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; + } + } + // Houses metadata about node TFTransform output from DataPreprocessor. message NodeMetadataOutput{ // The field in output TFRecords which references the node identifier. @@ -23,6 +53,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. + 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 8f2f818f3..9f377e2b1 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\"\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') _PREPROCESSEDMETADATA = DESCRIPTOR.message_types_by_name['PreprocessedMetadata'] +_PREPROCESSEDMETADATA_LINEARQUANTIZATIONSTATE = _PREPROCESSEDMETADATA.nested_types_by_name['LinearQuantizationState'] +_PREPROCESSEDMETADATA_CENTROIDQUANTIZATIONSTATE = _PREPROCESSEDMETADATA.nested_types_by_name['CentroidQuantizationState'] +_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'] @@ -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) + }) + , + + '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.FeatureQuantizationMetadata) + }) + , + '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.FeatureQuantizationMetadata) _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=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=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 2271a7fec..5f35dc8b4 100644 --- a/snapchat/research/gbml/preprocessed_metadata_pb2.pyi +++ b/snapchat/research/gbml/preprocessed_metadata_pb2.pyi @@ -20,6 +20,76 @@ 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 + """Lower clipping bound; dequantized value for linear code 0.""" + clip_max: builtins.float + """Upper clipping bound; dequantized value for max linear code.""" + 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 + + 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, + *, + neg_mean: builtins.float = ..., + pos_mean: builtins.float = ..., + ) -> 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 + + 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 + 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]: + """Model-visible feature names produced after dequantizing and appending.""" + 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_key: builtins.str = ..., + dequantized_feature_keys: collections.abc.Iterable[builtins.str] | None = ..., + 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_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): """Houses metadata about node TFTransform output from DataPreprocessor.""" @@ -34,6 +104,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 +125,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.FeatureQuantizationMetadata: + """Optional append-only quantized node feature sidecar metadata.""" def __init__( self, *, @@ -66,9 +140,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.FeatureQuantizationMetadata | 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): diff --git a/tests/test_assets/distributed/run_distributed_partitioner.py b/tests/test_assets/distributed/run_distributed_partitioner.py index 3bbc3406c..f557e283b 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_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..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, 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, 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, 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