[BUG] AMD: Windows ROCm image gen silently exhausts host RAM when the memory plan under-budgets

[LeoBorcherding]

What happens

On Windows ROCm an under-estimated diffusion memory plan does not fail. The generation keeps allocating, WDDM backs the overflow out of system RAM, and the machine becomes unusable (available RAM under 1.2 GB, pagefile past 50 GB). On Linux the same shortfall raises torch.OutOfMemoryError and the job dies cleanly.

The trigger is that the image memory plan is resolution-blind.

Why

estimate_image_runtime_mib (studio/backend/core/inference/diffusion_memory.py:273-292) scales headroom by pixel area against DEFAULT_IMAGE_WIDTH/HEIGHT, both 1024. Its only caller is _plan_memory (studio/backend/core/inference/diffusion.py:3006), and that caller passes no dimensions (diffusion.py:3130):

runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = variant_hint)

pixel_scale is therefore always exactly 1.0. Every image at every resolution is budgeted as one 1 MP frame. studio/backend/core/inference/video.py:1186 passes real values, so image is the odd one out.

The arithmetic

Z-Image-Turbo GGUF Q4_K_S at 1088x1920, 9 steps:

quantity	value
requested area	2,088,960 px, 1.99x of 1024^2
pixel_scale actually used	1.0
planned headroom	8192 x 1.0 x 0.85 = 6,963 MiB
headroom actually needed	~13,872 MiB
shortfall	~6.7 GiB

The planner picked its most conservative tier and still overran, because it was conservative about a workload half the real size:

diffusion.loaded: repo=unsloth/Z-Image-Turbo-GGUF base=Tongyi-MAI/Z-Image-Turbo
device=cuda offload=model tiling=True reasons=companions exceed budget;
whole-module offload of every component

Why Windows is worse than Linux

Measured from Windows performance counters on the backend process during the overrun:

counter	value
GPU Process Memory\Local Usage (dedicated VRAM)	13,215 MB
GPU Process Memory\Non Local Usage (system RAM as GPU memory)	14,025 MB
total held by the process	~27 GB on a 15.92 GiB card
system Available RAM	1,228 MB
Committed	80.5 GB of an 89.1 GB limit
pagefile	grown to 52.5 GB

Under WDDM the driver satisfies the overflow from shared system memory instead of failing. No OOM, no error, no failed job, just a desktop that stops responding.

Nothing in the inference path bounds this: torch.cuda.set_per_process_memory_fraction appears exactly once in the backend (studio/backend/core/training/worker.py:3734), inside run_training_process. No inference process reaches it.

Environment

• GPU: Radeon RX 9060 XT (gfx1200), discrete, 15.92 GiB
• OS: Windows 11
• torch 2.11.0+rocm7.13.0
• Model: unsloth/Z-Image-Turbo-GGUF, Q4_K_S, base Tongyi-MAI/Z-Image-Turbo

Reproduce

Load Z-Image-Turbo GGUF Q4_K_S on a ~16 GB Windows ROCm card, generate at 1088x1920. Watch GPU Process Memory\Non Local Usage for the backend process climb into the tens of GB while available RAM collapses, with no exception raised.

At 1024x1024 the same model completes, because that is the size the plan silently assumed.

Fix

Re-check at generate time, where the real resolution is known. The load-time width=None is defensible on its own (no request exists yet); the defect is that nothing revisited it. Compare the activation term alone against the current free budget: weights can be offloaded, activations cannot, so a runtime estimate over budget is an overrun at every tier rather than a tuning guess.

Longer term the ROCm OOM guard should extend to inference processes. Cross-importing from the training worker is already an established pattern (studio/backend/core/inference/llama_cpp.py:4864 does from core.training.worker import _rocm_classify_unified_memory).

Related

• [BUG] AMD: diffusion in studio falls back to whole-module offload and a 20 step image takes 48 minutes #8081 flags the same missing scaling term for the opposite failure. There the flat 8192 MiB is too large at 1 MP and below, so _group_fits() goes false and you get needless whole-module offload and a 48-minute image. Above 1 MP the same flat value is too small, which is this report. Note that a generate-time check does not close [BUG] AMD: diffusion in studio falls back to whole-module offload and a 20 step image takes 48 minutes #8081: that one needs the load-time plan itself changed, and these two want the estimate moved in opposite directions, so they should be fixed together rather than independently.
• fix(AMD): let diffusion load on Windows ROCm, and say what a gated repo needs #7981 (merged) fixed the xformers/torchao import wall on Windows ROCm and did not touch memory planning. Ruled out as a cause here.
• [BUG] AMD: image and video models in Unsloth Studio cannot load on Windows ROCm (torch.distributed has no attribute 'Work') #7992 (closed) was fixed by fix(AMD): let diffusion load on Windows ROCm, and say what a gated repo needs #7981.
• A separate video-gen OOM on this same host is a different bug: with xformers stubbed there is no memory-efficient attention on Windows ROCm, so _native_attention materialized a 66.54 GiB attention matrix. That one raised cleanly and was handled. Mentioned only to keep it out of scope.

Open question

It is not yet established whether the ~14 GB of non-local memory is caching-allocator bytes that WDDM placed in system RAM, or pinned host buffers created by enable_model_cpu_offload. This matters for the fix: set_per_process_memory_fraction bounds the former and not the latter.

Evidence leaning toward some of it sitting outside the allocator: the separate video OOM on this host reported "Of the allocated memory 9.09 GiB is allocated by PyTorch" while the process held substantially more GPU-addressable memory than that.

Settling it requires sampling torch.cuda.memory_reserved() against the live GPU Process Memory\Non Local Usage counter during a large generation.

[danielhanchen]

Confirmed, and you were right to insist this be fixed together with #8081. Every file and line citation checked out.

pixel_scale really is always exactly 1.0: _plan_memory is the only caller and it passes width = None, height = None, so every image at every resolution was budgeted as one 1 MP frame while video.py has always passed real values.

One caveat on the arithmetic, since it should not be read as a measurement. "Headroom actually needed ~13,872 MiB" is exactly what estimate_image_runtime_mib returns at 1088x1920 with the distilled discount: 8192 * 1.9921875 * 0.85 = 13872. It is the estimator's own output at the real resolution, not an observation of what the run consumed. It is still the right number to plan with, and it is reassuring that it matches your intuition, but the estimator has not been validated against a measured peak by this work.

Your proposed fix is what shipped: a generate-time re-check, comparing the activation term against the current free budget, on the grounds that weights can be offloaded and activations cannot. On Linux this is a nicer error than the OOM; on Windows it is the only protection there is, since WDDM never raises.

Three deviations worth flagging:

1. It only fires above what the load already budgeted. Treating the flat headroom as a hard limit at or below 1024x1024 would refuse work that succeeds today: an 8 GB card has a safe budget of 5898 MiB against a 6963 MiB default-resolution estimate, and those generations complete. So the refusal is confined to the resolution-driven overrun it is for.
2. It budgets one image, not the whole batch. The chunk loop already halves a failed multi-image forward down to singletons, so an oversized batch self-heals wherever torch raises. One image that does not fit is the case no backoff can rescue.
3. It counts the base overhead alongside the activations. Without that it was silent by a few hundred MiB on your exact card: 15.92 GiB gives a 14,254 MiB budget, which 13,872 MiB of activations fits and 13,872 + 2048 does not. With the overhead in, a 16 GB card at 1088x1920 is refused. This cannot cause a false refusal at or below the default resolution because of point 1.

On your open question, partly settled and honestly not more than that. I have no Windows ROCm card, so nobody sampled torch.cuda.memory_reserved() against the live counter. What the source does settle is one half: enable_model_cpu_offload goes through accelerate's cpu_offload_with_hook, which contains no pin_memory call anywhere. Diffusers' group offloading does pin, but your log says offload=model. So on that specific run there were no pinned host buffers from the offload path, which points at the allocator rather than at pinning. That is an inference from the code, not a measurement, and it does not generalise to the group tier.

Which turned out to matter. The #8081 fix moves cards like yours from whole-module offload onto a group tier, and that tier does pin. Converting a device-memory shortfall into ten-plus GB of unswappable host RAM on a machine whose available RAM you measured at 1.2 GB would have been a poor trade, so that tier passes low_cpu_mem_usage = True and declines to pin. The encoders it streams run once per call, so the slower copy is paid once rather than per step.

The guard is scoped to be correct either way: it bounds the request before allocation rather than trying to bound the allocator, so it does not depend on which pool the overflow would have come from. set_per_process_memory_fraction for inference processes is a reasonable separate ask and I have not done it here.

Before and after tables for both issues across card sizes, and the confirmation that neither fix reintroduces the other, are on the PR.

PR: #8224