use tagging checkpointing

[t-vi]

This enables the checkpointing through tags and switches the litgpt benchmarks to this pattern.

Follows on top of #1615

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

• Eager instead of thunder: Average iter time: 803.95 ms Memory used: 17.55 GB
• Thunder with this PR: Average iter time: 779.80 ms, Memory used: 18.34 GB
• Main: Average iter time: 788.88 ms, Memory used: 18.76 GB
• refactor recomputation to work with tags #1615: Average iter time: 686.02 ms, Memory used: 20.77 GB

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 

• Eager instead of thunder: Average iter time: 1256.91 ms, Memory used: 10.37 GB
• Thunder with this PR: Average iter time: 1229.80 ms, Memory used: 10.27 GB
• Main: Average iter time: 1279.73 ms, Memory used: 13.27
• refactor recomputation to work with tags #1615 Average iter time: 1146.84 ms, Memory used: 13.29 GB

@riccardofelluga

[t-vi]

Marking this draft because it's onto the #1615 branch.

To my mind, both would be good to have soonish.

There is more investigation to do for Qwen and rematerialization.

[lantiga]

Stamped!