Gemma3 NF4 on GPU: quantized encoder, batch encoding, byte-LUT eager dequant - #39
Merged
Conversation
gemma3_quantize_nf4() streams the HF snapshot into an NF4 artifact: the 7 projection weights per layer (~11B of 12B params) pack to NF4, embeddings and norms stay at the resident dtype, and vision-tower / projector weights are dropped (the text encoder never uses them). ~8 GB artifact vs 45 GB of host RAM for the fp32 CPU path, small enough to sit on a 16 GB card during the encode phase. load_gemma3_nf4() builds the skeleton at the compute dtype, swaps the projections for ltx23_nf4_linear modules, fills residents, and hard-errors on unfilled parameters; load_gemma3_text_encoder() dispatches to it when pointed at an artifact directory. The eager NF4 dequant (ltx23_nf4_dequantize) now uses the byte-LUT shared with the jit block stack: one 1-based int64 index per packed byte and a single gather straight into the output buffer, replacing the nibble-unpack / per-value-index / float32-staging chain. Scratch shrinks to one int64 buffer. This speeds every eager NF4 forward (Gemma3, FLUX.1, and the LTX fallback paths). Tests: quantize/load round-trip through a tiny HF-layout snapshot (hermetic fp32 residents), NF4 module swap, forward parity within quantization tolerance (cosine > 0.99), loader dispatch, and the unfilled-parameter hard error.
txt2vid with decode_video = TRUE, decode_audio = FALSE, and a filename crashed in write_wav: result$audio partial-matched result$audio_latents (no audio element exists on that path), passing the raw latents tensor. Exact [[ ]] indexing restores the intended NULL -> no-audio mux.
Encodes a character vector of prompts in batch_size chunks (bounding activation VRAM) and, with cache_dir, writes each prompt's result to gemma3-<md5>.pt and returns the paths - already-cached prompts are skipped, so interrupted batches resume. Without cache_dir it returns the per-prompt embedding list. Budget note in the docs: each result is the full hidden-state stack the LTX connectors consume (~0.4 GB at 1024 tokens). Includes a workaround for an R torch serialization bug found while testing: torch_save of a storage-offset view (e.g. narrow of a contiguous batch) serializes from the BASE storage, so every row saves as row 1; narrow()$contiguous() does not help because the view is already contiguous. clone() forces fresh storage. Minimal repro: torch_save(x$narrow(1,2,1)$contiguous()) loads x[1]'s data. Tests: batched-vs-solo parity through a namespace-stubbed tokenizer, cache path stability, resume without re-encode, cache round-trip content (would have caught the offset-view bug), and incremental cache growth.
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.
Summary
Puts the Gemma3 12B text encoder on the GPU via NF4 quantization, replacing the 45 GB fp32 CPU path for LTX prompt encoding.
gemma3_quantize_nf4()streams the HF snapshot into a ~8 GB artifact (one-time 52 s; 336 projection weights cast, vision tower dropped, resident embeddings/norms at bf16).load_gemma3_nf4()builds the skeleton at compute dtype, swaps projections for NF4 modules, hard-errors on unfilled params;load_gemma3_text_encoder()dispatches automatically when pointed at an artifact.gemma3_encode_batch()encodes a prompt vector in sub-batches with optional per-prompt disk caching (resumable, returns paths) — encode a whole batch of prompts once, swap the encoder off, render from cache.txt2vid_ltx2(decode_audio = FALSE, filename =)crashed via R partial matching feeding raw audio latents to the WAV writer; andtorch_saveof a storage-offset view serializes from the base storage (upstream torch bug —narrow()$contiguous()doesn't help because the view is already contiguous; the batch helper clones).Measurements (RTX 5060 Ti, gemma-3-12b-it)
Embedding drift vs fp32: cosine 0.989, rel RMS 0.19 — same-seed 768x512x49 renders come out scene-equivalent (same creature, stage, crowd, lighting; framing nudge comparable to sampler jitter). Fresh-prompt render cost drops from ~70 s to ~54 s.
Validation