update checkpointing support for jit

[t-vi]

I'll add switches and testing of checkpointing, but here is the material code changes.

[t-vi]

There are a number of things still to be fixed.
There is a failure with the memory calculation because we don't update the initial collection proxy.
I'll look into these.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[t-vi]

To my mind, the remaining missing bit is the handling of uniform -> uniform_philox conversion (currently handled by rematerialization) and the interaction with recomputation. I'll look into it. In the meantime, I'd be keen to hear complaints and/or success stories about memory / performance impact.

[riccardofelluga]

I was trying to check memory savings but it looks like the following hangs:

python thunder/benchmarks/benchmark_litgpt.py --model_name stablecode-completion-alpha-3b --compile thunder --checkpoint_activations True --low_precision_mode none --micro_batch_size 1

:(

[t-vi]

So one needs to enable checkpointing layers with compiler==thunder for this. Then the memory profile of the backward is still terrible:
Running @riccardofelluga 's benchmark:
Eager with checkpointing and 8 layers needs 12.59GB, thunder with checkpointing and 8 layers needs 38GB.
This is not the saved for backwards, which is (for thunder): Saved for backward size: 1761.89 MiB Saved for backward number of tensors: 103

I think we need to look more closely at the memory over time.

[t-vi]

Doing the following:

• apply del last used to any trace of the last_backward_traces
• run the memory examine tool (so awesome @kiya00 !)

We see that we are doing much better at first than after transform for execution (this is for 4 layers), so we still have reordering that hurts us.

#the following trace uses ~9.84GB memory
# Constructed by Saved for backward remat trace (took 20.46 milliseconds)
#the following trace uses ~12.39GB memory
# Constructed by Transform for execution (took 886 milliseconds)

Note that the difference of 2.55GB is smaller than the difference between thunder and eager (5.82GB).

[t-vi]

With the two latest bits

• dumb down fusion logic to not change ordering,
• splitting ThunderFunction to avoid the effect in Loss function backward output kept in memory longer than expected #1379,

I have that

python thunder/benchmarks/benchmark_litgpt.py --model_name stablecode-completion-alpha-3b --compile thunder --checkpoint_activations True --low_precision_mode none --micro_batch_size 1 --n_layer 4

is on par (even a little below) the same with --compile eager.
I will be looking at the failing tests and splitting out some bits that can be handled independently.