Summary of Changes

Hello @finbarrtimbers, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the DPO training module by introducing robust support for gradient accumulation and microbatching. These changes enable more efficient training with larger effective batch sizes, which is crucial for optimizing resource utilization and improving model performance, especially in distributed environments. The modifications involve restructuring the training loop, refining batch size calculations, and integrating mechanisms for proper metric accumulation across micro-batches.

Highlights

• Gradient Accumulation Support: Implemented gradient accumulation and microbatching for DPO training, allowing for larger effective batch sizes without increasing memory footprint.
• Batch Splitting Utility: Introduced a new utility function, split_batch_dpo(), to efficiently divide DPO batches into smaller micro-batches for sequential processing.
• Microbatch Size Configuration: Added the rank_microbatch_size parameter to DPOTrainModule's initialization, providing explicit control over the size of micro-batches processed per rank.
• Accumulator Tensors for Metrics: Integrated dedicated tensors (_total_loss, _total_chosen_logps, etc.) within DPOTrainModule to accurately accumulate metrics across multiple micro-batches before a single optimization step.
• Distributed Synchronization Control: Created _train_microbatch_context() to manage FSDP/DDP synchronization during gradient accumulation steps, ensuring gradients are only synchronized on the final micro-batch.
• Refactored Training Loop: Modified the train_batch() method to properly loop through micro-batches, ensuring correct loss scaling and metric accumulation for each micro-batch.
• Token Count Metric: Added token_count calculation in DataCollatorForSeq2SeqDPO and integrated its recording as a training metric in DPOTrainModule.
• Batch Size Calculation Fix: Corrected the calculation of rank_batch_size to properly incorporate gradient_accumulation_steps, ensuring accurate effective batch size computation.
• Optimized Cache Batch Size: Updated the cache_batch_size calculation for inference to a more efficient dp_world_size * 2.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.