Fix DDP "marked ready twice" for VLMs with CPU offload + TiledMLP#4240
Fix DDP "marked ready twice" for VLMs with CPU offload + TiledMLP#4240danielhanchen wants to merge 4 commits into
Conversation
for more information, see https://pre-commit.ci
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
for more information, see https://pre-commit.ci
Summary of ChangesHello, 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 significantly enhances the Distributed Data Parallel (DDP) compatibility for Unsloth's models, particularly those using Vision Language Models (VLMs) with CPU offloading and TiledMLP. The changes address the "marked ready twice" issue by introducing a robust non-reentrant checkpointing mechanism that integrates CPU activation offloading, thereby maintaining memory efficiency. Furthermore, it refines the TiledMLP's backward pass to correctly interact with DDP, preventing redundant hook firings. These targeted improvements ensure stable and efficient multi-GPU training for complex models within the Unsloth framework. Highlights
Changelog
Activity
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 introduces two important fixes for DDP compatibility in Vision Language Models, specifically when using CPU offloading and TiledMLP. The first fix correctly implements non-reentrant checkpointing with CPU activation offloading. The second fix provides a DDP-safe backward pass for TiledMLP by strategically using torch.autograd.grad and torch.autograd.backward to control DDP hook firing. The changes are well-implemented and gated behind is_distributed() to avoid affecting single-GPU training. My review includes a suggestion for improving robustness in the gradient handling and another for improving code maintainability by refactoring a complex loop.
| for i, x_split in enumerate(x_splits): | ||
| x_split = x_split.unsqueeze(0) | ||
| split_size = x_split.numel() | ||
| x_grad_slice = ( | ||
| x_gradients.view(-1) | ||
| .narrow( | ||
| dim = 0, | ||
| start = start_idx, | ||
| length = split_size, | ||
| ) | ||
| .view_as(x_split) | ||
| ) | ||
| grad_shard = ( | ||
| grad_output.view(-1) | ||
| .narrow( | ||
| dim = 0, | ||
| start = start_idx, | ||
| length = split_size, | ||
| ) | ||
| .view_as(x_split) | ||
| ) | ||
|
|
||
| x_split.requires_grad_(True) | ||
| with torch.enable_grad(): | ||
| outputs = TiledMLP.handle_output( | ||
| ctx.mlp_forward(x_split), | ||
| extra_outputs, | ||
| ) | ||
|
|
||
| if i < n_chunks - 1: | ||
| # Functional grad: no DDP hooks fired | ||
| grads = torch.autograd.grad( | ||
| outputs, | ||
| [x_split] + mlp_params, | ||
| grad_shard, | ||
| allow_unused = True, | ||
| ) | ||
| x_grad_slice.copy_(grads[0]) | ||
| for p, g in zip(mlp_params, grads[1:]): | ||
| if g is not None: | ||
| if p.grad is None: | ||
| p.grad = g | ||
| else: | ||
| p.grad.add_(g) | ||
| else: | ||
| # Last chunk: .backward() fires DDP hooks once | ||
| x_split.grad = x_grad_slice | ||
| torch.autograd.backward(outputs, grad_shard) | ||
|
|
||
| start_idx += split_size |
There was a problem hiding this comment.
The logic inside this for loop is quite complex and makes the _ddp_safe_backward function long (over 80 lines). To improve readability and maintainability, consider refactoring the loop body into a separate helper function. This helper could encapsulate the logic for processing a single chunk, including gradient calculation and accumulation, making the main function's flow easier to understand.
| grad_shard, | ||
| allow_unused = True, | ||
| ) | ||
| x_grad_slice.copy_(grads[0]) |
There was a problem hiding this comment.
For robustness, it's a good practice to check if grads[0] is not None before calling .copy_(). While x_split is expected to have a gradient since requires_grad is set to True, torch.autograd.grad with allow_unused=True can return None for inputs that don't receive gradients. If grads[0] were None for some reason (e.g., an unexpected graph structure where x_split is not used), this would raise an error. Adding this check makes the code more robust.
| x_grad_slice.copy_(grads[0]) | |
| if grads[0] is not None: | |
| x_grad_slice.copy_(grads[0]) |
This comment was marked as low quality.
This comment was marked as low quality.
1 similar comment
This comment was marked as low quality.
This comment was marked as low quality.
3 participants
Replacement for #4077 due to Studio rebasing
Improves the existing DDP compatibility block (from PR #3751) with two targeted fixes that preserve Unsloth's memory optimizations:
Non-reentrant checkpointing with CPU activation offloading via saved_tensors_hooks. PR Fix VLM + DDP checkpointing #3751 switched to non-reentrant but dropped CPU offloading entirely.
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(), so single-GPU training is completely unaffected.
Tested on Qwen3-VL-4B + LoRA with 8x L40S GPUs. These changes successfully enabled multi-GPU training.