diff --git a/DESCRIPTION b/DESCRIPTION index 92dfaa1..af2e1c6 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.1.0.12 +Version: 0.1.0.13 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), diff --git a/NAMESPACE b/NAMESPACE index 7232a59..3c5e787 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -66,6 +66,8 @@ export(flux2_unpack_latents_with_ids) export(flux2_unpatchify_latents) export(flux2_vae_decoder) export(gemma3_config_ltx2) +export(gemma3_encode_batch) +export(gemma3_quantize_nf4) export(gemma3_text_model) export(gemma3_tokenizer) export(img2img) @@ -74,6 +76,7 @@ export(latents_to_video) export(load_decoder_safetensors) export(load_decoder_weights) export(load_flux2_vae_decoder) +export(load_gemma3_nf4) export(load_gemma3_text_encoder) export(load_model_component) export(load_pipeline) diff --git a/NEWS.md b/NEWS.md index 26e4d24..85594ed 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,25 @@ +# diffuseR 0.1.0.13 (development) + +* `gemma3_quantize_nf4()` + `load_gemma3_nf4()` put the Gemma3 12B text + encoder on the GPU: the 336 projection weights quantize to a ~8 GB + NF4 artifact (one-time 52 s; vision tower dropped), which loads in + ~12 s and encodes in ~7 s on CUDA vs ~30 s load + ~24 s encode for + the fp32 CPU path, at 8 GB host RAM instead of 45. Pinned phase swap + costs 0.3 s on / 0 s off after a one-time 3.3 s page-lock. + `load_gemma3_text_encoder()` dispatches to the artifact + automatically. Embedding drift (cosine 0.989 vs fp32) renders + scene-equivalent same-seed videos. +* The eager NF4 dequant now shares the byte-LUT with the jit block + stack (one index per packed byte, gather straight into the output), + speeding every eager NF4 forward (Gemma3, FLUX.1, LTX fallbacks). +* Fixed `txt2vid_ltx2(decode_audio = FALSE, filename = )`: partial + matching handed the raw audio latents to the WAV writer. +* `gemma3_encode_batch()` encodes a prompt vector in sub-batches with + optional per-prompt disk caching (resumable; returns cache paths): + encode a whole episode list once, then swap the encoder off and + render from the cached embeddings. Works around a torch_save bug + where storage-offset views serialize from the base storage. + # diffuseR 0.1.0.12 (development) * LTX video decode runs untiled when the estimated activation cost diff --git a/R/gemma3_text_encoder.R b/R/gemma3_text_encoder.R index 53caa99..466eeb6 100644 --- a/R/gemma3_text_encoder.R +++ b/R/gemma3_text_encoder.R @@ -550,9 +550,33 @@ gemma3_config_ltx2 <- function() { # Weight Loading # ----------------------------------------------------------------------------- +# Model config with HF defaults, shared by the checkpoint and NF4 +# artifact loaders (the raw text_config travels in the NF4 manifest) +.gemma3_build_config <- function(config_raw) { + list( + vocab_size = config_raw$vocab_size %||% 262208L, + hidden_size = config_raw$hidden_size %||% 3840L, + intermediate_size = config_raw$intermediate_size %||% 15360L, + num_hidden_layers = config_raw$num_hidden_layers %||% 48L, + num_attention_heads = config_raw$num_attention_heads %||% 16L, + num_key_value_heads = config_raw$num_key_value_heads %||% 8L, + head_dim = config_raw$head_dim %||% 256L, + max_position_embeddings = config_raw$max_position_embeddings %||% 131072L, + rms_norm_eps = config_raw$rms_norm_eps %||% 1e-6, + rope_theta = config_raw$rope_theta %||% 1000000.0, # Gemma3 uses 1M + rope_scaling = list(factor = config_raw$rope_scaling$factor %||% 8.0), + sliding_window = config_raw$sliding_window %||% 1024L, + sliding_window_pattern = config_raw$sliding_window_pattern %||% 6L, + attn_logit_softcapping = config_raw$attn_logit_softcapping, # NULL = no softcapping (HF default) + query_pre_attn_scalar = config_raw$query_pre_attn_scalar %||% 256L + ) +} + #' Load Gemma3 Text Model from safetensors #' #' Loads pre-trained Gemma3 weights from HuggingFace safetensors files. +#' An NF4 artifact directory (from \code{\link{gemma3_quantize_nf4}}) +#' dispatches to \code{\link{load_gemma3_nf4}}. #' #' @param model_path Character. Path to directory containing model files. #' @param device Character. Device to load model to. @@ -562,6 +586,15 @@ gemma3_config_ltx2 <- function() { #' @export load_gemma3_text_encoder <- function(model_path, device = "cpu", dtype = "float16", verbose = TRUE) { + # An NF4 artifact directory (gemma3_quantize_nf4) loads through the + # quantized path; bf16 compute unless the caller says otherwise + if (file.exists(file.path(model_path, "manifest.json"))) { + if (identical(dtype, "float16")) { + dtype <- "bfloat16" + } + return(load_gemma3_nf4(model_path, device = device, dtype = dtype, + verbose = verbose)) + } # Load config config_path <- file.path(model_path, "config.json") if (!file.exists(config_path)) { @@ -575,23 +608,7 @@ load_gemma3_text_encoder <- function(model_path, device = "cpu", config_raw <- config_raw$text_config } - config <- list( - vocab_size = config_raw$vocab_size %||% 262208L, - hidden_size = config_raw$hidden_size %||% 3840L, - intermediate_size = config_raw$intermediate_size %||% 15360L, - num_hidden_layers = config_raw$num_hidden_layers %||% 48L, - num_attention_heads = config_raw$num_attention_heads %||% 16L, - num_key_value_heads = config_raw$num_key_value_heads %||% 8L, - head_dim = config_raw$head_dim %||% 256L, - max_position_embeddings = config_raw$max_position_embeddings %||% 131072L, - rms_norm_eps = config_raw$rms_norm_eps %||% 1e-6, - rope_theta = config_raw$rope_theta %||% 1000000.0, # Gemma3 uses 1M - rope_scaling = list(factor = config_raw$rope_scaling$factor %||% 8.0), - sliding_window = config_raw$sliding_window %||% 1024L, - sliding_window_pattern = config_raw$sliding_window_pattern %||% 6L, - attn_logit_softcapping = config_raw$attn_logit_softcapping, # NULL = no softcapping (HF default) - query_pre_attn_scalar = config_raw$query_pre_attn_scalar %||% 256L - ) + config <- .gemma3_build_config(config_raw) if (verbose) { message(sprintf("Creating Gemma3 model: %d layers, hidden_size=%d", @@ -866,3 +883,106 @@ if (!exists("%||%", mode = "function")) { x } } + +#' Batch-encode prompts with Gemma3, cached to disk +#' +#' Encodes a vector of prompts in sub-batches (bounding activation +#' VRAM) and optionally caches each prompt's result under +#' \code{cache_dir}, keyed by the prompt text and sequence length. +#' Already-cached prompts are skipped, so an interrupted batch resumes +#' where it stopped. Embeddings land on the CPU either way; the +#' renderer moves them per phase. +#' +#' Budget note: each prompt's result is the full hidden-state stack the +#' LTX connectors consume - roughly 0.4 GB at 1024 tokens - so caching +#' N prompts needs ~0.4N GB of disk. +#' +#' @param prompts Character vector. +#' @param model,tokenizer As in \code{\link{encode_with_gemma3}}. +#' @param batch_size Integer. Prompts per forward pass (default 4; +#' raise on cards with headroom, lower if the encode OOMs). +#' @param cache_dir Optional directory. When given, each prompt is +#' written to \code{gemma3-.pt} and the return value is the +#' character vector of those paths (load one with +#' \code{torch::torch_load}); when NULL, the return value is a list +#' of \code{list(prompt_embeds, prompt_attention_mask)} in prompt +#' order. +#' @param max_sequence_length,device,verbose As in +#' \code{\link{encode_with_gemma3}}. +#' +#' @return Character vector of cache paths (with \code{cache_dir}) or +#' a list of per-prompt embedding results (without). +#' +#' @export +gemma3_encode_batch <- function(prompts, model = NULL, tokenizer = NULL, + batch_size = 4L, cache_dir = NULL, + max_sequence_length = 1024L, device = "cuda", + verbose = TRUE) { + prompts <- as.character(prompts) + n <- length(prompts) + if (!n) { + stop("No prompts given") + } + + cache_paths <- NULL + if (!is.null(cache_dir)) { + dir.create(cache_dir, recursive = TRUE, showWarnings = FALSE) + cache_paths <- vapply(prompts, function(p) { + tf <- tempfile() + writeLines(c(p, as.character(max_sequence_length)), tf) + key <- unname(tools::md5sum(tf)) + unlink(tf) + file.path(cache_dir, paste0("gemma3-", key, ".pt")) + }, character(1), USE.NAMES = FALSE) + todo <- which(!file.exists(cache_paths)) + if (verbose && length(todo) < n) { + message(sprintf("%d/%d prompts already cached", n - length(todo), + n)) + } + } else { + todo <- seq_len(n) + } + + if (is.null(cache_dir)) { + results <- vector("list", n) + } else { + results <- NULL + } + done <- 0L + for (chunk in split(todo, ceiling(seq_along(todo) / batch_size))) { + enc <- encode_with_gemma3(prompts[chunk], model = model, + tokenizer = tokenizer, + max_sequence_length = max_sequence_length, + device = device, verbose = FALSE) + embeds <- enc$prompt_embeds$cpu() + mask <- enc$prompt_attention_mask$cpu() + for (j in seq_along(chunk)) { + # clone(), not contiguous(): a narrow of a contiguous + # tensor is already contiguous, so contiguous() keeps the + # storage-offset view - and torch_save serializes offset + # views from the base storage (every row would save as + # row 1). clone() forces fresh storage. + one <- list( + prompt_embeds = embeds$narrow(1L, j, 1L)$clone(), + prompt_attention_mask = mask$narrow(1L, j, 1L)$clone() + ) + if (is.null(cache_dir)) { + results[[chunk[j]]] <- one + } else { + torch::torch_save(one, cache_paths[chunk[j]]) + } + } + rm(enc, embeds, mask) + gc(verbose = FALSE) + done <- done + length(chunk) + if (verbose) { + message(sprintf(" encoded %d/%d prompts", done, length(todo))) + } + } + + if (is.null(cache_dir)) { + results + } else { + cache_paths + } +} diff --git a/R/nf4_ltx23.R b/R/nf4_ltx23.R index aad6133..cac2ab4 100644 --- a/R/nf4_ltx23.R +++ b/R/nf4_ltx23.R @@ -82,42 +82,32 @@ ltx23_nf4_dequantize <- function(packed, absmax, shape, if (is.null(out)) { out <- torch::torch_empty(shape, dtype = dtype, device = packed$device) } - out_flat <- out$view(-1L) + out_pairs <- out$view(c(-1L, 2L)) + out_blocks <- out$view(c(-1L, block)) n_bytes <- packed$shape[1] bytes_per_chunk <- max(chunk_elements %/% 2L, block) scratch <- .ltx23_get_dequant_scratch(min(bytes_per_chunk, n_bytes), packed$device) + # Byte LUT [256, 2] in the target dtype (shared with the jit block + # stack): one 1-based index per PACKED byte and a single gather of + # (hi, lo) pairs written straight into out - no nibble unpack, no + # per-value int64 index, no float32 staging copy + tbl <- .ltx23_jit_table(packed$device, dtype) + am <- absmax$to(dtype = dtype) start <- 1L torch::with_no_grad({ while (start <= n_bytes) { len <- min(bytes_per_chunk, n_bytes - start + 1L) - chunk <- packed$narrow(1L, start, len) - - # Fully in-place nibble unpack into persistent scratch: - # hi = byte %/% 16, lo = byte - 16 * hi - hi <- scratch$hi$narrow(1L, 1L, len) - lo <- scratch$lo$narrow(1L, 1L, len) - hi$copy_(chunk)$div_(16L, rounding_mode = "floor") - lo$copy_(hi)$mul_(-16L)$add_(chunk) - - # Interleave into the (1-based) int64 index scratch - idx <- scratch$idx$narrow(1L, 1L, len * 2L) - idx_pairs <- idx$view(c(-1L, 2L)) - idx_pairs$narrow(2L, 1L, 1L)$squeeze(2L)$copy_(hi) - idx_pairs$narrow(2L, 2L, 1L)$squeeze(2L)$copy_(lo) - idx$add_(1L) - - vals <- scratch$vals$narrow(1L, 1L, len * 2L) - .ltx23_index_select_into(vals, scratch$table, idx) + idx <- scratch$idx$narrow(1L, 1L, len) + idx$copy_(packed$narrow(1L, start, len))$add_(1L) + .ltx23_index_select_into(out_pairs$narrow(1L, start, len), tbl, idx) block_start <- ((start - 1L) * 2L) %/% block + 1L n_blocks <- (len * 2L) %/% block - scales <- absmax$narrow(1L, block_start, n_blocks) - vals$view(c(-1L, block))$mul_(scales$unsqueeze(2L)) - - out_flat$narrow(1L, (start - 1L) * 2L + 1L, len * 2L)$copy_(vals) + out_blocks$narrow(1L, block_start, n_blocks)$ + mul_(am$narrow(1L, block_start, n_blocks)$unsqueeze(2L)) start <- start + len } }) @@ -159,18 +149,7 @@ ltx23_nf4_dequantize <- function(packed, absmax, shape, if (is.null(scratch) || scratch$n_bytes < n_bytes) { scratch <- list( n_bytes = n_bytes, - hi = torch::torch_empty(n_bytes, dtype = torch::torch_uint8(), - device = device), - lo = torch::torch_empty(n_bytes, dtype = torch::torch_uint8(), - device = device), - idx = torch::torch_empty(n_bytes * 2L, - dtype = torch::torch_long(), - device = device), - vals = torch::torch_empty(n_bytes * 2L, - dtype = torch::torch_float32(), - device = device), - table = torch::torch_tensor(.ltx23_nf4_table, - dtype = torch::torch_float32(), + idx = torch::torch_empty(n_bytes, dtype = torch::torch_long(), device = device) ) .ltx23_dequant_scratch[[key]] <- scratch diff --git a/R/quantize_gemma3.R b/R/quantize_gemma3.R new file mode 100644 index 0000000..d1d9371 --- /dev/null +++ b/R/quantize_gemma3.R @@ -0,0 +1,215 @@ +#' Quantize a Gemma3 text encoder to NF4 shards +#' +#' Streams the HuggingFace Gemma3 checkpoint tensor by tensor. The +#' language model's projection weights (q/k/v/o and gate/up/down, ~11B +#' of the 12B parameters) are stored as NF4 (packed uint8 + +#' \code{_absmax} float32 blocks); embeddings and norms are copied +#' at the resident dtype. Vision-tower and projector weights are +#' dropped - the text encoder never uses them. The result is a ~8 GB +#' artifact that fits a 16 GB card during the encode phase (vs 45 GB +#' of host RAM for the fp32 CPU path). +#' +#' Keys are stored normalized (\code{language_model.} / \code{model.} +#' prefixes stripped), matching the module tree of +#' \code{\link{gemma3_text_model}}. +#' +#' @param model_path HuggingFace snapshot directory (config.json + +#' model-*.safetensors). +#' @param output_dir Output directory for shards + manifest (default: +#' \code{gemma3-nf4} under \code{tools::R_user_dir}). +#' @param shard_bytes Numeric. Target shard size in bytes; the 1.9e9 +#' default keeps shards readable by stock CRAN safetensors. +#' @param force Logical. Re-quantize even if a valid manifest exists. +#' @param verbose Logical. +#' +#' @return Invisibly, the manifest list. +#' +#' @export +gemma3_quantize_nf4 <- function(model_path, output_dir = NULL, + shard_bytes = 1.9e9, force = FALSE, + verbose = TRUE) { + model_path <- path.expand(model_path) + if (is.null(output_dir)) { + output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + "gemma3-nf4") + } + + manifest_path <- file.path(output_dir, "manifest.json") + if (!force && file.exists(manifest_path)) { + manifest <- jsonlite::fromJSON(manifest_path) + if (identical(manifest$format, "nf4") && + all(file.exists(file.path(output_dir, manifest$shards)))) { + if (verbose) { + message("NF4 artifact already present: ", output_dir) + } + return(invisible(manifest)) + } + } + + config_path <- file.path(model_path, "config.json") + if (!file.exists(config_path)) { + stop("Config file not found: ", config_path) + } + config_raw <- jsonlite::fromJSON(config_path) + if (!is.null(config_raw$text_config)) { + config_raw <- config_raw$text_config + } + + resident_dtype <- if (.st_can_write("bfloat16")) { + torch::torch_bfloat16() + } else { + message("safetensors cannot write bfloat16; storing resident ", + "tensors as float32 (larger artifact, same results)") + torch::torch_float32() + } + + opened <- .flux_open_sharded_dir(model_path, "model") + dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) + + shard <- list() + shard_size <- 0 + shard_files <- character(0) + n_cast <- 0L + + flush_shard <- function() { + if (!length(shard)) { + return() + } + fname <- sprintf("gemma3-nf4-%05d.safetensors", + length(shard_files) + 1L) + safetensors::safe_save_file(shard, file.path(output_dir, fname)) + shard_files[[length(shard_files) + 1L]] <<- fname + if (verbose) { + message(sprintf(" wrote %s (%.2f GB, %d tensors)", fname, + shard_size / 1e9, length(shard))) + } + shard <<- list() + shard_size <<- 0 + gc(verbose = FALSE) + } + + for (key in opened$keys) { + if (grepl("^(vision_tower|multi_modal_projector|lm_head)\\.", key)) { + next + } + norm_key <- sub("^model\\.", "", sub("^language_model\\.", "", key)) + tensor <- opened$handle$get_tensor(key) + + if (.gemma3_is_quant_key(norm_key)) { + torch::with_no_grad({ + q <- ltx23_nf4_quantize(tensor) + shard[[norm_key]] <- q$packed + shard[[paste0(norm_key, "_absmax")]] <- q$absmax + }) + shard_size <- shard_size + prod(tensor$shape) * 0.5625 + n_cast <- n_cast + 1L + } else { + shard[[norm_key]] <- tensor$to(dtype = resident_dtype) + shard_size <- shard_size + prod(tensor$shape) * 2 + } + rm(tensor) + if (shard_size >= shard_bytes) { + flush_shard() + } + } + flush_shard() + + manifest <- list(model = "gemma3", format = "nf4", shards = shard_files, + n_cast = n_cast, config = config_raw) + jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE, + pretty = TRUE) + if (verbose) { + message(sprintf("Quantized %d weights to NF4 across %d shards: %s", + n_cast, length(shard_files), output_dir)) + } + invisible(manifest) +} + +# The 7 per-layer projection weights carrying ~11B of the 12B params +.gemma3_is_quant_key <- function(key) { + grepl(paste0("^layers\\.[0-9]+\\.(self_attn\\.(q|k|v|o)_proj|", + "mlp\\.(gate|up|down)_proj)\\.weight$"), key) +} + +#' Load a Gemma3 text encoder from an NF4 artifact +#' +#' Builds the model as a skeleton at the compute dtype, swaps the +#' projection linears for NF4 modules filled from the artifact +#' (dequantized per forward through the shared byte-LUT), copies the +#' residents, and hard-errors on any parameter the artifact does not +#' fill. +#' +#' @param artifact_dir Directory produced by +#' \code{\link{gemma3_quantize_nf4}}. +#' @param device "cuda" or "cpu". +#' @param dtype Compute dtype ("bfloat16" default). +#' @param verbose Logical. +#' +#' @return A \code{gemma3_text_model} ready for +#' \code{\link{encode_with_gemma3}}. +#' +#' @export +load_gemma3_nf4 <- function(artifact_dir, device = "cuda", + dtype = "bfloat16", verbose = TRUE) { + ckpt <- ltx23_open_fp8_checkpoint(artifact_dir) + manifest <- jsonlite::fromJSON(file.path(artifact_dir, "manifest.json")) + if (!identical(manifest$model, "gemma3")) { + stop("Not a Gemma3 NF4 artifact (manifest model is ", + manifest$model %||% "absent", "): ", artifact_dir) + } + config <- .gemma3_build_config(manifest$config) + + torch_dtype <- switch(dtype, + "float32" = torch::torch_float32(), + "float16" = torch::torch_float16(), + "bfloat16" = torch::torch_bfloat16(), + torch::torch_bfloat16()) + + if (verbose) { + message(sprintf("Creating Gemma3 NF4 model: %d layers, hidden_size=%d", + config$num_hidden_layers, config$hidden_size)) + } + model <- .construct_skeleton(gemma3_text_model, config, dtype = torch_dtype) + + cast_keys <- grep("_absmax$", ckpt$keys, value = TRUE) + cast_keys <- sub("_absmax$", "", cast_keys) + torch::with_no_grad({ + for (key in cast_keys) { + segs <- strsplit(sub("\\.weight$", "", key), ".", fixed = TRUE)[[1]] + # layers... + layer <- model$layers[[as.integer(segs[2]) + 1L]] + parent <- layer[[segs[3]]] + old <- parent[[segs[4]]] + nf4 <- ltx23_nf4_linear(old$weight$shape[1], old$weight$shape[2], + bias = FALSE) + nf4$set_nf4_weight(ckpt$handle$get_tensor(key), + ckpt$handle$get_tensor(paste0(key, "_absmax"))) + do.call(`$<-`, list(parent, segs[4], nf4)) + } + + dests <- model$parameters + filled <- character(0) + for (key in setdiff(ckpt$keys, c(cast_keys, + paste0(cast_keys, "_absmax")))) { + dest <- dests[[key]] + if (is.null(dest)) { + next + } + dest$copy_(ckpt$handle$get_tensor(key)) + filled <- c(filled, key) + } + unfilled <- setdiff(names(dests), filled) + if (length(unfilled)) { + stop("Gemma3 NF4 load: ", length(unfilled), + " parameters not filled, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", ")) + } + }) + + model$to(device = device) + model$eval() + if (verbose) { + message("Gemma3 NF4 encoder ready on ", device) + } + model +} diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 5829d0b..44106f1 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -852,10 +852,13 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, if (!is.null(filename) && decode_video) { phase_t0 <- Sys.time() - save_video_ltx23(result$video, filename, + # [[ ]] not $: with decode_audio = FALSE there is no result$audio + # and partial matching silently hands write_wav the raw + # audio_latents tensor + save_video_ltx23(result[["video"]], filename, fps = frame_rate, - audio = result$audio, - sample_rate = result$sample_rate, + audio = result[["audio"]], + sample_rate = result[["sample_rate"]], verbose = verbose ) result$filename <- filename diff --git a/inst/tinytest/test_gemma3_batch.R b/inst/tinytest/test_gemma3_batch.R new file mode 100644 index 0000000..8689993 --- /dev/null +++ b/inst/tinytest/test_gemma3_batch.R @@ -0,0 +1,112 @@ +# gemma3_encode_batch: sub-batched encoding matches solo encodes, disk +# caching returns paths, resumes without re-encoding, and round-trips +# through torch_load. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +config <- gemma3_config_ltx2() +for (nm in names(cfg <- list( + vocab_size = 64L, hidden_size = 16L, intermediate_size = 32L, + num_hidden_layers = 2L, num_attention_heads = 2L, + num_key_value_heads = 1L, head_dim = 8L, query_pre_attn_scalar = 8L, + sliding_window = 4L, sliding_window_pattern = 2L +))) config[[nm]] <- cfg[[nm]] +torch::torch_manual_seed(7) +model <- gemma3_text_model(config) +model$eval() + +# A stub tokenizer: encode_with_gemma3 only needs tokenize_gemma3 to +# work, which needs ids within vocab; use the real path via a fake +# minimal bpe? Simpler: call encode_with_gemma3 through its tokens by +# monkeying is overkill - instead drive the batch helper with a real +# tokenizer-free path: fake tokenizer via the exported tokenize seam is +# not available, so tokenize through gemma3_tokenizer would need vocab +# files. Drive the model directly instead for parity, and the batch +# helper end-to-end with a trivial tokenizer object. + +# Minimal tokenizer standing in for gemma3_tokenizer: the helper only +# passes it through to encode_with_gemma3 -> tokenize_gemma3, which +# calls encode_bpe; emulate the contract with a closure-based double. +fake_tok <- structure(list(), class = "fake_tok") +tokenize_real <- diffuseR:::tokenize_gemma3 + +# Patch tokenize_gemma3 inside the namespace for this test: map each +# prompt deterministically to ids from its characters (padded left) +fake_tokenize <- function(tokenizer, prompts, max_length = 16L, + padding = "max_length") { + n <- length(prompts) + ids <- torch::torch_zeros(n, max_length, dtype = torch::torch_long()) + mask <- torch::torch_zeros(n, max_length) + for (i in seq_len(n)) { + v <- utf8ToInt(prompts[[i]]) %% 60L + 1L + v <- utils::tail(v, max_length) + ids[i, (max_length - length(v) + 1L):max_length] <- + torch::torch_tensor(v, dtype = torch::torch_long()) + mask[i, (max_length - length(v) + 1L):max_length] <- 1 + } + list(input_ids = ids, attention_mask = mask) +} +# No on.exit here: at test-file top level tinytest evaluates each +# expression in its own frame, so on.exit would fire immediately and +# restore the real tokenizer before the calls below. Restored at the +# end of the file instead. +assignInNamespace("tokenize_gemma3", fake_tokenize, ns = "diffuseR") + +prompts <- c("a cat on a mat", "neon fog rolling", "a robot dancing") + +# --- no cache: list results match solo encodes ------------------------------------- + +batch <- gemma3_encode_batch(prompts, model = model, tokenizer = fake_tok, + batch_size = 2L, max_sequence_length = 16L, + device = "cpu", verbose = FALSE) +expect_equal(length(batch), 3L) +solo <- encode_with_gemma3(prompts[2], model = model, tokenizer = fake_tok, + max_sequence_length = 16L, device = "cpu", + verbose = FALSE) +expect_equal(as.integer(batch[[2]]$prompt_embeds$shape), + as.integer(solo$prompt_embeds$shape)) +expect_true(as.numeric((batch[[2]]$prompt_embeds - + solo$prompt_embeds)$abs()$max()) < 1e-5) +expect_true(as.numeric((batch[[2]]$prompt_attention_mask - + solo$prompt_attention_mask)$abs()$max()) == 0) + +# --- cache: paths returned, resume skips, round-trip matches ----------------------- + +cache <- file.path(tempdir(), "gemma3-batch-cache") +unlink(cache, recursive = TRUE) +paths <- gemma3_encode_batch(prompts, model = model, tokenizer = fake_tok, + batch_size = 2L, cache_dir = cache, + max_sequence_length = 16L, device = "cpu", + verbose = FALSE) +expect_equal(length(paths), 3L) +expect_true(all(file.exists(paths))) + +loaded <- torch::torch_load(paths[2]) +expect_true(as.numeric((loaded$prompt_embeds - + batch[[2]]$prompt_embeds)$abs()$max()) < 1e-6) + +# Resume: files untouched on the second call (no re-encode) +mt <- file.mtime(paths) +paths2 <- gemma3_encode_batch(prompts, model = model, tokenizer = fake_tok, + batch_size = 2L, cache_dir = cache, + max_sequence_length = 16L, device = "cpu", + verbose = FALSE) +expect_equal(paths2, paths) +expect_equal(file.mtime(paths), mt) + +# A new prompt encodes into the same cache without touching the others +paths3 <- gemma3_encode_batch(c(prompts, "one more scene"), model = model, + tokenizer = fake_tok, batch_size = 2L, + cache_dir = cache, + max_sequence_length = 16L, device = "cpu", + verbose = FALSE) +expect_equal(length(paths3), 4L) +expect_equal(paths3[1:3], paths) +expect_true(file.exists(paths3[4])) + +unlink(cache, recursive = TRUE) +assignInNamespace("tokenize_gemma3", tokenize_real, ns = "diffuseR") diff --git a/inst/tinytest/test_gemma3_nf4.R b/inst/tinytest/test_gemma3_nf4.R new file mode 100644 index 0000000..df3068e --- /dev/null +++ b/inst/tinytest/test_gemma3_nf4.R @@ -0,0 +1,112 @@ +# gemma3_quantize_nf4 + load_gemma3_nf4: round-trip through a tiny +# on-disk HF-layout checkpoint, projection swap to NF4 modules, forward +# parity within NF4 quantization tolerance, dispatch through +# load_gemma3_text_encoder, and the unfilled-parameter hard error. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +config <- list( + vocab_size = 64L, hidden_size = 16L, intermediate_size = 32L, + num_hidden_layers = 2L, num_attention_heads = 2L, + num_key_value_heads = 1L, head_dim = 8L, query_pre_attn_scalar = 8L, + sliding_window = 4L, sliding_window_pattern = 2L +) + +ref_config <- gemma3_config_ltx2() +for (nm in names(config)) ref_config[[nm]] <- config[[nm]] +torch::torch_manual_seed(23) +ref <- gemma3_text_model(ref_config) +ref$eval() + +# HF-layout snapshot: config.json, one shard, and the weight-map index +# the quantizer's sharded-dir opener requires +src <- file.path(tempdir(), "gemma3-nf4-src") +dir.create(src, showWarnings = FALSE) +jsonlite::write_json(config, file.path(src, "config.json"), auto_unbox = TRUE) +params <- ref$parameters +weights <- stats::setNames( + lapply(params, function(p) p$detach()$contiguous()), + paste0("model.", names(params)) +) +shard_name <- "model-00001-of-00001.safetensors" +safetensors::safe_save_file(weights, file.path(src, shard_name)) +jsonlite::write_json( + list(weight_map = stats::setNames( + as.list(rep(shard_name, length(weights))), names(weights))), + file.path(src, "model.safetensors.index.json"), auto_unbox = TRUE +) + +# --- quantize ----------------------------------------------------------------------- + +art <- file.path(tempdir(), "gemma3-nf4-art") +unlink(art, recursive = TRUE) +# Pin residents to float32 so the exactness assertions below are +# hermetic across CRAN (no bf16 write) and fork safetensors +options(diffuseR.st_caps = list(bfloat16 = FALSE)) +manifest <- gemma3_quantize_nf4(src, art, verbose = FALSE) +options(diffuseR.st_caps = NULL) +expect_equal(manifest$model, "gemma3") +# 2 layers x 7 projections +expect_equal(manifest$n_cast, 14L) +expect_true(file.exists(file.path(art, "manifest.json"))) + +# --- load + forward parity ---------------------------------------------------------- + +enc <- load_gemma3_nf4(art, device = "cpu", dtype = "float32", + verbose = FALSE) +expect_true(inherits(enc$layers[[1]]$self_attn$q_proj, "ltx23_nf4_linear")) +expect_true(inherits(enc$layers[[2]]$mlp$down_proj, "ltx23_nf4_linear")) +# Residents are exact +expect_true(as.numeric((enc$parameters[["embed_tokens.weight"]] - + params[["embed_tokens.weight"]])$abs()$max()) == 0) + +ids <- torch::torch_randint(1L, 64L, size = c(1L, 5L), dtype = torch::torch_long()) +mask <- torch::torch_ones(1L, 5L) +torch::with_no_grad({ + out_ref <- ref(ids, attention_mask = mask, output_hidden_states = TRUE) + out_nf4 <- enc(ids, attention_mask = mask, output_hidden_states = TRUE) +}) +h_ref <- out_ref$last_hidden_state +h_nf4 <- out_nf4$last_hidden_state +rel <- as.numeric((h_nf4 - h_ref)$abs()$max()) / + as.numeric(h_ref$abs()$max()) +# 4-bit weights: outputs agree to quantization noise, not float noise. +# On this tiny 2-layer model the worst element sits ~0.15 relative +# (measured 0.1505); the direction check below is the real quality +# gate, and full-model embedding parity is validated on GPU. +expect_true(rel < 0.25, info = sprintf("relative max diff %.4f", rel)) +cs <- as.numeric(torch::nnf_cosine_similarity( + h_ref$flatten()$unsqueeze(1L), h_nf4$flatten()$unsqueeze(1L))) +expect_true(cs > 0.99, info = sprintf("cosine %.5f", cs)) + +# --- dispatch through load_gemma3_text_encoder -------------------------------------- + +enc2 <- load_gemma3_text_encoder(art, device = "cpu", dtype = "float32", + verbose = FALSE) +expect_true(inherits(enc2$layers[[1]]$self_attn$q_proj, "ltx23_nf4_linear")) + +# --- an artifact missing a resident is a hard error --------------------------------- + +art2 <- file.path(tempdir(), "gemma3-nf4-art2") +unlink(art2, recursive = TRUE) +dir.create(art2) +m2 <- jsonlite::fromJSON(file.path(art, "manifest.json")) +for (s in m2$shards) { + w <- safetensors::safe_load_file(file.path(art, s), framework = "torch") + w[["norm.weight"]] <- NULL + safetensors::safe_save_file(w, file.path(art2, s)) +} +jsonlite::write_json(m2, file.path(art2, "manifest.json"), auto_unbox = TRUE) +expect_error( + load_gemma3_nf4(art2, device = "cpu", verbose = FALSE), + "not filled" +) + +unlink(c(src, art, art2), recursive = TRUE) diff --git a/man/gemma3_encode_batch.Rd b/man/gemma3_encode_batch.Rd new file mode 100644 index 0000000..f8a897d --- /dev/null +++ b/man/gemma3_encode_batch.Rd @@ -0,0 +1,45 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma3_encode_batch} +\alias{gemma3_encode_batch} +\title{Batch-encode prompts with Gemma3, cached to disk} +\usage{ +gemma3_encode_batch(prompts, model = NULL, tokenizer = NULL, batch_size = 4L, + cache_dir = NULL, max_sequence_length = 1024L, + device = "cuda", verbose = TRUE) +} +\arguments{ +\item{prompts}{Character vector.} + +\item{batch_size}{Integer. Prompts per forward pass (default 4; +raise on cards with headroom, lower if the encode OOMs).} + +\item{cache_dir}{Optional directory. When given, each prompt is +written to \code{gemma3-.pt} and the return value is the +character vector of those paths (load one with +\code{torch::torch_load}); when NULL, the return value is a list +of \code{list(prompt_embeds, prompt_attention_mask)} in prompt +order.} + +\item{model,tokenizer}{As in \code{\link{encode_with_gemma3}}.} + +\item{max_sequence_length,device,verbose}{As in +\code{\link{encode_with_gemma3}}.} +} +\value{ +Character vector of cache paths (with \code{cache_dir}) or + a list of per-prompt embedding results (without). +} +\description{ +Encodes a vector of prompts in sub-batches (bounding activation +VRAM) and optionally caches each prompt's result under +\code{cache_dir}, keyed by the prompt text and sequence length. +Already-cached prompts are skipped, so an interrupted batch resumes +where it stopped. Embeddings land on the CPU either way; the +renderer moves them per phase. +} +\details{ +Budget note: each prompt's result is the full hidden-state stack the +LTX connectors consume - roughly 0.4 GB at 1024 tokens - so caching +N prompts needs ~0.4N GB of disk. + +} diff --git a/man/gemma3_quantize_nf4.Rd b/man/gemma3_quantize_nf4.Rd new file mode 100644 index 0000000..f4c6817 --- /dev/null +++ b/man/gemma3_quantize_nf4.Rd @@ -0,0 +1,41 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma3_quantize_nf4} +\alias{gemma3_quantize_nf4} +\title{Quantize a Gemma3 text encoder to NF4 shards} +\usage{ +gemma3_quantize_nf4(model_path, output_dir = NULL, shard_bytes = 1.9e+09, + force = FALSE, verbose = TRUE) +} +\arguments{ +\item{model_path}{HuggingFace snapshot directory (config.json + +model-*.safetensors).} + +\item{output_dir}{Output directory for shards + manifest (default: +\code{gemma3-nf4} under \code{tools::R_user_dir}).} + +\item{shard_bytes}{Numeric. Target shard size in bytes; the 1.9e9 +default keeps shards readable by stock CRAN safetensors.} + +\item{force}{Logical. Re-quantize even if a valid manifest exists.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the manifest list. +} +\description{ +Streams the HuggingFace Gemma3 checkpoint tensor by tensor. The +language model's projection weights (q/k/v/o and gate/up/down, ~11B +of the 12B parameters) are stored as NF4 (packed uint8 + +\code{_absmax} float32 blocks); embeddings and norms are copied +at the resident dtype. Vision-tower and projector weights are +dropped - the text encoder never uses them. The result is a ~8 GB +artifact that fits a 16 GB card during the encode phase (vs 45 GB +of host RAM for the fp32 CPU path). +} +\details{ +Keys are stored normalized (\code{language_model.} / \code{model.} +prefixes stripped), matching the module tree of +\code{\link{gemma3_text_model}}. + +} diff --git a/man/load_gemma3_nf4.Rd b/man/load_gemma3_nf4.Rd new file mode 100644 index 0000000..f6d7b84 --- /dev/null +++ b/man/load_gemma3_nf4.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_gemma3_nf4} +\alias{load_gemma3_nf4} +\title{Load a Gemma3 text encoder from an NF4 artifact} +\usage{ +load_gemma3_nf4(artifact_dir, device = "cuda", dtype = "bfloat16", + verbose = TRUE) +} +\arguments{ +\item{artifact_dir}{Directory produced by +\code{\link{gemma3_quantize_nf4}}.} + +\item{device}{"cuda" or "cpu".} + +\item{dtype}{Compute dtype ("bfloat16" default).} + +\item{verbose}{Logical.} +} +\value{ +A \code{gemma3_text_model} ready for + \code{\link{encode_with_gemma3}}. +} +\description{ +Builds the model as a skeleton at the compute dtype, swaps the +projection linears for NF4 modules filled from the artifact +(dequantized per forward through the shared byte-LUT), copies the +residents, and hard-errors on any parameter the artifact does not +fill. +} diff --git a/man/load_gemma3_text_encoder.Rd b/man/load_gemma3_text_encoder.Rd index bffba3d..1ee2d4c 100644 --- a/man/load_gemma3_text_encoder.Rd +++ b/man/load_gemma3_text_encoder.Rd @@ -20,4 +20,6 @@ Initialized gemma3_text_model with loaded weights. } \description{ Loads pre-trained Gemma3 weights from HuggingFace safetensors files. +An NF4 artifact directory (from \code{\link{gemma3_quantize_nf4}}) +dispatches to \code{\link{load_gemma3_nf4}}. }