Optimize VRAM Allocation on 16GB GPUs to Prevent CUDA Out-of-Memory Errors

[purvanshjoshi]

Technical Overview

Fine-tuning Whisper-Medium with standard sequence lengths on 16GB VRAM GPUs (like Tesla T4) regularly triggers Out-of-Memory (OOM) failures due to the size of the attention map allocations. Training requires aggressive memory optimizations to maintain a stable effective batch size.

Affected Modules

• File: src/train.py (Line 110-140)
• System: VRAM allocations on 16GB GPUs.

Acceptance Criteria

1. Continuous Execution: The training script completes the full 4000 steps without hitting CUDA OOM constraints.
2. Stable Batch Representation: Maintains an effective batch size of 16 using gradient accumulation without exceeding 13GB VRAM allocation.
3. Optimized Throughput: Minimizes VRAM allocation while maximizing samples processed per second.

Proposed Implementation Approach

1. Enable gradient_checkpointing=True to reconstruct the backward pass activations dynamically, saving substantial forward pass memory.
2. Pin training parameters to fp16=True mixed precision.
3. Enforce 8-bit quantization mapping in combination with LoRA parameters target modules (q_proj, v_proj).
4. Set parameter configurations: per_device_train_batch_size=4 paired with gradient_accumulation_steps=4.

Severity & Priority

• Severity: High (Limits model scaling and training stability)
• Priority: P1

[Saksham-dev-del]

Hi @purvanshjoshi,

I would like to work on this issue. I will update src/train.py (lines 110–140) to apply the proposed memory optimizations:

Enable gradient_checkpointing=True.
Set fp16=True for mixed precision.
Apply 8-bit quantization with LoRA targeting q_proj and v_proj.
Configure per_device_train_batch_size=4 with gradient_accumulation_steps=4 for an effective batch size of 16.

I will validate the full 4000-step run completes on a 16GB GPU (e.g., Tesla T4) without exceeding 13GB VRAM.

Could you please assign this issue to me? I will submit a PR for review once complete.