Import HTA to Chakra to extract synchronization dependency

[JoongunPark]

Summary

This PR is to process synchronization dependency between the Chakra nodes.
In order to do that, we use CriticalPathAnalyzer in Holistic Trace Analysis (https://github.com/facebookresearch/HolisticTraceAnalysis/blob/main/hta/analyzers/critical_path_analysis.py).

Please note that,

1. Command has been changed. We need to specify rank number by --rank.
2. Kineto trace profiler should collect nodes with their 'cat' field is 'cuda_sync'. Please follow the instruction here ([profiler] add option for kineto synchronization events in the trace pytorch/pytorch#105187).

Test Plan

Download and Install HTA

git clone https://github.com/facebookresearch/HolisticTraceAnalysis.git
cd HolisticTraceAnalysis
git submodule update --init
pip install -r requirements.txt
pip install -e .

Run Chakra et_converter

python3 tools/trace_link.py --pytorch-et-file result/eg.rank_0.pt.trace.json --kineto-file result/kineto.rank_0_step_5.1708449344148840892.pt.trace.json --rank 0 --output-file result/rank_0.json

Test input with Resnet-50 with 2 GTX1070 (rank 0)
eg.rank_0.pt.trace.json
kineto.rank_0_step_5.1708449344148840892.pt.trace.json

Test result with Resnet-50 with 2 GTX1070 (rank 0)
rank_0.json

Test Result with Megatron (No Sync dependency)

I've observed that this update will not cause any changes in result with trace which has no synchronization dependency.

Original
sys[4] finished, 607252677 cycles
sys[5] finished, 607253196 cycles
sys[6] finished, 607253715 cycles
sys[7] finished, 607254234 cycles
sys[0] finished, 607254753 cycles
sys[1] finished, 607255272 cycles
sys[2] finished, 607255791 cycles
sys[3] finished, 607256310 cycles

New
sys[4] finished, 607252677 cycles
sys[5] finished, 607253196 cycles
sys[6] finished, 607253715 cycles
sys[7] finished, 607254234 cycles
sys[0] finished, 607254753 cycles
sys[1] finished, 607255272 cycles
sys[2] finished, 607255791 cycles
sys[3] finished, 607256310 cycles

[TaekyungHeo]

I am not sure whether it is the best name to describe your method. Could you please justify the method name or rename the method name?

[JoongunPark]

I tried to make similar name with 'enforce_inter_thread_order'.

[TaekyungHeo]

I want to know whether we can assume that the annotation is always 'ProfilerStep'

[JoongunPark]

I think we can. That is always assumed in HTA testing example. HTA also describes in that way.

            annotation (str): a trace annotation to limit the analysis to,
                for example "ProfilerStep" would match all annotations that
                match this string (ProfilerStep#100, ProfilerStep#101 etc)

[TaekyungHeo]

What is the meaning of instance_id? Why is it set to zero?

[JoongunPark]

According to HTA, instance_id is used to classify which annotation to consider.

            instance_id: can be either of the following
                (int) - specify which instance of the annotation to consider.
                        Defaults to the first instance.
                (Tuple(int, int)) - considers a range of annotation instances start to end,
                        inclusive of both start and end instance.

[TaekyungHeo]

The tool fails with the following command.

$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_0.json --kineto-file ~/Downloads/megatron_kineto_0.json --output-file ~/rank0.json --rank 0 
[2024-03-15 17:16:12,875] trace_file.py:36 [WARNING]: No trace file is found in /Users/theo/Downloads
[2024-03-15 17:16:12,875] trace_file.py:88 [WARNING]: failed to create rank to trace map
[2024-03-15 17:16:12,875] trace.py:735 [INFO]: ranks=[]
[2024-03-15 17:16:12,875] trace.py:741 [ERROR]: The list of ranks to be parsed is empty.
Traceback (most recent call last):
  File "/Users/theo/param-jupark/train/compute/python/tools/trace_link.py", line 1196, in <module>
    main()
  File "/Users/theo/param-jupark/train/compute/python/tools/trace_link.py", line 1177, in main
    cpa = CriticalPathAnalyzer(
  File "/Users/theo/param-jupark/train/compute/python/tools/trace_link.py", line 39, in __init__
    self.event_sync_trace = TraceAnalysis(trace_dir = kineto_file)
  File "/Users/theo/HolisticTraceAnalysis/hta/trace_analysis.py", line 37, in __init__
    self.t.load_traces(include_last_profiler_step)
  File "/Users/theo/HolisticTraceAnalysis/hta/common/trace.py", line 614, in load_traces
    self.align_and_filter_trace(include_last_profiler_step)
  File "/Users/theo/HolisticTraceAnalysis/hta/common/trace.py", line 754, in align_and_filter_trace
    self._align_all_ranks()
  File "/Users/theo/HolisticTraceAnalysis/hta/common/trace.py", line 873, in _align_all_ranks
    self.min_ts = min(trace_df["ts"].min() for trace_df in self.traces.values())
ValueError: min() arg is an empty sequence

[JoongunPark]

The tool fails with the following command.

$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_0.json --kineto-file ~/Downloads/megatron_kineto_0.json --output-file ~/rank0.json --rank 0 
[2024-03-15 17:16:12,875] trace_file.py:36 [WARNING]: No trace file is found in /Users/theo/Downloads
[2024-03-15 17:16:12,875] trace_file.py:88 [WARNING]: failed to create rank to trace map
[2024-03-15 17:16:12,875] trace.py:735 [INFO]: ranks=[]
[2024-03-15 17:16:12,875] trace.py:741 [ERROR]: The list of ranks to be parsed is empty.
Traceback (most recent call last):
  File "/Users/theo/param-jupark/train/compute/python/tools/trace_link.py", line 1196, in <module>
    main()
  File "/Users/theo/param-jupark/train/compute/python/tools/trace_link.py", line 1177, in main
    cpa = CriticalPathAnalyzer(
  File "/Users/theo/param-jupark/train/compute/python/tools/trace_link.py", line 39, in __init__
    self.event_sync_trace = TraceAnalysis(trace_dir = kineto_file)
  File "/Users/theo/HolisticTraceAnalysis/hta/trace_analysis.py", line 37, in __init__
    self.t.load_traces(include_last_profiler_step)
  File "/Users/theo/HolisticTraceAnalysis/hta/common/trace.py", line 614, in load_traces
    self.align_and_filter_trace(include_last_profiler_step)
  File "/Users/theo/HolisticTraceAnalysis/hta/common/trace.py", line 754, in align_and_filter_trace
    self._align_all_ranks()
  File "/Users/theo/HolisticTraceAnalysis/hta/common/trace.py", line 873, in _align_all_ranks
    self.min_ts = min(trace_df["ts"].min() for trace_df in self.traces.values())
ValueError: min() arg is an empty sequence

This failure occurs when the HTA can not find the files in the directory. Could you check if the files are really there.