Skip to content

Add sd21 to the residency contract - #63

Merged
TroyHernandez merged 7 commits into
mainfrom
feat/sd21-resident
Aug 24, 2026
Merged

Add sd21 to the residency contract#63
TroyHernandez merged 7 commits into
mainfrom
feat/sd21-resident

Conversation

@TroyHernandez

Copy link
Copy Markdown
Contributor

The sixth resident family. Stacked on #62 — based on feat/sdxl-resident,
because it uses the generic pieces that PR introduces (.resident_gpu_set(),
.devices_for_pipeline(), the "cuda:N" dtype fix, the allocator pre-warm).
Retarget to main before #62 merges.

With that machinery already in place this is a loader shim, two switch
branches and one entry in .resident_families.

Verified end to end at 512x512: load 14.0 s (4.678 GB pinned), activate
0.54 s, real image, same seed reproducing after a full GPU round trip,
reserved.peak 6.807 GB. Comfortable on a 12 GB card. Suite 1197
assertions, 0 failures.

Two places SD 2.1 disagrees with SDXL

The UNet is float32, not float16, and this is a trap rather than a
preference. SD 2.1's attention overflows in float16 and the pipeline
returns all-NaN instead of raising, so a float16 resident would load,
activate, generate and hand back a blank image with nothing anywhere to
catch it. txt2img_sd21()'s native path already defaults to float32 for
this reason; the resident loader has to make the same call independently,
because it fixes the dtype before the weights are pinned and nothing
downstream can undo it. Note setup_dtype() answers float16 for a CUDA
device, so sd21_load_pipeline() deliberately does not delegate to it.

The device list is keyed by get_required_components() rather than by
the pipeline's own fields. SD 2.1 requires an encoder — the VAE encoder
img2img needs — that the text-to-image pipeline never builds, and
standardize_devices() refuses a list without it. It follows the decoder,
which is where it would live anyway. This also changes the sdxl injection,
which now supplies its encoder explicitly instead of relying on
standardize_devices()'s fill-in; the sdxl end-to-end was re-run
unchanged afterwards (18/18, identical numbers).

Why gpu_components = "unet", for a different reason than SDXL

SDXL is undone by its float32 VAE decode. SD 2.1 never reaches the decode:
the float32 requirement makes the denoise itself expensive. At the 768x768
default it wants 11.35 GB allocated plus 3.1 GB of allocator slack, which
loses on a 15.47 GiB card.

That is not a residency artifact, and I checked rather than assumed — a
plain txt2img_sd21() with the same devices and dtype OOMs at 11.494 GB,
within 0.04 GB of the resident run. The pre-warm and the bulk onload are
both exonerated. auto_devices() already answers "unet on cuda, the rest
on cpu" for this model at every strategy, so residency is agreeing with the
package rather than inventing a placement.

Moving the encoders off the card saves only 0.11 GB at 768, because the
peak is denoise activations rather than weights. Resolution is what decides
it:

resolution allocated.peak allocator slack result
512x512 6.807 GB 0.011 GB passes
768x768 11.35 GB 3.13 GB OOM on 15.47 GiB

A pre-existing issue this surfaced

SD 2.1's native float32 path does not fit its own default resolution on a
16 GB card, with or without residency. That wants its own fix — the
expandable_segments:True the allocator itself suggests is the obvious
lever, and diffuseR already uses it on other paths. It is recorded in NEWS
rather than papered over by quietly lowering a default, since changing
img_dim's default is a user-visible decision and not mine to make inside
a residency PR.

sdxl becomes the fifth resident family, and ltx stops leaking its NF4
dequantisation scratch when a handle deactivates.

The staging layer needed nothing: R/staging.R walks any nn_module's
$parameters/$buffers, so SDXL's four components pin, onload and offload
through the same path the other families use. Verified end to end on an
RTX 5060 Ti -- load 40.7 s (8.013 GB pinned), activate 2.51 s, a real
1024x1024 image, and the same seed reproducing bit-for-bit after a full
GPU round trip.

Four things did need doing:

* Only the UNet goes to the card. Bulk-onloading all four components fits
  the 8.0 GB of weights and then OOMs in the VAE decode, which runs
  1024x1024 in float32 while the UNet is still resident: measured 14.38
  GiB of a 15.47 GiB card, dying on a further 512 MiB request. txt2img_sdxl
  cannot evict the UNet before decode (its "phase cleanup" is gc plus
  empty_cache, which cannot move a live module), so the pipeline declares
  gpu_components = "unet" and the text encode and decode run on the host
  from the same pinned copies. The fit check is charged for what actually
  travels, not for the whole pinned set.

* The UNet dtype is fixed at load. sdxl_pipeline_from_safetensors derives
  it from the component device, so loading to CPU for pinning would
  page-lock a float32 UNet and then render in float32. A resident handle
  has to decide from where it will compute, not from where the weights are
  parked.

* resident_generate states the placement. txt2img_sdxl ignores the
  pipeline's actual devices and, with its default devices = "auto", calls
  auto_devices() afresh; on a 12 GB card that answers "unet cuda, encoders
  cpu" and the text encoder call then dies on a device mismatch. An
  explicit devices= from the caller still wins.

* A bulk activation pre-warms the allocator. A cold onload grows the pool
  one cudaMalloc per tensor and the syscalls dominate: 24.16 s against
  0.32 s warm, a 74x ratio. Same technique the NF4 LTX loader already uses.
  It is not only wall clock -- a broker with a startup deadline reads 24 s
  inside a first activate as a wedged worker.

Two fixes fell out of making that work, both of which affected callers
beyond residency:

* setup_dtype() rejected an ordinal-qualified device. resident_load binds
  an explicit "cuda:N" so transitions cannot drift, and "cuda:0" matched
  neither branch and hit "Invalid device".

* txt2img_sdxl demanded TorchScript .pt files it never opens. models2devices
  ends by verifying them, which is right when it is about to load them and
  wrong when the caller already holds a native safetensors pipeline;
  download_models = FALSE does not avoid the check. Callers with a pipeline
  now take .devices_for_pipeline() and skip it.

Separately, resident_deactivate now releases the NF4 dequantisation
buffers for ltx. Those live in a package-level environment rather than in
the module, so offloading the weights does not free them, and
txt2vid_ltx2 deliberately skips its own release while the transformer is
resident -- correct within a render, and it leaves the scratch on the card
once the render ends. Nothing else could reclaim it: the environment still
holds a reference, so gc() and cuda_empty_cache() cannot. Unconditional on
`release`, since a broker passing release = FALSE to keep the pool warm
for the next tenant is exactly the caller that must not be handed a budget
short by this scratch.
The sixth resident family, and the cheap one: the generic pieces landed
with sdxl, so this is a loader shim, two switch branches and one entry in
.resident_families.

Verified end to end at 512x512: load 14.0 s (4.678 GB pinned), activate
0.54 s, real image, same seed reproducing after a full GPU round trip,
reserved.peak 6.807 GB. Comfortable on a 12 GB card.

Two places SD 2.1 disagrees with SDXL:

* The UNet is float32, not float16, and this one is a trap rather than a
  preference. SD 2.1's attention overflows in float16 and the pipeline
  returns all-NaN instead of raising, so a float16 resident would load,
  activate, generate and hand back a blank image with nothing anywhere to
  catch it. txt2img_sd21's native path already defaults to float32 for this
  reason; the resident loader has to make the same call independently,
  because it fixes the dtype before the weights are pinned and nothing
  downstream can undo it. Notably setup_dtype() answers float16 for a CUDA
  device, so sd21_load_pipeline deliberately does not delegate to it.

* The device list is keyed by get_required_components() rather than by the
  pipeline's own fields. SD 2.1 requires an `encoder` -- the VAE encoder
  img2img needs -- that the text-to-image pipeline never builds, and
  standardize_devices() refuses a list without it. It follows the decoder,
  which is where it would live anyway. This also changes the sdxl injection,
  which now supplies its `encoder` explicitly instead of relying on
  standardize_devices' fill-in.

gpu_components is "unet", as for sdxl, but for a different reason. SDXL is
undone by its float32 VAE decode; SD 2.1 never reaches the decode. The
float32 requirement makes the denoise itself expensive: at the 768x768
default it wants 11.35 GB allocated plus 3.1 GB of allocator slack, which
loses on a 15.47 GiB card. That is not a residency artifact -- a plain
txt2img_sd21 with the same devices and dtype OOMs at 11.494 GB, within
0.04 GB of the resident run, so the pre-warm and the bulk onload are both
exonerated. auto_devices() already answers "unet on cuda, the rest on cpu"
for this model at every strategy, so residency agrees with the package
rather than inventing a placement.

Moving the encoders off the card saves only 0.11 GB at 768, because the
peak is denoise activations rather than weights. The resolution is what
decides it: 6.807 GB at 512 with 0.011 GB of slack, against 11.35 GB with
3.13 GB of slack at 768. SD 2.1's native float32 path not fitting its own
default resolution on a 16 GB card is a pre-existing property of the
package and wants its own fix; it is recorded here rather than papered over
by quietly lowering a default.

txt2img_sd21 also now skips the TorchScript .pt verification when it is
handed a pipeline, not only when it is handed a diffusers_dir -- the same
fix sdxl needed, and the path a resident handle arrives by.
@TroyHernandez
TroyHernandez changed the base branch from feat/sdxl-resident to main August 24, 2026 16:15
# Conflicts:
#	DESCRIPTION
#	NAMESPACE
#	NEWS.md
#	R/resident.R
#	inst/tinytest/test_resident_sdxl.R
#	man/resident_load.Rd
@TroyHernandez
TroyHernandez merged commit 2c0b504 into main Aug 24, 2026
2 checks passed
@TroyHernandez
TroyHernandez deleted the feat/sd21-resident branch August 24, 2026 16:40
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