Feature Quantization [5/5]: Distributed Integration

gigl/distributed/base_dist_loader.py

@@ -245,6 +245,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

gigl/distributed/base_sampler.py

@@ -1,7 +1,7 @@
 import asyncio
 import traceback
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Optional, Union
 
 import torch
@@ -20,6 +20,7 @@
 from gigl.common.logger import Logger
 from gigl.distributed.sampler import (
     NEGATIVE_LABEL_METADATA_KEY,
+    NODE_PACKED_FEATURES_METADATA_KEY,
     POSITIVE_LABEL_METADATA_KEY,
     ABLPNodeSamplerInput,
 )
@@ -110,9 +111,28 @@ def __init__(self, *args, **kwargs) -> None:
         which GLT's event loop would swallow the same way it swallows the
         original sampling exception.
         """
+        data = kwargs.get("data")
         super().__init__(*args, **kwargs)
         self._sampling_error_sent: bool = False
 
+        self.dist_node_quantized_feature: Optional[DistFeature] = 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,
@@ -357,18 +377,63 @@ async def _collate_fn(
                             ]
             if self.dist_node_feature is not None:
                 if self.use_all2all:
-                    sorted_ntype = sorted(self.dist_node_feature.feature_pb.keys())
+                    sorted_ntype = sorted(self.dist_node_feature.local_feature.keys())
+                    # GLT get_all2all() iterates every type in output.node, not
+                    # just sorted_ntype. feature_pb contains partition books for
+                    # every node type, while local_feature contains only types
+                    # registered in this feature store, such as when only some
+                    # heterogeneous node types have quantized features.
+                    feature_output = replace(
+                        output,
+                        node={
+                            ntype: nodes
+                            for ntype, nodes in output.node.items()
+                            if ntype in sorted_ntype
+                        },
+                    )
                     nfeat_dict = self.dist_node_feature.get_all2all(
-                        output, sorted_ntype
+                        feature_output, sorted_ntype
                     )
                     for ntype, nfeats in nfeat_dict.items():
                         result_map[f"{as_str(ntype)}.nfeats"] = nfeats
                 else:
                     for ntype, nodes in output.node.items():
+                        if ntype not in self.dist_node_feature.local_feature:
+                            continue
                         nodes = nodes.to(torch.long)
                         futs[f"{as_str(ntype)}.nfeats"] = wrap_torch_future(
                             self.dist_node_feature.async_get(nodes, ntype)
                         )
+            if self.dist_node_quantized_feature is not None:
+                if self.use_all2all:
+                    sorted_ntype = sorted(
+                        self.dist_node_quantized_feature.local_feature.keys()
+                    )
+                    feature_output = replace(
+                        output,
+                        node={
+                            ntype: nodes
+                            for ntype, nodes in output.node.items()
+                            if ntype in sorted_ntype
+                        },
+                    )
+                    quantized_nfeat_dict = self.dist_node_quantized_feature.get_all2all(
+                        feature_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:
+                    for ntype, nodes in output.node.items():
+                        if ntype not in self.dist_node_quantized_feature.local_feature:
+                            continue
+                        nodes = nodes.to(torch.long)
+                        futs[
+                            f"#META.{NODE_PACKED_FEATURES_METADATA_KEY}.{as_str(ntype)}"
+                        ] = wrap_torch_future(
+                            self.dist_node_quantized_feature.async_get(nodes, ntype)
+                        )
             if self.dist_edge_feature is not None and self.with_edge:
                 for etype in self.edge_types:
                     if self.edge_dir == "in":
@@ -416,6 +481,10 @@ async def _collate_fn(
                 futs["nfeats"] = wrap_torch_future(
                     self.dist_node_feature.async_get(output.node)
                 )
+            if self.dist_node_quantized_feature is not None:
+                futs[f"#META.{NODE_PACKED_FEATURES_METADATA_KEY}"] = wrap_torch_future(
+                    self.dist_node_quantized_feature.async_get(output.node)
+                )
             if self.dist_edge_feature is not None:
                 eids = result_map["eids"]
                 futs["efeats"] = wrap_torch_future(

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,

gigl/distributed/dist_ablp_neighborloader.py

@@ -38,6 +38,7 @@
     extract_edge_type_metadata,
     extract_metadata,
     labeled_to_homogeneous,
+    materialize_quantized_node_features,
     set_missing_features,
     shard_nodes_by_process,
     strip_label_edges,
@@ -607,6 +608,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,
             ),
         )
@@ -796,6 +798,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=dataset.fetch_node_quantization_metadata(),
                 edge_dir=edge_dir,
             ),
             backend_key,
@@ -964,6 +967,12 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
 
         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.
         for key, value in metadata.items():