Fix VLM + DDP checkpointing#3751
Conversation
Summary of ChangesHello @djsaunde, I'm Gemini Code Assist1! 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 addresses a critical stability issue in Vision-Language Models (VLMs) when trained with Distributed Data Parallel (DDP) by modifying how gradient checkpointing is handled. It specifically disables re-entrant checkpointing in DDP environments to prevent errors, ensuring that training can proceed without interruption. This change involves adapting both Unsloth's internal checkpointing mechanisms and external PyTorch/Hugging Face utilities to operate in a non-reentrant mode when DDP is active, providing a functional solution for VLM DDP training. Highlights
🧠 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. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses a checkpointing issue with Vision Language Models (VLMs) under Distributed Data Parallel (DDP) by disabling re-entrant gradient checkpointing. The approach involves monkey-patching torch.utils.checkpoint.checkpoint and transformers.modeling_utils.checkpoint to force use_reentrant=False when in a distributed environment.
My review focuses on the implementation of this monkey-patching. I've identified a potential bug where the checkpointing function could be wrapped multiple times if the patching code is executed more than once, and I've provided a suggestion to prevent this. I also raised a concern about the global nature of the patch and suggested considering a more scoped approach or providing a way to revert the patch in a future, more permanent solution.
Improves the existing DDP compatibility block (from PR unslothai#3751) with two targeted fixes that preserve Unsloth's memory optimizations: 1. Non-reentrant checkpointing with CPU activation offloading via saved_tensors_hooks. PR unslothai#3751 switched to non-reentrant but dropped CPU offloading entirely, increasing memory by ~10 GiB. 2. DDP-safe TiledMLP backward: uses functional torch.autograd.grad() for all-but-last sequence chunk (no DDP hooks fired), then .backward() for the final chunk (fires hooks exactly once). Both fixes are gated behind is_distributed() — single-GPU training is completely unaffected. Tested on Qwen3-VL-4B + LoRA with 8x L40S GPUs: - 128K tokens: 31.0 GiB peak, 300s/step (was OOM without these fixes) - 64K tokens: 33.8 GiB peak, 171s/step - 32K tokens: 24.8 GiB peak, 69s/step Fixes unslothai#3713 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2 participants
Closes #3713.
Unfortunately, this makes DDP training somewhat slower then expected for VLM models. I see this as a kind of short-term fix; we should try to really root cause this so we can still use our gradient checkpointing impl.
Searching online, I see a few different training libs had also run into this problem, and the fix was to fallback to
use_reentrant=Falsecheckpointing.See the linked issue for more details.