Model residency: pinned host weights, disposable GPU copies - #48
Merged
Merged
Conversation
The staging mechanism (pin host copies once, onload by non-blocking DMA, offload by pointer swap) is about to be shared by the FLUX-family image loaders, so drop the ltx23 prefix: staging_ltx23.R -> staging.R, and .ltx23_pin_host/.ltx23_pin_component/.ltx23_staged_onload/ .ltx23_staged_offload -> .pin_host/.pin_component/.staged_onload/ .staged_offload. Doc topic staging_ltx23 -> staging (alias kept). Pure rename across the LTX pipeline, the Gemma3 encoder, and the fp8 loader; no behavior change.
flux1/flux2/zimage phase-swap the transformer, VAE decoder, and text encoder(s) CPU<->GPU every generation, through pageable host memory. Pin them once at load (the LTX staging mechanism) so onload runs at full PCIe rate and offload is a pointer swap. Payoff: flux1 T5 encode ~7-12 s -> sub-second, flux2/zimage Qwen3 swap ~2-4 s -> ~0.4 s, plus faster transformer/decoder moves. - Shared pin gate: extract .pin_decision (host RAM >= 2x pinned set) from recommend() and add .resolve_pin (explicit pin= > option > decision). Loaders gain pin=NULL. "full" checkpoints map to the bf16 row. - .flux_build_staging pins the listed components; the resident-fp8 transformer (flux2/zimage) also pins its plain weight_fp8/weight_scale fields via .flux_fp8_collect (set_data keeps the field reference live). Generation skips .flux_fp8_to_device when staging covers the transformer, else the field reassignment would orphan the staged pairs. - flux1 now GPU-encodes T5 (bf16, pinned) on 14 GB+ cards where its ~9.8 GB encode phase fits; nf4 tiers keep the CPU-fp32 encode. text_device default "cpu" -> NULL (profile-resolved); gc footprint sized to the T5 phase. An explicit text_device="cpu" still encodes in place (also fixes a latent flux2/zimage OOM on that path). - VAE decoder cast to compute dtype once at load so staged onload is device-only. onload/offload closures route by component name. - .pinned_set_gb flux1 row 43/31/26 -> 34/22/17 (T5 host copy now bf16). - Tests: test_flux_staging (gate logic, .flux_fp8_collect, pin round trip) and a zimage end-to-end smoke test.
The 0.2.0 submission was returned for four items, none of which
R CMD check implements (there is no missing-\value check anywhere in
tools, no \dontrun policy check, and no DESCRIPTION quoting or
whitespace check), so every local check had been green.
* \value on all 50 exported .Rd files that had a \usage block and no
return documentation, chiefly the nn_module generators for the FLUX,
FLUX.2, Z-Image, LTX-2.3 and Gemma3 ports. Doc-only topic pages have
no \usage and are correctly exempt.
* Examples: 14 of the 23 \dontrun{} blocks now run during check, and
were rewritten to be self-contained rather than referencing objects
that were never defined. The 9 that remain need model weights on
disk; cran-comments.md itemises them.
* DESCRIPTION: 'Python', 'Stable Diffusion' and 'Hugging Face' (with
its URL) single-quoted; trailing whitespace removed. Every
continuation line had ended in a space since the first commit, and
DCF folding turned each one into a double space.
Making the examples run surfaced two dead defaults. ddim_scheduler_create()
never passed beta_schedule through match.arg(), so switch() errored on the
length-3 default, and its device default was a length-2 vector that
torch_tensor() rejects; ddim_scheduler_step() had the same missing
match.arg() on prediction_type. Every internal caller passes these
explicitly, which is why the documented defaults were never exercised.
device now defaults to torch_device("cpu").
Also refreshes the safetensors messaging: mlverse/safetensors#11, #13
and #14 merged upstream on 2026-07-31 without a version bump, so the
advice now points at upstream rather than the fork, and the comments
record why the gates are runtime probes (the version number cannot
distinguish the two 0.2.1 builds). The install.packages()/
install_github() literals in hint text are reworded, since CRAN's
scanner flags the token regardless of context.
R CMD check --as-cran runs the examples a second time with \donttest{}
enabled, feeding diffuseR-Ex.R to R on stdin. The mp4 branch of
save_video() hands off to an ffmpeg process that inherits that stdin and
consumes a byte of the script, so every later example parsed one
character short: the run died on
base::assign(".ptime", proc.time(), pos = "CheckExEv")
Error in as.environment(pos) :
no item called "CheckExEv" on the search list
with the generated file plainly containing "CheckExEnv" (81 occurrences,
none of them mangled) and R evaluating a string one byte shorter. It
reproduced at the same offset with the GPU hidden, so it was never
CUDA-related.
The frames backend already exercises save_video() live and needs no
external encoder, so the encoder call is now a commented example with
the reason recorded next to it.
Implements the residency contract whisper (PRs #27/#28) and chatterbox (PR #25) already present, so the three packages offer one interface to a residency broker. diffuseR had the hard part already: staging.R pins component weights as page-locked host tensors and moves them by DMA. What was missing is the lifecycle handle around it. resident_load() load once, weights pinned, no VRAM held resident_activate() DMA the pinned copies onto the bound GPU resident_deactivate() re-point at the pinned copies, free the VRAM resident_generate() dispatch to the family's txt2img/txt2vid resident_status() state, components, pinned and GPU bytes resident_unload() terminal States are inactive -> activating -> active -> deactivating -> inactive. A failed transition rolls back to pinned host state and verifies it tensor-by-tensor; a rollback that cannot be verified leaves the handle broken, where only status and unload work. This sits ABOVE the per-generation phase offloading in the txt2img_* functions: those swap one component at a time within a render, residency decides who owns the card between renders. Components are discovered by scanning the pipeline for nn_module fields rather than hard-coded, since the families disagree (FLUX.1 has two text encoders, LTX adds a video VAE, an audio VAE and a vocoder). No gpu.ctl dependency, matching whisper and chatterbox: the contract is what a broker drives, not something that needs the broker present. Two API errors caught by check rather than by review, both from trusting a snippet instead of the namespace: torch has no cuda_memory_allocated() (the numbers are under cuda_memory_stats()$allocated_bytes$all$current) and no cuda_set_device(). A sweep of all 107 distinct torch:: calls in R/ now reports zero unexported names. Not yet exercised end to end against a real multi-GB pipeline: this box has ~3 GB of VRAM free with an RStudio session holding the rest. The state machine, component discovery, byte accounting and the full pinned<->GPU round trip are covered by 46 tinytest assertions.
Verified the contract against every family's local artifact on the 16 GB card, and two of the four failed. Both failures were mine. FLUX.1 OOMed on resident_activate(): its pinned set is 15.73 GB and the card has 15.47 GiB, so bulk-onloading every component cannot fit even though the phased render fits comfortably. The rollback worked, but the transfer should never have been attempted. LTX then hit the same path for a different reason: it takes phase_offload again at generate time and stores nothing on the pipeline, so reading pipeline$phase_offload returned NULL and was misread as "stays resident" -- it demanded 18.41 GB. So activation now follows how the pipeline was loaded. With phase_offload = TRUE (the default, and what every txt2img_* expects) there is no bulk transfer: the render moves each component on as its phase begins and back off as it ends, from the same pinned copies, so pre-loading them was undone within one phase. Only phase_offload = FALSE copies wholesale, and .resident_check_fits() refuses that up front with both figures named rather than letting libtorch report an allocator error. The phase_offload choice is captured in resident_load() instead of read back off the pipeline, since only the FLUX family keeps the field. Removing the redundant transfer made renders faster, not just safer: flux2 12.6 -> 9.1 s, zimage 27.3 -> 19.5 s, flux1 OOM -> 40.1 s. resident_status() gains components_on_gpu, because state and reality legitimately diverge: a render returns everything to pinned host memory as it finishes, so an active handle can hold nothing. state is the ownership claim, components_on_gpu is the measurement, and a broker should schedule on the measurement. Measured (RTX 5060 Ti, local artifacts): flux2 11.22 GB pinned / 9.1 s, flux1 15.73 GB / 40.1 s, zimage 13.45 GB / 19.5 s, ltx 18.41 GB across 5 components. All three image models reproduce bit-for-bit across a deactivate/activate cycle.
# Conflicts: # DESCRIPTION # NEWS.md
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
1 participant
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.
Implements the residency contract whisper (PRs #27/#28) and chatterbox
(PR #25) already present, so the three packages offer one interface to a
residency broker. diffuseR had the hard part already —
staging.Rpinscomponent weights as page-locked host tensors and moves them by DMA. What
was missing is the lifecycle handle around it.
States are
inactive → activating → active → deactivating → inactive. Afailed transition rolls back to pinned host state and verifies it
tensor-by-tensor; a rollback that cannot be verified leaves the handle
broken, where only status and unload work.
No
gpu.ctldependency, matching whisper and chatterbox: the contract iswhat a broker drives, not something that needs the broker present.
Verified against every family, and two of four failed first
FLUX.1 OOMed on
resident_activate(): its pinned set is 15.73 GB and thecard has 15.47 GiB, so bulk-onloading every component cannot fit even
though the phased render fits comfortably. LTX then hit the same path for
a different reason — it takes
phase_offloadagain at generate time andstores nothing on the pipeline, so reading
pipeline$phase_offloadreturned
NULLand was misread as "stays resident".So activation now follows how the pipeline was loaded. With
phase_offload = TRUE(the default, and what everytxt2img_*expects)there is no bulk transfer: the render moves each component on as its phase
begins and back off as it ends, from the same pinned copies, so pre-loading
them was undone within one phase. Only
phase_offload = FALSEcopieswholesale, and that path is refused up front with both figures named rather
than letting libtorch report an allocator error.
Removing the redundant transfer made renders faster, not just safer:
resident_status()gainscomponents_on_gpu, because state and realitylegitimately diverge: a render returns everything to pinned host memory as
it finishes, so an active handle can hold nothing.
stateis the ownershipclaim,
components_on_gpuis the measurement, and a broker should scheduleon the measurement.
Measured on an RTX 5060 Ti against local artifacts: flux2 11.22 GB pinned,
flux1 15.73 GB, zimage 13.45 GB, ltx 18.41 GB across 5 components. All
three image models reproduce bit-for-bit across a deactivate/activate
cycle.
R CMD check: 0 errors, 0 warnings, 2 NOTEs (New submission + the localtorch/lantern artifact).