Skip to content

Allow splitting flow to two GPUs - #14

Open
amv wants to merge 2 commits into
microsoft:mainfrom
amv:two-gpu
Open

Allow splitting flow to two GPUs#14
amv wants to merge 2 commits into
microsoft:mainfrom
amv:two-gpu

Conversation

@amv

@amv amv commented Jul 22, 2026

Copy link
Copy Markdown

Allows running the model with two consumer grade GPUs:

mage-flow-app --device cuda:0 --text-device cuda:1

The "--text-device" card needs to have roughly 10GB of free VRAM and the "--device" card needs roughly 12GB of free VRAM.

Change was done with the help of a local Qwen3.6 27b model and reviewed by manually. No copyrights reserved.

@amv

amv commented Jul 22, 2026

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree

@amv

amv commented Jul 22, 2026

Copy link
Copy Markdown
Author

I'm not really sure if this is the best way to do the split, but it seemed the easiest way to get it working with a limited amount of memory, and this might be useful for someone even if it is not merged.

The LLM also gave an analysis that since the text model and the diffusion model are not actually running at the same time, it would theoretically be possible to allow running this also with a 12GB VRAM card by doing the text inference first, unloading the text model, and then loading the image models. But this seemed a bit slow, especially when doing multiple images, so I did not pursue it.

I also witnessed some issues with the normal model switching where the previous model was not properly unloaded before the another one was loaded in, so it probably would have required a more sizeable patch to impleement the unloading and switching. It might also be that the model switching OOM situation is also somehow a result of my patch, but it doesn't seem obvious based on the code that it would be, and I don't have a card with more than 16GB of VRAM so I couldn't try it out.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant