Pin FLUX.2's revision and fix the partial-onload wedge in resident LTX - #67
Merged
Merged
Conversation
Two things stood between the resident API and a host serving it.
`flux2_load_pipeline()` reached the hub at hfhub's default revision,
the branch `main` -- resolved through `refs/main` and then over the
network. A read-only bind of one snapshot carries neither, so the VAE,
the Qwen3 encoder and the tokenizer all failed to resolve inside a
container. `revision` now threads to all four `.flux2_cached` calls and
a branch name is refused rather than passed through.
`resident_load("ltx")` loaded a pipeline that could not generate.
`ltx23_load_pipeline()` does not load the text encoder and
`txt2vid_ltx2()` takes it per call, so the handle activated fine and
the first generation had no encoder. It now loads Gemma3 pinned on the
CPU and injects it, which also stops a serving caller re-reading 7.6 GB
of encoder for every clip.
The encoder is deliberately kept OUT of `staging`: that list is what
activation moves to the device, and the encoder belongs on the host.
Its pinned bytes are still counted, so the handle's declared host
footprint is the real one.
txt2vid_ltx2() chose the encode device with
`if (is.character(text_encoder)) device else "cpu"`, so a preloaded
encoder -- the resident/gpuhost path, resident_load("ltx", text_encoder
= ...) -- always encoded the prompt on CPU. That path loads the encoder
with pin = TRUE precisely so encode_with_gemma3() can DMA it onto the
card per encode and free it after, but the `else "cpu"` forced CPU and
left the pinned staging unused: every prompt paid the ~24 s CPU encode
instead of the ~7 s staged-GPU one, and on the gpuhost path that is once
per chunk.
The decision moves to .ltx23_text_encode_device(): a path loads onto the
asked-for device; a preloaded encoder that carries a `staging` set and
got a cuda request stages to the card; a bare CPU-resident object still
degrades to CPU, because a cuda request without staging would send the
tokens to the card while the weights sat on the host. Split out so the
rule is asserted without a GPU or a real encode
(test_text_encode_device.R, 8 cases), the way .resident_gen_args is.
The in-process backend is unaffected: it precomputes connector_embeds
and never reaches this branch.
The Gemma3 encoder's staged encode and the LTX pipeline's per-phase onload both decided "already on the card" by probing the FIRST staging pair. An onload that dies partway -- device memory runs out with most of the encoder already copied -- leaves exactly that pair on the card and the rest on the host, so every later call skipped the onload and failed on the first matrix multiply with "mat2 is on cpu", on every request, until the process ended. That is what took the gpuhost's ltx-2.3 entry down for USA 20260912 on 2026-09-10: one failed encoder onload at track 2, then the same refusal from every request after it, with nothing left in the process that could move the encoder back. Three changes. `.staged_on()` asks every pair, not the first one. `.staged_onload()` is idempotent per pair: a resident tensor is left where it is (no re-transfer over itself, which is what the first-pair probe existed to avoid) and a half-copied component is completed rather than restarted. And `encode_with_gemma3()` arms its offload BEFORE the onload, so a failed transfer is undone on the way out and the next encode starts from a clean host copy and reports the real error. The LTX onload closure drops its own probe and lets the per-pair rule decide; plain modules without staging keep the `$parameters` probe they had. Covered by test_staged_on.R (pure fakes, no torch) and a partial round trip added to test_staging.R (CUDA).
Two findings from the 2026-09-11 review of the staging fix.
`.staged_on()` and `.staged_onload()` compared device TYPE only, so a
request for "cuda:1" counted a tensor on cuda:0 as resident, skipped
the transfer, and left the weights on the wrong card. The old first-pair
probe had the same blind spot; the new helper just made it explicit.
`.device_spec()` now carries the index when the caller named one, and
`.on_device()` requires it to match. Bare "cuda" still accepts any card.
`resident_unload()` cleared `staging` and set pinned_bytes to 0 but
never touched `res$text_encoder`, which resident_load("ltx") keeps
outside `staging` on purpose. An unloaded handle therefore held the
Gemma3 encoder's pinned buffers while reporting none. Unload now
offloads the encoder's staging and drops the encoder and tokenizer.
Both covered without a GPU: the card cases in test_staged_on.R and an
unload-with-encoder case in test_resident.R.
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.
feat/pinned-revision, 0.2.2.7 to 0.2.2.10. The gpuhost image has run this tree since 2026-09-14; merging makes the next build from main carry it.resident_load("ltx")its text encoder and tokenizer (0.2.2.8)..staged_on()asks every staging pair,.staged_onload()is idempotent per pair, andencode_with_gemma3()arms its offload before the onload. This is the "mat2 is on cpu" wedge that took the ltx-2.3 entry down twice (0.2.2.10).resident_unload()drops the LTX text encoder.Tests:
test_staged_on.R(pure fakes), additions totest_resident.Randtest_text_encode_device.R, and a CUDA partial round trip intest_staging.R.Check status: both CI jobs pass the workflow, and both end
R CMD checkwithStatus: 2 WARNINGs, aboutvignettes/performance-levers.mdhaving no built copy underinst/doc. main's last run (#66, 25d95f5) ends with the same two; r-ci builds with--no-build-vignettesby default and this workflow does not override it. This branch does not touchvignettes/.