Fix Chakra Converter Handling of NCCL All-to-All, Logging Consistency, and HTA Dependency Issues

[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. Support for PyTorch >2.5.0 and Updated Chakra Schema

The nightly version of PyTorch >2.5.0 appears to require support for schema version 1.1.1-chakra.0.0.4.
Updated Chakra to include the most recent schema version in its supported list.

5. Proper Encoding of pg_name in Collective Operations

Identified an issue where 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.

6. 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.

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 ✍️ ✅