Lightning-AI / Lightning-AI/lightning-thunder

Use nvFuser executor decisions to pass on op execution to a different backend and retire hybrid `torch_compile_cat_ex` executor.

Open
#446 2 comments 2 reactions 0 assignees View on GitHub
enhancement executors torch.compile
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature

The feature request is to add the decision making capabilities inside nvFuser executor that allows the nvFuser executor to reject/pass on certain op executions where other backends/executors are known to have better performance.

### Motivation

The motivation for this is primarily to fix the longstanding performance issue #256

In brief, the hybrid `torch_compile_cat_ex` executor was introduced to execute the RoPE module as it contains a bunch of concat operations that nvFuser is known to not have optimized performance for. However, we don't have any direct pattern matching capabilities yet that can map just RoPE to a different backend.

To solve that issue, the hybrid executor tries to find concat operations and tries to fuse operations around it in a single region and pass that onto Torch.Compile. This helps improve performance but in some cases (like Dollyv2, Phi, Pythia etc.), the model trace is such that this hybrid executor ends up breaking apart a potential single nvFuser region into multiple by consuming an op in the middle of the graph which makes performance worse.

Summary - Disabling the hybrid executor gives us poor performance on RoPE, enabling it gives us poor performance due to smaller nvFuser regions.

### Pitch

We now have a natively integrated TorchCompile executor that is capable of taking any Thunder subgraph and generating TorchInductor optimized kernels for execution.

The pitch is for Thunder trace to pass on entire potential subgraphs to nvFuser including RoPE operations and nvFuser passes on the operations it chooses to not execute to this native TorchCompile executor. In this way, we are not relying on the hybrid executor to make good choices based on assumptions which may break from model architecture to another. This also might enable potential future choices like passing on certain ops to another backend based on tensor shapes etc.

### Alternatives

1. Fix the hybrid executor. This would involve changing the way the hybrid executor consumes operations from the Thunder trace. For example, we might assume that the output of RoPE goes into SDPA and force the hybrid executor to not consume any ops that feed into nvFuser regions or only feed into Matmuls/SDPA.
Whatever we do here will be a hack that might break on a different architecture.
2. Add a pattern matching solution for graphs (not limited to singular ops) such that we can create a pattern match replacement scheme for RoPE module to be executed by another backend like TorchCompile, APEX RoPE etc. This is a valuable feature on its own merit and serves a much broader scope. Naturally, it is also much more challenging to implement in the short term.
3. As suggested by Tom in a slack thread, a robust autotuning algorithm (with or without ML techniques) which allows selecting the best executors for the full graph. Similar to no. 2 above, it is a much broader feature and could be more challenging to implement.

cc @apaz-cli @lantiga @tfogal @mruberry @IvanYashchuk @jjsjann123 @kevinstephano @t-vi

[Long list of CC as it touches many pieces of Thunder]

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by tracing the nvFuser executor, the hybrid `torch_compile_cat_ex` executor, and the natively integrated TorchCompile executor to understand how Thunder subgraphs are assigned. Done means nvFuser can pass selected operations to TorchCompile without fragmenting useful nvFuser regions, and the hybrid executor is retired while preserving RoPE performance.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.