Fix error encoding METADATA node

mlcommons · Dec 1, 2024 · cbcce6a · cbcce6a
1 parent 4678f31
commit cbcce6a
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/src/converter/pytorch_converter.py b/src/converter/pytorch_converter.py
@@ -11,6 +11,7 @@
     COMM_RECV_NODE,
     COMM_SEND_NODE,
     COMP_NODE,
+    METADATA_NODE,
     REDUCE_SCATTER,
     GlobalMetadata,
 )
@@ -338,6 +339,8 @@ def get_protobuf_node_type_from_json_node(
         Returns:
             int: The corresponding Chakra node type.
         """
+        if json_node.is_metadata_op():
+            return METADATA_NODE
         if json_node.is_gpu_op():
             if "ncclDevKernel_SendRecv" in json_node.name:
                 parent_node = json_node_map[json_node.parent]

diff --git a/src/converter/pytorch_node.py b/src/converter/pytorch_node.py
@@ -137,6 +137,15 @@ def get_op_type(self) -> PyTorchNodeType:
         else:
             return PyTorchNodeType.LABEL
 
+    def is_metadata_op(self) -> bool:
+        """
+        Check if the node is a METADATA operator.
+
+        Returns
+            bool: True if the node is a METADATA operator, False otherwise.
+        """
+        return self.get_op_type() == PyTorchNodeType.METADATA
+
     def is_cpu_op(self) -> bool:
         """
         Check if the node is a CPU operator.

diff --git a/tests/converter/test_pytorch_converter.py b/tests/converter/test_pytorch_converter.py
@@ -167,6 +167,7 @@ def test_write_chakra_et(mock_file: MagicMock, sample_pytorch_data: Dict) -> Non
 @pytest.mark.parametrize(
     "pytorch_node_data, expected_type",
     [
+        ({"name": "process_group:init", "is_gpu_op": False}, METADATA_NODE),
         ({"name": "ncclKernel", "is_gpu_op": True}, COMM_COLL_NODE),
         ({"name": "ncclDevKernel", "is_gpu_op": True}, COMM_COLL_NODE),
         ({"name": "c10d::all_reduce", "is_gpu_op": True}, COMP_NODE),