Multi GPU inference for Qwen-Image-Edit-2509

[jingieboy]

How do i utilise multi GPU inference for qwen image edit 2509. I have 4 A6000 GPU's, with 48 GB VRAM in each one.

Only are supported:

device_map='balanced'
device_map='cuda'

No device_map='auto', sorry I am new to the huggingface ecosystem. When I load it in balanced and cuda, either errors like illegal memory is accessed or CUDA ran out of memory

[Samielakkad]

device_map="balanced" is model/component sharding for fitting one pipeline; it is not data parallelism and does not guarantee that temporary activation memory is evenly distributed. A single large transformer block, VAE step, or high-resolution temporary can still OOM one GPU even when aggregate VRAM looks sufficient.

Start with current compatible releases of Diffusers, Transformers, Accelerate, and PyTorch, use torch.bfloat16, and give Accelerate an explicit budget that leaves headroom on every card, for example about 44 GiB per A6000 rather than the full reported capacity. Inspect pipe.hf_device_map after loading to confirm where the transformer, text encoders, and VAE actually landed.

Important rules:

• Do not call pipe.to("cuda") after loading with a device map; that tries to collapse the placement. Do not combine a device map with CPU offload hooks unless you first reset the device map as documented.
• Test one image at a smaller resolution first. Enable VAE tiling/slicing where the pipeline supports it, then increase resolution/batch size separately.
• If pipeline-level placement leaves the large transformer on one card, load that transformer component separately with its own device_map/max_memory and inject it into the pipeline, following the model-level device-map guide.

An illegal memory access poisons the CUDA context, so restart the Python process before testing another configuration. Run the minimal failing case once with CUDA_LAUNCH_BLOCKING=1 to identify the first failing operation; later errors may only be consequences. Include the exact package versions, hf_device_map, resolution, and first synchronous traceback in an issue if it persists.

Finally, if your goal is throughput rather than fitting one request, run one pipeline process per GPU and distribute independent requests. That normally scales more predictably than sharding a single generation across four GPUs.

[rootkiller6788]

Fwiw, the device_map list you're hitting isn't specific to Qwen — DiffusionPipeline.from_pretrained only accepts
"balanced" / "cuda" / "cpu" (no "auto", no dict), and "balanced" places each top-level component on a GPU as a whole.
It does not shard one big model across cards. So with Qwen-Image-Edit you get something like: the ~40 GB (bf16)
transformer on one A6000, the 7B VL text encoder on another, small stuff wherever. The card holding the transformer is
left with ~8 GB for activations, which is a very OOM-friendly place to be at 1024² — even though 4×48 GB looks like
plenty.

Things to check first:

• Load in bfloat16. Without torch_dtype=torch.bfloat16 that transformer is ~80 GB in fp32 and won't fit on any single
  card at all — that alone produces exactly these errors.
• Pass max_memory that leaves headroom (e.g. {"cuda:0": "44GiB", ...}) and inspect pipe.hf_device_map afterwards to
  see where things actually landed.
• Don't call pipe.to("cuda") after loading with a device map.

And if your real goal is to spread that transformer over the four GPUs, whole-pipeline device_map won't get you there
— load the transformer separately with its own device_map/max_memory (that path does shard it) and pass it into the
pipeline. Component-per-card is the ceiling for pipeline-level balanced.

If it still crashes after bf16 + max_memory, the traceback is what actually matters — illegal memory access at that
point is usually a kernel/version problem, not RAM. Versions plus the first failure under CUDA_LAUNCH_BLOCKING=1 would
pin it down.