Gradient Checkpointing Trade-off Lab: Compressing Memory by Crediting Compute

# === Selective gradient checkpointing — sadece attn'i recompute et ===
import torch
from torch.utils.checkpoint import checkpoint
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
 
# Monkey-patch decoder layer'ı: sadece self_attn'ı checkpoint, MLP'yi değil
original_forward = LlamaDecoderLayer.forward
 
def selective_forward(self, hidden_states, attention_mask=None, position_ids=None,
                       past_key_value=None, output_attentions=False, use_cache=False, **kw):
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
 
    # Self-attention'ı checkpoint
    def attn_fn(h):
        return self.self_attn(
            hidden_states=h, attention_mask=attention_mask,
            position_ids=position_ids, past_key_value=past_key_value,
            output_attentions=output_attentions, use_cache=use_cache,
        )
 
    attn_out = checkpoint(attn_fn, hidden_states, use_reentrant=False)
    if isinstance(attn_out, tuple):
        hidden_states = attn_out[0]
    else:
        hidden_states = attn_out
    hidden_states = residual + hidden_states
 
    # FFN'ı checkpoint'siz çalıştır — memory'sini sakla, recompute'a girme
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states
 
    return (hidden_states,)
 
LlamaDecoderLayer.forward = selective_forward
# Now training with selective ckpt