Activation recompute can be skipped on the final instance of the recomputed modules to improve performance.

**Is your feature request related to a problem? Please describe.**

- While benchmarking Megatron-FSDP and FSDP2 in Automodel, we noticed that the final instance of the modules specified for recompute in MCore doesn't just store its activations in preparation for the backward pass. Can we add an argument to MCore that supports this? It will improve performance across the board. Are we doing something wrong?

@NVIDIA/mcore-oncall @NVIDIA/mcore-mlperf

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

- Ideally, we just keep the final activations in memory as they will be used immediately in the backward pass.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

N/A

**Additional context**
Add any other context or screenshots about the feature request here.

@Autumn1998 @ZhiyuLi-Nvidia for more context.