Add sd21 to the residency contract - #63
Merged
Merged
Conversation
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.
# Conflicts: # DESCRIPTION # NAMESPACE # NEWS.md # R/resident.R # inst/tinytest/test_resident_sdxl.R # man/resident_load.Rd
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
mainbefore #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.peak6.807 GB. Comfortable on a 12 GB card. Suite 1197assertions, 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 forthis 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 CUDAdevice, so
sd21_load_pipeline()deliberately does not delegate to it.The device list is keyed by
get_required_components()rather than bythe pipeline's own fields. SD 2.1 requires an
encoder— the VAE encoderimg2img 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
encoderexplicitly instead of relying onstandardize_devices()'s fill-in; the sdxl end-to-end was re-rununchanged afterwards (18/18, identical numbers).
Why
gpu_components = "unet", for a different reason than SDXLSDXL 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 reston 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:
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:Truethe allocator itself suggests is the obviouslever, 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 insidea residency PR.