Skip to content

Conversation

@JoongunPark
Copy link
Collaborator

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.

7. Getter in ETFeeder

Updated ETFeeder to have getter functions of I/O attributes.
The I/O attributes include value/shape/type for the node.

Node that this feature is also required in other code in Feeder ( json_node.cpp json_node.h wrapper_node.cpp wrapper_node.h) which can be done after we decide details of JSON format.

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

rvinaybharadwaj and others added 30 commits October 7, 2024 17:44
fix lint errors

fix lint errors

fix lint errors
Without specifying the kineto filepath explicitly, HTA may pick
arbitrary files from the `trace_dir` and either provide incorrect analysis
results, or fail in some weird ways.
…neto-file-explicitly

Specify the kineto filepath explicitly when running HTA analysis
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants