Fix Chakra Errors

[JoongunPark]

Summary

This PR addresses multiple issues in the Chakra converter:

1. Improper Handling of NCCL All-to-All Communication

Chakra incorrectly distinguishes between point-to-point and collective communication. In NCCL, all-to-all is implemented as point-to-point communication, but Chakra's current logic treats these as distinct, leading to an incorrect type for PyTorchNode. More details on NCCL point-to-point can be found here.

2. Logging Inconsistency

There was a mismatch in logging levels: sync dependencies log via logging.info, while other dependencies use logging.debug. This PR resolves the inconsistency by standardizing the logging approach.

3. False Positive Dependencies from HTA

HTA returns false positives for sync dependencies, leading to invalid later op -> earlier op dependencies. This causes Chakra to fail in certain traces. The Chakra converter was found to encounter two critical failures:

• Cycle dependencies
• Stack overflows (due to call stacks exceeding 1000 levels)

4. Update trace_linker to use external_id for finding GPU op's parent CPU op

There were many operations matched with wrong parent CPU during trace linking.
This PR solves this problem using external_id instead of ev_idx.

5. Handling HTA Errors in Chakra

The trace linker was terminating unexpectedly due to errors in HTA. Although this may stem from trace inconsistencies, the issue does not occur when HTA is excluded.
Updated Chakra to handle these errors by raising exceptions instead of terminating the trace linker.

6. Proper Encoding of pg_name in Collective Operations

Identified an issue where SendRecv, Reduce-Scatter and All-Gather operations do not correctly encode pg_name following updates on the PyTorch side.
Modified Chakra to ensure proper encoding of pg_name in these collective operations.

Test Plan

I tested the fixes using Mixtral 8x3B traces collected with the NeMo framework (NVIDIA).
traces_device_0.zip

#!/bin/bash
# Set the result path
PATH="~/scratch/results/mixtral_8x3b/results"

# Loop through trace ranks
for i in 0
do
    echo "Start linking trace: $i"
    chakra_trace_link \
        --chakra-host-trace $PATH/host_$i.json \
        --chakra-device-trace $PATH/device_$i.json  \
        --rank $i \
        --output-file $PATH/rank_$i.json

    echo "Start converting trace: $i"
    chakra_converter PyTorch \
        --input $PATH/rank_$i.json \
        --output $PATH/rank_$i.et
done

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[spandoescode]

#160 was closed accidentally, this PR was made instead. Looks good to me!

[jinsun-yoo]

This string name nccl::all_to_all is hardcoded below:

https://github.com/pytorch/pytorch/blob/7c71ab1d40992ea2660bb124152e95b7a9a5119d/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L4974

There is a possibility that once this changes in PT, this breaks.

[jinsun-yoo]

c10d calls multiple send & receive calls, but we believe this is coalesced into one ncclDevKernel_SendReceive kenel, although we are not - yet- sure exactly where

c10d::allToall https://github.com/pytorch/pytorch/blob/7c71ab1d40992ea2660bb124152e95b7a9a5119d/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L4912

c10d::collective::fn
https://github.com/pytorch/pytorch/blob/7c71ab1d40992ea2660bb124152e95b7a9a5119d/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L3353C8-L3353C32
torch::cuda::nccl::all2all
https://github.com/pytorch/pytorch/blob/7c71ab1d40992ea2660bb124152e95b7a9a5119d/torch/csrc/cuda/nccl.cpp#L965

• designated here: https://github.com/pytorch/pytorch/blob/7c71ab1d40992ea2660bb124152e95b7a9a5119d/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L4958

[tushar-krishna]

This first issue (All-to-All) seems to be because of coalescing of multiple send-recvs (for All to All) into a ncclDevKernel_SendRecv ..
The confusion was that the name "SendRecv" implied a single src-dst pair. However, we confirm that one "SendRecv" nccl kernel can be the result of fusing multiple send and receive messages.

[tushar-krishna]

Looks good to merge!