Fix multi-node distributed training with single GPU per node

[Maxusmusti]

Fixes #4142

When each node has only one visible GPU (DEVICE_COUNT == 1), the _prepare_backend patch incorrectly disables distributed training — even in multi-node setups where WORLD_SIZE > 1.

This adds a WORLD_SIZE check so the patch is only applied when genuinely on a single device in a single-node setup.

Changes

One-line change in unsloth/models/_utils.py:

 -if DEVICE_COUNT == 1:
 +if DEVICE_COUNT == 1 and int(os.environ.get("WORLD_SIZE", "1")) <= 1:

Context

DEVICE_COUNT comes from torch.cuda.device_count(), which only counts locally visible GPUs. In a multi-node setup with one GPU per node (or CUDA_VISIBLE_DEVICES=0), DEVICE_COUNT is 1, but WORLD_SIZE is > 1. Without this fix, accelerate.state.PartialState._prepare_backend is patched to return DistributedType.NO, which prevents distributed communication from initializing.

[gemini-code-assist]

Summary of Changes

Hello, 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 addresses a critical bug affecting multi-node distributed training configurations, specifically when each node utilizes a single GPU. The change ensures that the system correctly identifies and initializes distributed communication, preventing unintended disabling of distributed training and allowing models to leverage multi-node resources as intended.

Highlights

• Distributed Training Fix: Corrected an issue where accelerate.state.PartialState._prepare_backend was incorrectly patched to disable distributed training in multi-node setups where each node had only one visible GPU.
• Conditional Patching Logic: Modified the condition for applying the _prepare_backend patch to include a check for WORLD_SIZE, ensuring it only applies when genuinely on a single device in a single-node environment.

🧠 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.

Changelog

• unsloth/models/_utils.py
  • Imported the os module to access environment variables.
  • Modified the conditional statement for patching _prepare_backend to check both DEVICE_COUNT and WORLD_SIZE.

Activity

• The pull request was created to fix issue [Bug] Multi-node distributed training broken only in one-GPU-per-node setting (one-line fix) #4142, which described incorrect distributed training behavior.
• The author, Maxusmusti, provided a clear description of the problem and the proposed solution, including a diff snippet.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly fixes an issue in multi-node distributed training setups with a single GPU per node. The change introduces a check for the WORLD_SIZE environment variable, preventing distributed training from being incorrectly disabled. The logic is sound. I've added one comment to suggest removing a redundant import statement for better code quality.

[gemini-code-assist]

The os module is already imported at the top of this file (line 97). This import os statement is redundant and can be removed.

[Maxusmusti]

Good catch, fixed!

[Maxusmusti]

Fix is tested: now in multi-node, single-gpu setting, correctly launch training with Data Parallel GPUs = 2, and runs as expected.

[Datta0]

This does make sense to me and seems quite a simple change

[mmathew23]

There may be some edge cases where are config sets a default value as empty string, but we can patch it if that becomes an issue.

[Maxusmusti]

@mmathew23 yeah thats a fair point, I was also thinking about that earlier. I figured since its a torchrun override its very rarely set manually, and when it is, throwing an error on a non-integer value would be good to let the user know their env has a bad value set for world size.

I didnt test the empty-string case though 😅, for some reason I was thinking since its an env var empty = unset and it'd pick the default, but thinking now I dont think that's true lol. If that becomes an issue someone hits in the future, I am happy to throw up a quick PR to patch that case!