jina-ai / dalle-flow

🌊 A Human-in-the-Loop workflow for creating HD images from text
grpcs://dalle-flow.dev.jina.ai
2.83k stars 211 forks source link

Multi-GPU doesn't work due to hard coded replicas value #126

Closed Kadah closed 1 year ago

Kadah commented 1 year ago

Setting the various CUDA_VISIBLE_DEVICES to more than 1 GPU doesn't result in them getting used due to replicas being hardcoded to 1. Setting replicas to >1 allows flow to make use of more GPUs.

JoanFM commented 1 year ago

As explained here https://docs.jina.ai/fundamentals/flow/topologies/#replicate-on-multiple-gpus

This is expected, setting this env var will affect how each replica is assigned a GPU. So I would say this works as expected