[Feature] Implement Activation Checkpointing for Context Length Extension

[Sarkar-AGI]

Description

Currently, training with long context lengths increases activation memory significantly, leading to potential CUDA OOM errors. We need to implement activation checkpointing (gradient checkpointing) for our custom attention mechanism.

To-Do

• Identify bottleneck layers in the forward pass.
• Implement torch.utils.checkpoint or custom CUDA backward hooks to recalculate activations during backward propagation.
• Benchmark VRAM saving vs compute overhead trade-off.