Thunder's horizontal fusion is memory inefficient for backward functions with activation checkpointing

@kiya00 is already working on this.

Since Oct 23, 2024.

TL;DR: Thunder's fusion pass needs to change to consider the memory usage of the operations and the intermediate tensors. It should avoid fusing operations that increase peak memory usage. Use `memory_peak_efficient_func` as the target function for optimization.

Let's take a look at the following code snippet:

```python
import torch
import gc

def memory_peak_efficient_func(t0s, a):
for t0 in t0s:
t1 = torch.nn.functional.relu(t0); del t0
t2 = torch.matmul(t1, t1); del t1
t3 = torch.nn.functional.relu(t2); del t2
a = torch.matmul(t3, a); del t3
return a

N_PARALLEL_PATHS = 10
t0s = [torch.randn(256, 256, device="cuda") for _ in range(N_PARALLEL_PATHS)] # 0.25 MiB * N_PARALLEL_PATHS
a = torch.randn(256, 256, device="cuda") # 0.25 MiB
memory_peak_efficient_func(t0s, a)

# Record peak memory usage
gc.collect(0)
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
before_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024)
memory_peak_efficient_func(t0s, a)
max_allocated_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024) - before_in_MiB
print(f"Peak memory usage diff: {max_allocated_in_MiB:.2f} MiB")
```
```plaintext
Peak memory usage diff: 0.75 MiB
```
The code snippet above is a simple example of a memory-efficient function that processes a list of tensors sequentially. The function avoids unnecessary memory allocations due to the order of operations and deleting intermediate tensors after they are no longer needed.

The code can be further "optimized" by fusing some operations into a single region. Before fusion let's take a look at the differently structured code snippet below:

```python
import torch
import gc

def memory_peak_inefficient_func(t0s, a):
# Can be fused into a single region
t1s = []
for t0 in t0s:
t1 = torch.nn.functional.relu(t0); del t0
t1s.append(t1)
del t0s

t2s = []
while t1s:
t1 = t1s.pop()
t2 = torch.matmul(t1, t1); del t1
t2s.append(t2)
del t1s

# Can be fused into a single region
t3s = []
while t2s:
t2 = t2s.pop()
t3 = torch.nn.functional.relu(t2); del t2
t3s.append(t3)
del t2s

while t3s:
t3 = t3s.pop()
a = torch.matmul(t3, a); del t3
del t3s

return a

N_PARALLEL_PATHS = 10
t0s = [torch.randn(256, 256, device="cuda") for _ in range(N_PARALLEL_PATHS)] # 0.25 MiB * N_PARALLEL_PATHS
a = torch.randn(256, 256, device="cuda") # 0.25 MiB
memory_peak_inefficient_func(t0s, a)

# Record peak memory usage
gc.collect(0)
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
before_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024)
memory_peak_inefficient_func(t0s, a)
max_allocated_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024) - before_in_MiB
print(f"Peak memory usage diff: {max_allocated_in_MiB:.2f} MiB")
```
```plaintext
Peak memory usage diff: 2.75 MiB
```

The code snippet above is a modified version of the previous code snippet computing the same result. It precomputes intermediate tensors and stores them in lists to be used later. This version of the code is less memory-efficient because it stores intermediate tensors in memory, which increases peak memory usage.

Let's apply Thunder on the inefficient code snippet to see if it can optimize the memory usage.

```python
import thunder

jit_memory_peak_inefficient_func = thunder.jit(memory_peak_inefficient_func)
gc.collect(0)
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
before_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024)
jit_memory_peak_inefficient_func(t0s, a)
max_allocated_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024) - before_in_MiB
print(f"Peak memory usage diff: {max_allocated_in_MiB:.2f} MiB")
```
```plaintext
Peak memory usage diff: 2.25 MiB
```

Thunder was able to optimize the memory usage of the inefficient code snippet by fusing the operations into a single region. The peak memory usage decreased from 2.75 MiB to 2.25 MiB, but it is still higher than the memory-efficient version of the code.

What would happen if we apply Thunder on the memory-efficient code snippet?

```python
import thunder

jit_memory_peak_efficient_func = thunder.jit(memory_peak_efficient_func)
jit_memory_peak_efficient_func(t0s, a)
gc.collect(0)
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
before_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024)
jit_memory_peak_efficient_func(t0s, a)
max_allocated_in_MiB = torch.cuda.max_memory_allocated() / (1024 * 1024) - before_in_MiB
print(f"Peak memory usage diff: {max_allocated_in_MiB:.2f} MiB")
```
```plaintext
Peak memory usage diff: 2.25 MiB
```

The same memory usage as the inefficient code snippet! Thunder was not able to optimize the memory usage of the memory-efficient code snippet because it applies topological sorting to the computation graph preferring forming as big fusion groups as possible. In this case, the memory-efficient code snippet already has the optimal order of operations, and Thunder breaks it by fusing operations into a single region.

Here's the dataflow graph of the execution trace of the memory-efficient code snippet:

```py
from thunder.core.transform_common import unwrap_return_value
from thunder.examine import make_trace_dot

t = unwrap_return_value(thunder.last_traces(jit_memory_peak_efficient_func)[-3])
dot = make_trace_dot(t)
```

And here's how it looked before the horizontal fusion:

```py
from thunder.examine import make_trace_dot

t = thunder.last_traces(jit_memory_peak_efficient_func)[0]
dot = make_trace_dot(t)
```

A lot more freedom in terms of grouping and reordering of operations as there are many parallel paths.

Thunder's fusion pass needs to change to consider the memory usage of the operations and the intermediate tensors. It should avoid fusing operations that increase peak memory usage. Here are the lines of code where the main logic for grouping operations is implemented:

https://github.com/Lightning-AI/lightning-thunder/blob/79e59d0c5c5f8aa8ef80eb31f3fe918466d64c1c/thunder/executors/data_dependent_partition.py#L299-L303

We need to resolve this issue because it leads to memory inefficiency in the generated code with activation checkpointing applied. The `memory_peak_efficient_func` should be the target function for optimization because the same pattern can be found in activation-checkpointed backward functions.

This problem was discovered with Yan's work on enabling PyTorch-native activation checkpointing in https://github.com/Lightning-AI/lightning-thunder/pull/1261. Here are last backward traces https://gist.github.com/kiya00/3ae4890e1ae5abf442d475cccadaa9ec#file-ckp_longchat-7b-16k_traces-py-L2116-L2117.

cc @apaz-cli