Analyze impact of traceinsight vs old HTA discrepancies #256

Title: trace_link: sync dependencies depend on HTA's critical-path heuristics (exposed by the traceinsight 0.6 upgrade)

Summary

chakra_trace_link keeps only the HTA sync dependencies that lie on the profiled run's critical path. traceinsight 0.6, adopted in #233, changes two heuristics in HTA's critical-path graph. Neither changes which sync edges HTA creates, but both can change which path is critical. So both can change the sync dependencies in Chakra's .et without any change in Chakra.

The before/after comparison in #233 found identical outputs on its single-rank and multi-rank traces. So this is a latent sensitivity, not an observed regression.

How Chakra uses HTA

TraceLinker.load_sync_dependencies runs critical_path_analysis, iterates cp_graph.critical_path_edges_set, and keeps the SYNC_DEPENDENCY edges. critical_path_edges_set holds only the edges on the single longest path (nx.dag_longest_path), so any change to the graph's edges or weights can move it.

The two changes (old pin d731cc2e → traceinsight 0.6.1)

1. Kernel-to-kernel cutoff.
   • The old pin always added an edge from a kernel to the previous kernel on the same stream (old).
   • 0.6.1 adds the edge only if the gap is under 1.5 ms (KERNEL_KERNEL_DELAY_THRESHOLD_US, used here). Longer gaps are treated as an unrecognized dependency.
   • Affects traces with idle gaps over 1.5 ms on a stream: communication waits, data-loading stalls, CPU-bound phases.
2. Backward-thread merge.
   • If exactly one thread has forward events and exactly one has autograd events, 0.6.1 remaps the backward events onto the forward thread before building CPU dependency edges (new). The old pin built per-thread call stacks with no remap (old).
   • Applies to most PyTorch training traces.

Verified unchanged between the two versions

• Creation of SYNC_DEPENDENCY edges: same two sites, refactored only.
• ev_idx is still the event's position in the raw traceEvents list, which Chakra uses to look up External id.
• The thread remap happens only inside HTA's internal copy of the trace. The raw events Chakra reads are unaffected.

Why this matters beyond the upgrade

A sync is an ordering constraint whether or not it was binding in the profiled run. Chakra traces are replayed on other systems (e.g., ASTRA-sim), where a different sync can become binding. Filtering by the measured critical path builds the profiled machine's timing into the graph. It also makes Chakra's output depend on HTA's path heuristics and on tie-breaking in dag_longest_path.

Proposal

Keep every SYNC_DEPENDENCY edge in the graph, not only those on the critical path:

for _, _, data in cp_graph.edges(data=True):
    edge = data["object"]
    if edge.type == CPEdgeType.SYNC_DEPENDENCY:
        ...  # existing External id lookup unchanged

HTA has no public accessor for edges by type, so this relies on HTA's internal "object" edge attribute. It may be worth asking HTA upstream for an accessor.

Open questions

• Was the critical-path filter deliberate? It was introduced in aacfcff ("Leverage HTA to support synchronization dependency in trace linking") with no stated rationale. Possible reasons include limiting edge count or avoiding spurious edges.
• How many sync dependencies does the change add on typical traces, and does it change ASTRA-sim replay results?

Acceptance criteria

• A decision between critical-path-only and all sync edges, recorded in the load_sync_dependencies docstring.
• A unit test that pins the chosen behaviour: a mocked CPGraph with one sync edge off the critical path.
• If switching: a before/after comparison on single-rank and multi-rank traces, reporting the sync-dependency count for each mode, an .et diff confirming only edges were added, and an ASTRA-sim replay comparison on at least one trace.

Discrepancy discovered by @mihncalo while performing tests with traceinsight. To check if this has a real impact or not.