mlcommons / mlcommons/chakra

Communication Nodes Incorrectly Marked as COMP_NODE

Open
#172 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
196
Forks
84
PR merge metrics
No merged PRs in 30d

Description

I have followed the user guide and converted Pytorch ET + Kineto Trace into Chakra ET. However, it seems that communication nodes are incorrectly marked as computation nodes, for example:

{
  "id": "20",
  "name": "nccl:broadcast",
  "type": "COMP_NODE",
  "ctrlDeps": [
    "19"
  ],
  "dataDeps": [
    "19"
  ],
  "inputs": {
    "values": "[[21, 22, 0, 5, 8, 'cuda:7']]",
    "shapes": "[[5]]",
    "types": "['Tensor(long int)']"
  },
  "outputs": {
    "values": "[]",
    "shapes": "[]",
    "types": "[]"
  },
  "attr": [
    {
      "name": "rf_id",
      "int64Val": "19"
    },
    {
      "name": "fw_parent",
      "int64Val": "0"
    },
    {
      "name": "seq_id",
      "int64Val": "-1"
    },
    {
      "name": "scope",
      "int64Val": "7"
    },
    {
      "name": "tid",
      "int64Val": "1"
    },
    {
      "name": "fw_tid",
      "int64Val": "0"
    },
    {
      "name": "op_schema",
      "stringVal": ""
    },
    {
      "name": "is_cpu_op",
      "boolVal": true
    },
    {
      "name": "stream",
      "int64Val": "0"
    }
  ]
}

At https://github.com/mlcommons/chakra/blob/main/src/converter/pytorch_converter.py#L341

if json_node.is_gpu_op():
    if "ncclDevKernel_SendRecv" in json_node.name:
        parent_node = json_node_map[json_node.parent]
        keyword = (
            json_node_map[parent_node.parent].name
            if parent_node.name == "record_param_comms"
            else parent_node.name
        )
        if "send" in keyword:
            return COMM_SEND_NODE
        if "recv" in keyword:
            return COMM_RECV_NODE
    if "ncclKernel" in json_node.name or "ncclDevKernel" in json_node.name:
        return COMM_COLL_NODE
    return COMP_NODE

It seems that a node must be a GPU node before it becomes a communication node. However, the logic of json_node.is_gpu_op() is really confusing.
At https://github.com/mlcommons/chakra/blob/main/src/converter/pytorch_node.py#L149

def is_gpu_op(self) -> bool:
    """
    Check if the node is a GPU operator.
    
    Returns
        bool: True if the node is a GPU operator, False otherwise.
    """
    return self.cat is not None

However, it seems that the "cat" attribute would be dropped during the Pytorch ET + Kineto Trace link so that none of the nodes would be marked as a GPU node, and consequently marked as COMP_NODE.

I do not know which part is not as expected, but the logic of json_node.is_gpu_op() seems weird to me.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/converter/pytorch_converter.py around line 341 and src/converter/pytorch_node.py around line 149, then reproduce the reported PyTorch ET and Kineto Trace conversion. Trace how cat and is_gpu_op() affect classification of the nccl:broadcast example. Done means the conversion distinguishes communication nodes from computation nodes as intended and the reported output is corrected.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.