refactor recomputation to work with tags

[t-vi]

Another step of #1560

This refactors the recomputation (activation checkpointing):

• recomputation works with tags,
• automatically tag intermediates from autograd on decomposition as to be recomputed, (this matches the behaviour of rematerialize_forward_backward, I think),
• I needed to disable rematerialize_forward_backward because it ran into "infinite capacity". However, I think this is not needed after this PR (cc @IvanYashchuk),
• move the uniform -> get_and_update_random_state + uniform_philox transform to before the autograd
• guard against recomputing random ops,

This is expected to be memory/compute neutral (I'll report numbers in a bit). It does not yet do the checkpointing frontend for the jit (including using memory commparable to eager checkpointing), that will be a separate PR.

@riccardofelluga could you take a look? (should be no surprises relative to #1560)

[t-vi]

So there is quite a regression in memory use:

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 --max_iters 3 --warmup_iters 2  --dump_thunder_traces True 

gives

• Main: Average iter time: 1279.73 ms, Memory used: 13.27 GB
• Average iter time: 1146.84 ms, Memory used: 13.29 GB

So 10% faster and not that much more memory (not sure what the 0.02GB are...)

[t-vi]

So I'm not entirely happy with disabling the rematerialize_forward_and_backward, but I wonder if the need arises from proxies that are re-computed and thus appear twice in the joint_extrace. (Hello name collision.)
I'll try that either here or in a follow-up (likely by renaming all fw tensors to fw_... and all bw tensors bw_ ... or somesuch.

[t-vi]

Unfortunately, disabling the forward-backward-rematerialization adversely affects memory for Qwen, but happily we recover that when we do better checkpointing:

python thunder/benchmarks/benchmark_litgpt.py --model_name Qwen2.5-7B --compile thunder --checkpoint_activations True --low_precision_mode none --micro_batch_size 1 --n_layer 4 --max_iters 3 --warmup_iters 2 --block_size 4096 

• Main: Average iter time: 788.88 ms, Memory used: 18.76 GB
• This PR: Average iter time: 686.02 ms, Memory used: 20.77 GB
• PR use tagging checkpointing #1616 : Average iter time: 779.80 ms, Memory used: 18.34 GB
• Eager instead of thunder: Average iter time: 803.95 ms Memory used: 17.55 GB

[riccardofelluga]

Overall it looks great! Just a couple of nits and clarifications

[IvanYashchuk]

What is the default value now?

[IvanYashchuk]

automatically tag intermediates from autograd on decomposition as to be recomputed, (this matches the behaviour of rematerialize_forward_backward, I think),

No, it doesn't match the behavior of rematerialize_forward_backward. What to recompute in fusion regions is decided by a min-cut-based algorithm. This PR introduced a regression in peak used memory (checked for Llama 2 7B): #1621.

I needed to disable rematerialize_forward_backward because it ran into "infinite capacity".

The rematerialization code has assumptions on input traces to be functioning, these assumptions were violated resulting in the "infinite capacity" error. @riccardofelluga was hitting the same problem when working on the recomputation. The problems are supposed to be fixed with #1367. Riccardo, what's the current status of 1367?

[riccardofelluga]

#1367 was parked due to change of priorities, we could tho bring that back here adapting to the new saved-for-backward logic of #1615. Tho it can also be that the infinite capacity error is reached by a different cause