Repository: pytorch
Issue: Improving format of communication metadata in PyTorch Execution Trace.
Labels: oncall: profiler

### 🚀 The feature, motivation and pitch

## TLDR
[PyTorch (Chakra) execution traces](https://mlcommons.org/working-groups/research/chakra/) are industry standard tracing tool to study and replay AI models on diverse hardware. A major use-case of this in the ability to *replay communication* collectives to study scale up and scale out AI training..

Communication collective metadata like process group, input side, collective is today added in a tracer function called `record_param_comms(...)`. The various metadata are represented as positional arguments in this function - 
```
 {
 "id": 3234, "name": "record_param_comms", "ctrl_deps": 3233,
 "inputs": {"values": [[[3232,3215,0,2049000,4,"cuda:0"]],35,["0","default_pg"],0,"allreduce",[],[],0,1,1], 
 "shapes": [[[2049000]],[],[[],[]],[],[],[],[],[],[],[]], 
 "types": ["GenericList[Tensor(float)]","Int","Tuple[String,String]","Int","String","GenericList[]","GenericList[]","Int","Int","Int"]}
```
Since metadata is represented as positional argument this leads to fragile trace parsing and can lead to issues.
We propose adding this metadata as "attributes" in the ET node itself.

```
  "attrs": [
    {
      "name": "pg_name",
      "type": "str",
      "value": "default_pg"
    },
```

## Background.

PyTorch's implementation of Chakra Execution Trace relies on observing the PyTorch functions using the `record_function` callback capability.  This means most of the metadata for an operator has to either arrive in input/output arguments or be known to the PyTorch framework.
(See the implementation here
https://github.com/pytorch/pytorch/blob/5785b02ba68f7ffe29c5aac086c5ddc84fc587c5/torch/csrc/profiler/standalone/execution_trace_observer.cpp#L163)

For collective communication we instrument PyTorch by adding macros `RECORD_PARAM_COMMS` and `RECORD_PARAM_COMMS_DATA`, these internally call the record_function API to trace several data.
https://github.com/pytorch/pytorch/blob/5785b02ba68f7ffe29c5aac086c5ddc84fc587c5/torch/csrc/distributed/c10d/ParamCommsUtils.hpp#L123-L134

The various metadata are positionally encoded in the Execution Trace using this mechanism. This leads to some **parsing challenges** and assumptions on the positional meaning of each operation.

https://github.com/facebookresearch/param/blob/main/train/comms/pt/commsTraceParser.py#L259-L261

## Proposal
### 1. Encode communication collective attributes as "attr" fields in Execution Trace nodes
Example -
```
  "attrs": [
    {
      "name": "pg_name",
      "type": "str",
      "value": "default_pg"
    },
    {
      "name": "pg_id",
      "type": "int64_t",
      "value": 0
    },
 ```
 
The list of fields we would like to encode includes
1. collective_name: name of collective
1. pg_name (int) and pg_desc (str): identifiers of process group for the collective
1. in and out number of elements
1. dtype: data type
1. in and out split size
1. global rank
1. global stride
1. world size

### 2. Attributes can be tagged to collective rather than the record_param_comms
We could chose to not use data in `record_param_comms` operator itself and hoist it to the actual collective. Thus we do not need to emit the record param comms node itself in ET and reduce the overall trace size. 

The feasibility of this might depend on internals of the ET implementation.

### Core Requirements
* We **do not break existing param_comms parsing** so no inputs/output should be changed in short run.
* **Integration/unit tests should continue to pass**.
* **Align with industry wide Chakra schema**.

### Impact
Much simpler collective metadata parsing that is maintainable in the future.

## Implementation

This is just a sketch but it involves a few steps
1. (PyTorch) Parse these attributes from record_param_comms nodes and set the them in the right place in ET observer. This will be the only change in PyTorch repo.
1. (Param) Have execution_trace observer support attr that are not present in all nodes.
1. (Replay) Update the et_replay to directly use these attributes, fallback to paramComms parser for older traces.



### Alternatives

_No response_

### Additional context

_No response_

cc @robieta @chaekit @aaronenyeshi @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @sraikund16 @sanrise