Description

This PR introduces activation checkpointing for single-GPU fine-tuning, directly addressing issue #835. This feature allows developers to train models on hardware with less VRAM or to use larger batch sizes, improving throughput and hardware utilization.

Closes #835

Motivation

Currently, activation checkpointing in the repository is coupled with FSDP. By enabling it for standard single-GPU workflows, we unlock several benefits:

• Fine-tune on smaller GPUs: Makes fine-tuning accessible on a wider range of hardware.
• Increase batch sizes: Allows for larger batch sizes on capable GPUs, potentially improving training stability and MFU.
• Boost throughput: The memory savings can lead to higher overall training throughput despite the slight overhead per step.

Implementation Details

• Primary Method: Leverages Hugging Face's native gradient_checkpointing_enable() API for maximum compatibility and robustness.
• Memory Efficiency: Fully supports use_reentrant=False, which is the more memory-efficient checkpointing implementation.
• Supporting Utilities: Adds memory_utils.py for easy monitoring of CPU/GPU memory consumption during training runs.

Summary of Changes

• ✅ Added the --enable_activation_checkpointing flag to the training arguments.
• ✅ Integrated activation checkpointing into the single-GPU training loop.
• ✅ Added memory monitoring utilities for tracking process-level CPU and GPU usage.
• ✅ Maintained backward compatibility with existing MemoryTrace usage.
• ✅ Included comprehensive unit tests and an updated example script.
• ✅ Added documentation detailing the new feature and its usage.

Testing

The feature was tested locally and validated against the behavior described in the issue.

Local Testing (macOS, CPU):

• All new unit tests pass successfully.
• The fine-tuning example script runs to completion.
• Memory monitoring utilities report expected process memory.

Expected Results on GPU (based on issue #835):

• VRAM usage reduction: ~40%
• Potential batch size increase: ~70%
• Per-step overhead: +20-30%
• Net throughput (MFU) improvement: ~10-15%

Usage Example

To enable activation checkpointing, simply add the --enable_activation_checkpointing flag to your fine-tuning command.

torchrun --nnodes 1 --nproc_per_node 1 -m llama_cookbook.finetuning \
    --model_name meta-llama/Llama-3.1-8B-Instruct \
    --enable_activation_checkpointing \
    --enable_memory_monitoring \
    --use_peft \
    --peft_method lora \
    --batch_size_training 4 \
    --dataset alpaca_dataset \
    --output_dir "output/llama3-8b-checkpointed"

## Files Changed

- `src/llama_cookbook/utils/activation_checkpointing.py` (new)
- `src/llama_cookbook/utils/memory_utils.py` (new) 
- `src/llama_cookbook/finetuning.py` (modified)
- `src/llama_cookbook/configs/training.py` (modified)
- `src/llama_cookbook/utils/__init__.py` (modified)
- `tests/test_activation_checkpointing.py` (new)
- `examples/single_gpu_activation_checkpointing.py` (new)
- `docs/single_gpu_activation_checkpointing.md` (new)
- `.gitignore` (modified)