[bug fix] fix multi-gpu model loading bug

[billcao2000]

As mentioned in the following issue, the current code has model loading bug when running on multi-gpus. When launching the script with torchrun, the model loading logic incorrectly places all model weights on GPU 0 (cuda:0). This centralizes the memory load on a single GPU and results in an OOM error.
#1075

I've modified the code in [diffsynth/core/vram/disk_map.py] . I updated the safetensors model loading method within the DiskMap class to ensure proper GPU device mapping when loading the model for multi-GPU inference.
Root cause:
The safetensors library interprets the device 'cuda' as Device 0 (cuda:0) by default. Unlike native PyTorch functions, it does not automatically follow the current device context set by torch.cuda.set_device().

[gemini-code-assist]

Summary of Changes

Hello @billcao2000, 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 resolves a critical bug affecting multi-GPU model loading, specifically when using safetensors with torchrun. The core problem was that safetensors would default to cuda:0 for all weights, causing memory bottlenecks and OOM errors. The implemented solution dynamically assigns the correct CUDA device during model loading, enabling proper distributed inference and preventing memory issues.

Highlights

• Multi-GPU Model Loading Fix: Addresses a critical bug where safetensors incorrectly loaded all model weights onto cuda:0 by default in multi-GPU setups, leading to Out-of-Memory (OOM) errors. The fix ensures proper device mapping.
• Dynamic Device Assignment: Implemented logic within the DiskMap class to dynamically assign the correct CUDA device (cuda:{torch.cuda.current_device()}) when torch.distributed is initialized, preventing memory centralization on a single GPU.

🧠 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 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 addresses a critical bug in multi-GPU model loading where all weights were incorrectly placed on GPU 0, leading to OOM errors. The fix correctly identifies when running in a distributed environment and specifies the correct device for safetensors. The change is logical and effectively solves the issue. I have one suggestion regarding the placement of the new logic to improve code structure and maintainability.

[gemini-code-assist]

This logic correctly resolves the device for multi-GPU setups. However, placing it in flush_files means it will be executed every time this method is called (e.g., when the buffer is full), not just during initialization. Since the device for a distributed process is determined at startup and doesn't change, this check is only needed once.

For better code organization and to avoid redundant checks, I recommend moving this block to the __init__ method, right after self.device is assigned. This would centralize the instance's setup logic and make the purpose of flush_files clearer.

[Feng0w0]

Hello, I have some suggestions that you can take a look at:

1. Your code did not take into account the situation when use_disk_map=False, where the load_state_dict function is used to load weights of type safetensors. I suggest you also add this code to the load_state_dict_from_safetensors function.
2. This issue also occurs on NPU. Can the code be modified as follows to adapt to NPU at the same time:

import torch.distributed as dist
from diffsynth.core.device.npu_compatible_device import get_device_name, IS_NPU_AVAILABLE, IS_CUDA_AVAILABLE
if dist.is_available() and dist.is_initialized() and (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE):
    device = get_device_name()

[billcao2000]

@Feng0w0
Thanks for your advice!

1. Function load_state_dict_from_safetensors is called not only during the base model loading but also during LoRA loading. Modifying this function is of critical importance. Given your recommendations, I have added the corresponding code to resolve the GPU mapping issue.
2. Your suggested modifications for NPU support are correct. I believe the new code should run properly on NPU devices. However, since I don't have access to an NPU, I have only tested it on CUDA devices.

[Feng0w0]

@billcao2000
Thank you for adopting my suggestion.I have tested it on NPU device and your code can fix this issue on NPU.