fix: multi-GPU training support for vision models #3809
Conversation
Summary of ChangesHello @Vinayyyy7, 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 resolves critical issues preventing successful multi-GPU training for vision models when using automatic or balanced device mapping. By intelligently detecting distributed training environments, the system now defaults to a data-parallel strategy, ensuring each GPU receives a complete model copy. This prevents model splitting across devices, which previously led to gradient calculation errors and 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://cdn.statically.io/img/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 critical issue with multi-GPU training for vision models by correctly handling device_map in distributed environments. The fix prevents model splitting across GPUs, which was causing runtime errors. The overall approach is sound. However, the implementation introduces code duplication between FastModel.from_pretrained and FastBaseModel.from_pretrained. My review includes suggestions to refactor this by removing the redundant code and centralizing the logic in the base model, which will improve maintainability and efficiency.
| # Fix for multi-GPU distributed training | ||
| # When using distributed training (e.g., 2x T4 on Kaggle), device_map="auto"/"balanced" | ||
| # splits the model across GPUs which can cause gradient device mismatch errors. | ||
| # Instead, use data-parallel approach where each GPU gets a full model copy. | ||
| from .loader_utils import prepare_device_map, is_distributed | ||
|
|
||
| if is_distributed(): | ||
| prepared_device_map, _ = prepare_device_map() | ||
| if device_map in ("auto", "balanced", "balanced_low_0"): | ||
| import warnings | ||
|
|
||
| warnings.warn( | ||
| f"Unsloth: Multi-GPU device splitting (device_map='{device_map}') is not supported " | ||
| f"for vision/multimodal models in distributed training. " | ||
| f"Using device_map={prepared_device_map} instead (data-parallel mode). " | ||
| f"Each GPU will load a full copy of the model.", | ||
| stacklevel = 2, |
There was a problem hiding this comment.
This logic for handling device_map in distributed training is duplicated from FastBaseModel.from_pretrained in unsloth/models/vision.py. Since FastModel.from_pretrained calls FastBaseModel.from_pretrained later in its execution, this check is performed twice. This is inefficient and creates a maintenance burden. To adhere to the DRY (Don't Repeat Yourself) principle, this entire block should be removed. The logic in FastBaseModel will handle this for all relevant models, especially after generalizing its warning message as suggested in the other comment.
| if device_map in ("auto", "balanced", "balanced_low_0"): | ||
| import warnings | ||
|
|
||
| warnings.warn( |
There was a problem hiding this comment.
To make this warning message more general and suitable for all models that use this base class, consider changing 'vision models' to 'vision/multimodal models'. This will allow you to remove the duplicated logic from FastModel.from_pretrained while keeping the warning accurate for all inheriting classes.
| warnings.warn( | |
| f"for vision/multimodal models in distributed training. " |
![chatgpt-codex-connector[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| f"Using device_map={prepared_device_map} instead (data-parallel mode). " | ||
| f"Each GPU will load a full copy of the model.", | ||
| stacklevel = 2, | ||
| ) | ||
| device_map = prepared_device_map |
There was a problem hiding this comment.
The new distributed guard now unconditionally replaces device_map with the per-rank map whenever is_distributed() is true. That code runs for any distributed launch (including inference with torchrun/accelerate), so runs that relied on device_map="auto"/"balanced" to shard a large vision model across GPUs now get forced into data-parallel replication on a single device, likely causing OOM or losing sharding entirely. Before this change, distributed inference respected the user-supplied sharded map. Consider limiting this override to training-only flows or providing an opt-out so multi-GPU inference continues to work.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
your concern: This fix might break distributed inference that relies on device_map="auto"/"balanced" for model sharding.
why this is unlikely to be an issue:
- Typical inference workflow in Unsloth notebooks:
After training, users typically run inference on a single GPU within the same notebook:
FastVisionModel.for_inference(model) # Switch to inference mode
inputs = tokenizer(image, text, ...).to("cuda") # Single GPU
model.generate(**inputs, max_new_tokens=128)this is single-process inference and is NOT affected by this fix. the
is_distributed()check returnsFalsein this case.
- For production inference with large models:
Users typically use vLLM (fast_inference=True) ORtransformers, NOT distributed device_map="balanced" with torchrun OR accelerate train.py scripts as unsloth docs suggest
The actual use case this fixes:
suppose user like me runs multi-GPU training on Kaggle 2x T4
torchrun --nproc_per_node=2 train.pywithout this fix, the model gets split across GPUs causing gradient device mismatch errors. fix enables data-parallel training where each GPU gets a full model copy.
How this was tested:
I tested this fix on Kaggle with 2x T4 GPUs by running code directly in notebook cells - no torchrun or accelerate launch commands. Just regular notebook execution with device_map="auto":
from unsloth import FastVisionModel
model, tokenizer = FastVisionModel.from_pretrained(
"Qwen/Qwen3-VL-2B-Instruct",
load_in_4bit=False, # 16-bit
device_map="auto",
)Training proceeds normally on 2 GPUs
trainer.train() # works now consistent with existing behavior:
Text generation models using FastLanguageModel already work fine with device_map="auto"/"balanced" on multi-GPU setups without requiring torchrun or accelerate commands. This fix brings the same capability to vision models using FastVisionModel which earlier crashed...
In a normal notebook colab/kaggle env (which is how most Unsloth users run training), the fix in cross_entropy_loss.py handles gradient device mismatch directly
inference is unaffected:
after training, users run inference on a single GPU:
FastVisionModel.for_inference(model)
inputs = tokenizer(image, text, ...).to("cuda")
model.generate(**inputs) # Single GPU inferenceThis is the standard workflow shown in Unsloth's official notebooks and remains completely unaffected.
But I do expect an actual human review from a maintainer for anything further, thanks...
|
Hey @Vinayyyy7 thanks for the contribution thoughts @danielhanchen ? |
|
agree that erroring out is cleaner than silently overriding. Users should explicitly knoww what's happening. clarifications: Splitting model across GPUs always crashed due to gradient mismatch errors, companion PR in unsloth-zoo (unslothai/unsloth-zoo#423) handles this by tracking the original device and moving gradients back. But even with that fix, model splitting has other issues with distributed training and quantization What device_map we can suggest users:
Proposed error message: raise ValueError(
f"Unsloth: device_map='{device_map}' is not supported for FastVisionModel in multi-GPU training. "
f"Model splitting across GPUs causes gradient device mismatch errors. "
f"For multi-GPU training, remove device_map or set device_map=None to use data-parallel mode "
f"where each GPU loads a full copy of the model. "
f"Note: This uses more VRAM per GPU but provides equivalent training to single GPU."
)Is this warning and
|
|
Hey @Vinayyyy7 |
|
it's just an if condition since we might warn users saying auto/balaned not supported in note: right know this is not implemented in PR, rn it just directly does data parallel if detects auto/balanced. |
2 participants
This PR fixes multi-GPU training for vision models when using
device_map="auto"ordevice_map="balanced".when running on multiple GPUs, setting
device_map="auto/balanced"withFastVisionModelClass causes model to be split across devices. During training, this results in hidden_states being on one GPU (e.g., cuda:1) while lm_head is on another (eg: cuda:0). The fused cross-entropy loss then computes gradients on the lm_head device but PyTorch expects them back on the original hidden_states device, causing a RuntimeError.Errors we see:
OR
this fix adds distributed training detection in
FastModel.from_pretrained()&FastBaseModel.from_pretrained().when distributed training is detected anddevice_mapis set toauto/balanced, it overrides to use data-parallel mode where each GPU loads a full copy of the model.Note: This PR works together with a corresponding fix in
unsloth-zoothat handles the gradient device mismatch in the fused CE loss.unslothai/unsloth-zoo#423
Tested on Kaggle with 2x T4 GPUs using
Qwen/Qwen3-VL-2B-InstructResult: IT WORKED