This PR introduces gradient checkpointing to reduce memory usage during model training.
Changes:
Block Class: Implemented gradient checkpointing using torch.utils.checkpoint, ensuring use_reentrant=False and requires_grad=True for inputs.
GPT Class: Applied gradient checkpointing to each transformer block.
Benefits:
Memory Reduction: Reduced GPU memory usage from around 19000MB to around 4000MB on the Shakespeare dataset with character tokenization, 512 context length and the softplus variation.
Feasibility: Enables training larger models on existing hardware.
Trade-offs:
Increased Computation Time: Activations are recomputed during the backward pass.
Added Complexity: Additional complexity in the model implementation and training loop.
Checklist:
[x] Implemented gradient checkpointing in Block and GPT classes.
[x] Tested for correct functionality and memory optimization.
[x] Verified memory usage reduction on Shakespeare character dataset.
Description:
This PR introduces gradient checkpointing to reduce memory usage during model training.
Changes:
torch.utils.checkpoint
, ensuringuse_reentrant=False
andrequires_grad=True
for inputs.Benefits:
Trade-offs:
Checklist:
Block
andGPT
classes.