From b7d9b00942eb2a4f43879cdb8a65ca30a0932f9f Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 29 Oct 2026 16:20:16 -0500 Subject: [PATCH 1/2] FLUX 3 in less host RAM: fp8 Qwen3-VL text encoder artifact --- CLAUDE.md | 10 ++ NAMESPACE | 1 + R/download_flux3.R | 19 ++- R/quantize_flux.R | 22 ++- R/quantize_qwen3.R | 243 ++++++++++++++++++++++++++++ R/qwen3_text_encoder.R | 66 ++++---- R/stream_flux3.R | 3 +- R/txt2img_flux3.R | 43 ++++- inst/tinytest/test_qwen3_quantize.R | 98 +++++++++++ man/download_flux3_klein.Rd | 11 +- man/flux3_load_pipeline.Rd | 6 + man/load_qwen3_text_encoder.Rd | 9 +- man/qwen3_quantize.Rd | 50 ++++++ 13 files changed, 534 insertions(+), 47 deletions(-) create mode 100644 R/quantize_qwen3.R create mode 100644 inst/tinytest/test_qwen3_quantize.R create mode 100644 man/qwen3_quantize.Rd diff --git a/CLAUDE.md b/CLAUDE.md index 758c5ef..c2fbff5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -479,6 +479,16 @@ ballast tensor on the 16 GB card: 8 GB runs 1024^2 t2i in ~30 s and a 1 MP edit in ~56 s; 12 GB (no streaming) ~27 s / ~53 s. 16 GB: t2i ~27 s (peak 9.1 GB), 1 MP edit ~53 s (10.7 GB). +Host RAM: `download_flux3_klein()` also quantizes the Qwen3-VL text +model to fp8 (`qwen3_quantize()`, 4.4 GB artifact +`flux3-klein-qwen3vl-fp8`), which `flux3_load_pipeline(text_precision += "auto")` picks up. Host high-water at 1024^2: bf16 text 20.2 GB, fp8 +15.5 GB, NF4 13.5 GB; same speed. fp8 text keeps the bf16 image (27 dB +PSNR, embeddings 2.7% rel. error); NF4 text (5%) changes the +composition for the same seed, so it is opt-in. Pinned host memory +rounds up in torch's CachingHostAllocator: RssShmem is 13.2 GB for +~9.7 GB of fp8 weights. + Gotchas that bit once (gc on small cards): torch for R runs R's gc only on an allocator miss, and not at all while reserved memory is below `torch.cuda_allocator_reserved_rate` x total; `ltx23_tune_gc` sets diff --git a/NAMESPACE b/NAMESPACE index 8fb40fb..4ccbece 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -208,6 +208,7 @@ export(preprocess_image) export(quant_conv) export(qwen_bpe_tokenizer) export(qwen3_encoder) +export(qwen3_quantize) export(recommend) export(reshard_safetensors) export(resident_activate) diff --git a/R/download_flux3.R b/R/download_flux3.R index 0950f70..84a6171 100644 --- a/R/download_flux3.R +++ b/R/download_flux3.R @@ -74,17 +74,26 @@ NULL #' @param output_dir Directory for the quantized artifact. #' @param text_encoders Logical. Also fetch the Qwen3-VL text encoder, #' tokenizer and VAE (~9.2 GB). +#' @param text_precision Quantize the Qwen3-VL text model too (with +#' \code{quantize}): "fp8" (default; a 4.4 GB artifact, host memory at +#' 1024x1024 20.2 -> 15.5 GB, images close to bf16 text: 27 dB PSNR), +#' "nf4" (2.8 GB, 13.5 GB host, but a visibly different image for the +#' same seed), or "bf16" (keep only the cached source). +#' \code{\link{flux3_load_pipeline}} picks the artifact up. #' @param verbose Logical. #' #' @return Invisibly, a list with \code{transformer_dir}, -#' \code{artifact_dir}, and \code{support} (named file paths). +#' \code{artifact_dir}, \code{support} (named file paths), and +#' \code{text_artifact_dir} when the text model was quantized. #' #' @export download_flux3_klein <- function(quantize = TRUE, precision = c("auto", "fp8", "nf4"), output_dir = NULL, text_encoders = TRUE, + text_precision = c("fp8", "nf4", "bf16"), verbose = TRUE) { precision <- match.arg(precision) + text_precision <- match.arg(text_precision) precision <- .flux_resolve_precision(precision, .flux3_artifact_prefix()) precision <- .st_graceful_precision(precision, mode = "write") output_dir <- output_dir %||% paste0(.flux3_artifact_prefix(), precision) @@ -151,6 +160,14 @@ download_flux3_klein <- function(quantize = TRUE, } result$support <- vapply(.flux3_support_files, .flux3_download, character(1)) + if (quantize && text_precision != "bf16") { + text_precision <- .st_graceful_precision(text_precision, + mode = "write") + result$text_artifact_dir <- .flux3_text_artifact(text_precision) + qwen3_quantize(dirname(result$support[["text_encoder/config.json"]]), + result$text_artifact_dir, format = text_precision, + verbose = verbose) + } } invisible(result) diff --git a/R/quantize_flux.R b/R/quantize_flux.R index 29bb261..4b27cc3 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -228,6 +228,15 @@ NULL args } +# float8_e4m3fn weight with a per-tensor absmax/448 float32 scale +.fp8_quantize <- function(tensor) { + scale <- tensor$abs()$max()$to(dtype = torch::torch_float32())$ + clamp(min = 1e-12) / 448 + list(weight = (tensor$to(dtype = torch::torch_float32()) / scale)$ + to(dtype = torch::torch_float8_e4m3fn()), + scale = scale) +} + # Move plain-field fp8 weights (and their scales) onto a device; used # for resident fp8 where the whole quantized model fits on the GPU .flux_fp8_to_device <- function(module, device) { @@ -294,6 +303,9 @@ NULL } else { .flux3_fp8_fields(pipe$fp8_split$resident) } + } else if (identical(nm, "text_encoder")) { + # An fp8 Qwen3 artifact's weights (empty list otherwise) + .flux_fp8_collect(module) } else { NULL } @@ -404,9 +416,6 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, ckpt <- flux_open_checkpoint(transformer_dir) hooks <- .flux_family_hooks(ckpt$config) - if (format == "fp8") { - fp8 <- torch::torch_float8_e4m3fn() - } shard <- list() shard_size <- 0 @@ -443,10 +452,9 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, shard[[paste0(key, "_absmax")]] <- q$absmax shard_size <- shard_size + prod(tensor$shape) * 0.5625 } else { - scale <- tensor$abs()$max()$to(dtype = torch::torch_float32())$ - clamp(min = 1e-12) / 448 - shard[[key]] <- (tensor$to(dtype = torch::torch_float32()) / scale)$to(dtype = fp8) - shard[[paste0(key, "_scale")]] <- scale + q <- .fp8_quantize(tensor) + shard[[key]] <- q$weight + shard[[paste0(key, "_scale")]] <- q$scale shard_size <- shard_size + prod(tensor$shape) } }) diff --git a/R/quantize_qwen3.R b/R/quantize_qwen3.R new file mode 100644 index 0000000..b175e91 --- /dev/null +++ b/R/quantize_qwen3.R @@ -0,0 +1,243 @@ +#' Quantize a Qwen3 text encoder to fp8 or NF4 shards +#' +#' Streams a transformers Qwen3 (FLUX.2 klein, Z-Image) or Qwen3-VL +#' (FLUX 3) text encoder tensor by tensor. The per-layer projection +#' weights (q/k/v/o and gate/up/down, 252 tensors in the 36-layer 4B) +#' are stored as float8_e4m3fn with a per-tensor absmax/448 +#' \code{_scale}, or as NF4 (packed uint8 + \code{_absmax}); +#' embeddings and norms are copied at bfloat16. The vision tower and the +#' tied LM head are dropped, and keys are written in the +#' \code{\link{qwen3_encoder}} layout (\code{model.layers.*}). The +#' manifest embeds the text config, so the loader needs no other file. +#' +#' Host memory: the 4B text model is ~8 GB in bfloat16, ~4.4 GB as fp8 +#' and ~2.8 GB as NF4 (the 0.8 GB embedding stays bfloat16 in both). +#' Load the artifact with \code{\link{load_qwen3_text_encoder}} on its +#' directory. +#' +#' @param model_path Source directory with \code{config.json} and +#' \code{model*.safetensors}. +#' @param output_dir Output directory for shards + manifest. +#' @param format "fp8" or "nf4". +#' @param shard_bytes Numeric. Target shard size in bytes (1.9e9 keeps +#' shards readable by safetensors older than 0.3.0). +#' @param force Logical. Re-quantize even if a valid manifest exists. +#' @param verbose Logical. +#' +#' @return Invisibly, the manifest list. +#' +#' @export +qwen3_quantize <- function(model_path, output_dir, format = c("fp8", "nf4"), + shard_bytes = 1.9e9, force = FALSE, + verbose = TRUE) { + format <- match.arg(format) + if (format == "fp8" && !.st_can_write("float8_e4m3fn")) { + stop("The installed safetensors package cannot write float8 ", + "tensors (released in safetensors 0.3.0). Run ", + "install.packages(\"safetensors\") to update, or use ", + "format = \"nf4\".", call. = FALSE) + } + manifest_path <- file.path(output_dir, "manifest.json") + if (!force && .is_qwen3_artifact(output_dir, format)) { + if (verbose) { + message(toupper(format), " Qwen3 artifact already present: ", + output_dir) + } + return(invisible(jsonlite::fromJSON(manifest_path, + simplifyVector = FALSE))) + } + + model_path <- path.expand(model_path) + config <- jsonlite::fromJSON(file.path(model_path, "config.json"), + simplifyVector = FALSE) + vl <- !is.null(config$text_config) + text_config <- if (vl) config$text_config else 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() + } + resident_bytes <- torch::torch_empty(1L, dtype = resident_dtype)$element_size() + dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) + # A stale manifest must not validate a half-written artifact + unlink(manifest_path) + + opened <- .flux_open_sharded_dir(model_path, "model") + shard <- list() + shard_size <- 0 + shard_files <- character(0) + n_cast <- 0L + flush_shard <- function() { + if (!length(shard)) { + return() + } + fname <- sprintf("qwen3-%s-%05d.safetensors", format, + 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) { + name <- .qwen3_map_key(key, vl) + if (is.na(name)) { + next + } + tensor <- opened$handle$get_tensor(key) + quant <- .qwen3_is_quant_key(name) + per_elem <- if (!quant) { + resident_bytes + } else if (format == "fp8") { + 1 + } else { + 0.5625 + } + bytes <- prod(tensor$shape) * per_elem + if (shard_size + bytes > shard_bytes) { + flush_shard() + } + torch::with_no_grad({ + if (quant && format == "fp8") { + q <- .fp8_quantize(tensor) + shard[[name]] <- q$weight + shard[[paste0(name, "_scale")]] <- q$scale + } else if (quant) { + q <- ltx23_nf4_quantize(tensor) + shard[[name]] <- q$packed + shard[[paste0(name, "_absmax")]] <- q$absmax + } else { + shard[[name]] <- tensor$to(dtype = resident_dtype) + } + }) + n_cast <- n_cast + quant + shard_size <- shard_size + bytes + rm(tensor) + # torch's native memory is invisible to R's gc trigger: collect + # per tensor so the source tensors and casts do not pile up + gc(verbose = FALSE) + } + flush_shard() + + manifest <- list(model = "qwen3", format = format, shards = shard_files, + cast = n_cast, config = text_config) + jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE, + pretty = TRUE, digits = NA, null = "null") + if (verbose) { + message(sprintf("Quantized %d Qwen3 weights to %s across %d shards: %s", + n_cast, format, length(shard_files), output_dir)) + } + invisible(manifest) +} + +# transformers key -> qwen3_encoder name; NA for keys the text model does +# not use (tied LM head, Qwen3-VL vision tower) +.qwen3_map_key <- function(key, vl) { + if (identical(key, "lm_head.weight")) { + return(NA_character_) + } + if (vl) { + if (startsWith(key, "model.visual.")) { + return(NA_character_) + } + key <- sub("^model\\.language_model\\.", "model.", key) + } + key +} + +# The quantized set: every decoder layer's attention and MLP projections +.qwen3_is_quant_key <- function(name) { + grepl(paste0("^model\\.layers\\.[0-9]+\\.", + "(self_attn\\.[qkvo]_proj|mlp\\.(gate|up|down)_proj)\\.weight$"), + name) +} + +# TRUE when `dir` holds a complete Qwen3 artifact (of `format`, if given) +.is_qwen3_artifact <- function(dir, format = NULL) { + manifest_path <- file.path(dir, "manifest.json") + if (!file.exists(manifest_path)) { + return(FALSE) + } + manifest <- jsonlite::fromJSON(manifest_path) + identical(manifest$model, "qwen3") && + (is.null(format) || identical(manifest$format, format)) && + all(file.exists(file.path(dir, manifest$shards))) +} + +# Build a qwen3_encoder from a quantized artifact: skeleton at the +# compute dtype, the cast-set linears swapped for fp8 or NF4 modules +# (dequantized per forward), residents copied, anything unfilled an +# error. fp8 weights are plain module fields: they follow the module to +# `device` here, and staging must collect them (.flux_fp8_collect). +.load_qwen3_quantized <- function(dir, device, dtype, verbose, ...) { + manifest <- jsonlite::fromJSON(file.path(dir, "manifest.json"), + simplifyVector = TRUE) + format <- manifest$format + if (identical(format, "fp8") && !.st_can_write("float8_e4m3fn")) { + stop("This fp8 artifact needs float8 support the installed ", + "safetensors lacks (released in safetensors 0.3.0).", + call. = FALSE) + } + args <- utils::modifyList(.qwen3_args(manifest$config), list(...)) + if (verbose) { + message(sprintf("Creating Qwen3 %s model: %d layers", format, + args$num_hidden_layers %||% 36L)) + } + model <- do.call(.construct_skeleton, + c(list(qwen3_encoder), args, list(dtype = .flux_dtype(dtype)))) + ckpt <- ltx23_open_fp8_checkpoint(dir) + sib_suffix <- if (format == "fp8") "_scale" else "_absmax" + cast_keys <- sub(paste0(sib_suffix, "$"), "", + grep(paste0("\\.weight", sib_suffix, "$"), ckpt$keys, + value = TRUE)) + + torch::with_no_grad({ + for (key in cast_keys) { + segs <- strsplit(key, ".", fixed = TRUE)[[1]] + parent <- .ltx23_walk_module(model, utils::head(segs, -2L)) + leaf <- segs[length(segs) - 1L] + shape <- parent[[leaf]]$weight$shape + w <- ckpt$handle$get_tensor(key) + sib <- ckpt$handle$get_tensor(paste0(key, sib_suffix)) + # Assign the module itself: the setters return the inner R6 + # object, which `$<-` does not register as a child + if (format == "fp8") { + q <- ltx23_fp8_linear(shape[1], shape[2], bias = FALSE) + q$set_fp8_weight(w, sib) + } else { + q <- ltx23_nf4_linear(shape[1], shape[2], bias = FALSE) + q$set_nf4_weight(w, sib) + } + do.call(`$<-`, list(parent, leaf, q)) + } + + dests <- c(model$named_parameters(), model$named_buffers()) + dests <- dests[!grepl("weight_(nf4|absmax)$", names(dests))] + residents <- setdiff(ckpt$keys, + c(cast_keys, paste0(cast_keys, sib_suffix))) + unmapped <- setdiff(residents, names(dests)) + unfilled <- setdiff(names(dests), residents) + if (length(unmapped) || length(unfilled)) { + stop("Qwen3 ", format, " load: ", length(unmapped), + " unmapped keys, ", length(unfilled), " unfilled params") + } + for (key in residents) { + dests[[key]]$copy_(ckpt$handle$get_tensor(key)) + } + }) + gc(verbose = FALSE) + + model$to(device = device) + if (format == "fp8" && device != "cpu") { + .flux_fp8_to_device(model, device) + } + model$eval() + model +} diff --git a/R/qwen3_text_encoder.R b/R/qwen3_text_encoder.R index e73ba64..40afc26 100644 --- a/R/qwen3_text_encoder.R +++ b/R/qwen3_text_encoder.R @@ -229,9 +229,14 @@ qwen3_encoder <- torch::nn_module( #' input with 1D positions, Qwen3-VL's interleaved MRoPE reduces to the #' plain Qwen3 rotary (all three axes carry the same position). #' +#' A directory written by \code{\link{qwen3_quantize}} loads with fp8 or +#' NF4 projection linears that dequantize per forward. fp8 weights are +#' plain module fields, not parameters: they move with \code{device} +#' here, but \code{$to()} on the returned model moves only the residents. +#' #' @param model_path Directory with \code{config.json} and #' \code{model*.safetensors} (FLUX.2-klein's or FLUX 3's -#' \code{text_encoder}). +#' \code{text_encoder}), or a \code{\link{qwen3_quantize}} artifact. #' @param device Character. Target device. #' @param dtype Character. "bfloat16" (GPU) or "float32" (CPU). #' @param verbose Logical. @@ -243,6 +248,9 @@ qwen3_encoder <- torch::nn_module( load_qwen3_text_encoder <- function(model_path, device = "cpu", dtype = "float32", verbose = TRUE, ...) { model_path <- path.expand(model_path) + if (.is_qwen3_artifact(model_path)) { + return(.load_qwen3_quantized(model_path, device, dtype, verbose, ...)) + } config <- NULL config_path <- file.path(model_path, "config.json") if (file.exists(config_path)) { @@ -254,24 +262,7 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", config <- config$text_config } - args <- list(vocab_size = config$vocab_size, - hidden_size = config$hidden_size, - intermediate_size = config$intermediate_size, - num_hidden_layers = config$num_hidden_layers, - num_attention_heads = config$num_attention_heads, - num_key_value_heads = config$num_key_value_heads, - head_dim = config$head_dim) - args <- Filter(function(x) !is.null(x) && length(x) > 0L, args) - args <- lapply(args, as.integer) - # transformers >= 5 writes rope_theta under rope_parameters - rope_theta <- config$rope_theta %||% config$rope_parameters$rope_theta - if (!is.null(rope_theta)) { - args$rope_theta <- as.numeric(rope_theta) - } - if (!is.null(config$rms_norm_eps)) { - args$rms_norm_eps <- as.numeric(config$rms_norm_eps) - } - args <- utils::modifyList(args, list(...)) + args <- utils::modifyList(.qwen3_args(config), list(...)) # Build at the target dtype: a 4B fp32 build before the bf16 cast # needs 16 GB of host RAM model <- do.call(.construct_skeleton, @@ -283,19 +274,8 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", version = NULL, config = config, path = model_path), class = "ltx23_checkpoint" ) - map_key <- function(key) { - if (identical(key, "lm_head.weight")) { - return(NA_character_) # tied to embed_tokens - } - if (vl) { - if (startsWith(key, "model.visual.")) { - return(NA_character_) - } - key <- sub("^model\\.language_model\\.", "model.", key) - } - key - } - res <- ltx23_load_group(ckpt, ckpt$keys, model, map_key = map_key, + res <- ltx23_load_group(ckpt, ckpt$keys, model, + map_key = function(key) .qwen3_map_key(key, vl), verbose = verbose) if (length(res$unmapped) || length(res$unfilled)) { stop("Qwen3 encoder load: ", length(res$unmapped), @@ -307,6 +287,28 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", model } +# qwen3_encoder constructor arguments from a (text) config +.qwen3_args <- function(config) { + args <- list(vocab_size = config$vocab_size, + hidden_size = config$hidden_size, + intermediate_size = config$intermediate_size, + num_hidden_layers = config$num_hidden_layers, + num_attention_heads = config$num_attention_heads, + num_key_value_heads = config$num_key_value_heads, + head_dim = config$head_dim) + args <- Filter(function(x) !is.null(x) && length(x) > 0L, args) + args <- lapply(args, as.integer) + # transformers >= 5 writes rope_theta under rope_parameters + rope_theta <- config$rope_theta %||% config$rope_parameters$rope_theta + if (!is.null(rope_theta)) { + args$rope_theta <- as.numeric(rope_theta) + } + if (!is.null(config$rms_norm_eps)) { + args$rms_norm_eps <- as.numeric(config$rms_norm_eps) + } + args +} + #' Encode prompts with the Qwen3 encoder for FLUX.2 #' #' Tokenizes with the chat template, runs the encoder with the padding diff --git a/R/stream_flux3.R b/R/stream_flux3.R index 4c76f50..2380121 100644 --- a/R/stream_flux3.R +++ b/R/stream_flux3.R @@ -40,7 +40,8 @@ .flux3_text_stream <- function(encoder, pin) { layers <- encoder$model$layers encoder$layer_stream <- lapply(seq_along(layers), function(i) { - .stage_pairs(c(layers[[i]]$parameters, layers[[i]]$buffers), pin) + .stage_pairs(c(layers[[i]]$parameters, layers[[i]]$buffers, + .flux_fp8_collect(layers[[i]])), pin) }) .stage_pairs(c(encoder$model$embed_tokens$parameters, encoder$model$norm$parameters), pin) diff --git a/R/txt2img_flux3.R b/R/txt2img_flux3.R index ac9c58e..8415193 100644 --- a/R/txt2img_flux3.R +++ b/R/txt2img_flux3.R @@ -26,6 +26,24 @@ NULL file.path(tools::R_user_dir("diffuseR", "data"), "flux3-klein-") } +# Quantized Qwen3-VL text model (see qwen3_quantize) +.flux3_text_artifact <- function(format) { + paste0(.flux3_artifact_prefix(), "qwen3vl-", format) +} + +# text_precision = "auto": a built artifact (fp8 first), else bf16 +.flux3_resolve_text_precision <- function(text_precision) { + if (!identical(text_precision, "auto")) { + return(text_precision) + } + for (p in c("fp8", "nf4")) { + if (.is_qwen3_artifact(.flux3_text_artifact(p), p)) { + return(p) + } + } + "bf16" +} + # Resolve a FLUX 3 file from the HuggingFace cache .flux3_cached <- function(file, revision = .flux3_revision) { if (!requireNamespace("hfhub", quietly = TRUE)) { @@ -83,6 +101,10 @@ NULL #' in inside each forward. NULL: all of it from 12 GB up, else #' \code{vram_gb - 6.5} (1.5 GB on an 8 GB card, enough headroom for #' 1024x1024 text-to-image and a 1 MP edit). fp8 only. +#' @param text_precision Qwen3-VL weights: "auto" (default: the fp8, +#' then NF4, artifact built by \code{\link{download_flux3_klein}}, +#' else bf16), "fp8" (~4.4 GB of host memory instead of ~8 GB), +#' "nf4" (~2.8 GB), or "bf16" (the cached source). #' @param verbose Logical. #' #' @return A \code{flux3_pipeline} list. @@ -95,6 +117,7 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", encoder = TRUE, revision = .flux3_revision, dtype = NULL, vram_gb = NULL, stream_text = NULL, resident_gb = NULL, + text_precision = c("auto", "fp8", "nf4", "bf16"), verbose = TRUE) { .flux2_rev(revision) precision <- match.arg(precision) @@ -199,11 +222,23 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading Qwen3-VL text encoder...") } - te_dir <- dirname(.flux3_cached("text_encoder/config.json", revision)) - for (f in c("text_encoder/model.safetensors.index.json", - sprintf("text_encoder/model-%05d-of-00002.safetensors", 1:2))) { - .flux3_cached(f, revision) + text_precision <- .flux3_resolve_text_precision(match.arg(text_precision)) + te_dir <- if (text_precision == "bf16") { + for (f in c("text_encoder/model.safetensors.index.json", + sprintf("text_encoder/model-%05d-of-00002.safetensors", 1:2))) { + .flux3_cached(f, revision) + } + dirname(.flux3_cached("text_encoder/config.json", revision)) + } else { + dir <- .flux3_text_artifact(text_precision) + if (!.is_qwen3_artifact(dir, text_precision)) { + stop("No ", text_precision, " Qwen3-VL artifact at ", dir, + "; run download_flux3_klein(text_precision = \"", + text_precision, "\") first.", call. = FALSE) + } + dir } + pipe$text_precision <- text_precision pipe$text_encoder <- load_qwen3_text_encoder( te_dir, device = if (phase_offload) "cpu" else text_device, dtype = if (text_device == "cpu") "float32" else "bfloat16", diff --git a/inst/tinytest/test_qwen3_quantize.R b/inst/tinytest/test_qwen3_quantize.R new file mode 100644 index 0000000..6659cbe --- /dev/null +++ b/inst/tinytest/test_qwen3_quantize.R @@ -0,0 +1,98 @@ +# Qwen3 text-encoder quantization (fp8 / NF4 artifacts) on the tiny +# random-init Qwen3-VL checkpoint: key mapping, the cast set, loading, +# outputs near the full-precision model, and layer streaming with fp8 +# fields. + +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) + +ckpt_dir <- system.file("tinytest", "fixtures", "flux3_qwen_tiny_ckpt", + package = "diffuseR") +if (ckpt_dir == "") ckpt_dir <- "fixtures/flux3_qwen_tiny_ckpt" +if (!dir.exists(ckpt_dir)) { + exit_file("flux3 qwen fixtures missing") +} + +# --- Key mapping and the cast set ------------------------------------------------- + +map <- diffuseR:::.qwen3_map_key +expect_equal(map("model.language_model.layers.0.mlp.up_proj.weight", TRUE), + "model.layers.0.mlp.up_proj.weight") +expect_true(is.na(map("model.visual.blocks.0.attn.qkv.weight", TRUE))) +expect_true(is.na(map("lm_head.weight", FALSE))) +expect_equal(map("model.norm.weight", FALSE), "model.norm.weight") + +is_q <- diffuseR:::.qwen3_is_quant_key +expect_true(all(is_q(c("model.layers.3.self_attn.q_proj.weight", + "model.layers.3.self_attn.o_proj.weight", + "model.layers.12.mlp.down_proj.weight")))) +expect_false(any(is_q(c("model.layers.3.self_attn.q_norm.weight", + "model.layers.3.input_layernorm.weight", "model.embed_tokens.weight", + "model.norm.weight")))) + +# --- Reference ---------------------------------------------------------------------- + +ref_model <- load_qwen3_text_encoder(ckpt_dir, device = "cpu", + dtype = "float32", verbose = FALSE) +n_layers <- length(ref_model$model$layers) +ids <- torch::torch_tensor(matrix(c(5L, 9L, 13L, 2L, 7L, 11L), nrow = 1), + dtype = torch::torch_long()) +ref <- torch::with_no_grad(ref_model(ids, out_layers = c(1L, 2L, n_layers))) +rel_err <- function(a, b) { + (a - b)$norm()$item() / b$norm()$item() +} + +formats <- if (diffuseR:::.st_can_write("float8_e4m3fn")) c("fp8", "nf4") else "nf4" +# Per-tensor fp8 keeps ~3% error on the real 4B; NF4 ~5% +tol <- c(fp8 = 0.1, nf4 = 0.25) +for (f in formats) { + dir <- file.path(tempdir(), paste0("qwen3-tiny-", f)) + unlink(dir, recursive = TRUE) + manifest <- qwen3_quantize(ckpt_dir, dir, format = f, verbose = FALSE) + # 7 projections per layer + expect_equal(manifest$cast, 7L * n_layers) + expect_true(diffuseR:::.is_qwen3_artifact(dir, f)) + expect_false(diffuseR:::.is_qwen3_artifact(dir, setdiff(c("fp8", "nf4"), f))) + keys <- diffuseR:::ltx23_open_fp8_checkpoint(dir)$keys + expect_false(any(grepl("visual|language_model|lm_head", keys))) + # A valid artifact is not rebuilt + again <- qwen3_quantize(ckpt_dir, dir, format = f, verbose = FALSE) + expect_equal(again$shards, manifest$shards) + + model <- load_qwen3_text_encoder(dir, device = "cpu", dtype = "float32", + verbose = FALSE) + q_proj <- model$model$layers[[1]]$self_attn$q_proj + expect_true(inherits(q_proj, if (f == "fp8") "ltx23_fp8_linear" else "ltx23_nf4_linear")) + expect_equal(model$rope_theta, ref_model$rope_theta) + out <- torch::with_no_grad(model(ids, out_layers = c(1L, 2L, n_layers))) + for (k in seq_along(ref)) { + err <- rel_err(out[[k]], ref[[k]]) + expect_true(err > 0 && err < tol[[f]], info = sprintf("%s state %d: %.4f", f, k, err)) + } + + # Layer streaming stages the fp8 fields with each layer + base <- diffuseR:::.flux3_text_stream(model, pin = FALSE) + per_layer <- length(model$model$layers[[1]]$parameters) + + length(model$model$layers[[1]]$buffers) + expected <- if (f == "fp8") per_layer + 14L else per_layer + expect_equal(length(model$layer_stream[[1]]), expected) + diffuseR:::.staged_onload(base, "cpu") + streamed <- torch::with_no_grad(model(ids, out_layers = c(1L, 2L, n_layers))) + for (k in seq_along(out)) { + expect_true((streamed[[k]] - out[[k]])$abs()$max()$item() == 0) + } + model$layer_stream <- NULL +} + +# --- Pipeline resolution ------------------------------------------------------------- + +resolve <- diffuseR:::.flux3_resolve_text_precision +expect_equal(resolve("bf16"), "bf16") +expect_equal(resolve("nf4"), "nf4") +expect_true(resolve("auto") %in% c("fp8", "nf4", "bf16")) diff --git a/man/download_flux3_klein.Rd b/man/download_flux3_klein.Rd index 2503a64..c2049a3 100644 --- a/man/download_flux3_klein.Rd +++ b/man/download_flux3_klein.Rd @@ -8,6 +8,7 @@ download_flux3_klein( precision = c("auto", "fp8", "nf4"), output_dir = NULL, text_encoders = TRUE, + text_precision = c("fp8", "nf4", "bf16"), verbose = TRUE ) } @@ -22,11 +23,19 @@ float8, else nf4), "fp8" (~6 GB), or "nf4".} \item{text_encoders}{Logical. Also fetch the Qwen3-VL text encoder, tokenizer and VAE (~9.2 GB).} +\item{text_precision}{Quantize the Qwen3-VL text model too (with +\code{quantize}): "fp8" (default; a 4.4 GB artifact, host memory at +1024x1024 20.2 -> 15.5 GB, images close to bf16 text: 27 dB PSNR), +"nf4" (2.8 GB, 13.5 GB host, but a visibly different image for the +same seed), or "bf16" (keep only the cached source). +\code{\link{flux3_load_pipeline}} picks the artifact up.} + \item{verbose}{Logical.} } \value{ Invisibly, a list with \code{transformer_dir}, - \code{artifact_dir}, and \code{support} (named file paths). + \code{artifact_dir}, \code{support} (named file paths), and + \code{text_artifact_dir} when the text model was quantized. } \description{ Skips work already done: a valid quantized manifest short-circuits diff --git a/man/flux3_load_pipeline.Rd b/man/flux3_load_pipeline.Rd index 56a44a0..4fe8aa7 100644 --- a/man/flux3_load_pipeline.Rd +++ b/man/flux3_load_pipeline.Rd @@ -17,6 +17,7 @@ flux3_load_pipeline( vram_gb = NULL, stream_text = NULL, resident_gb = NULL, + text_precision = c("auto", "fp8", "nf4", "bf16"), verbose = TRUE ) } @@ -65,6 +66,11 @@ in inside each forward. NULL: all of it from 12 GB up, else \code{vram_gb - 6.5} (1.5 GB on an 8 GB card, enough headroom for 1024x1024 text-to-image and a 1 MP edit). fp8 only.} +\item{text_precision}{Qwen3-VL weights: "auto" (default: the fp8, +then NF4, artifact built by \code{\link{download_flux3_klein}}, +else bf16), "fp8" (~4.4 GB of host memory instead of ~8 GB), +"nf4" (~2.8 GB), or "bf16" (the cached source).} + \item{verbose}{Logical.} } \value{ diff --git a/man/load_qwen3_text_encoder.Rd b/man/load_qwen3_text_encoder.Rd index a688a49..2f38ea2 100644 --- a/man/load_qwen3_text_encoder.Rd +++ b/man/load_qwen3_text_encoder.Rd @@ -14,7 +14,7 @@ load_qwen3_text_encoder( \arguments{ \item{model_path}{Directory with \code{config.json} and \code{model*.safetensors} (FLUX.2-klein's or FLUX 3's -\code{text_encoder}).} +\code{text_encoder}), or a \code{\link{qwen3_quantize}} artifact.} \item{device}{Character. Target device.} @@ -36,3 +36,10 @@ loads its text model only; the vision tower is skipped. For text-only input with 1D positions, Qwen3-VL's interleaved MRoPE reduces to the plain Qwen3 rotary (all three axes carry the same position). } +\details{ +A directory written by \code{\link{qwen3_quantize}} loads with fp8 or +NF4 projection linears that dequantize per forward. fp8 weights are +plain module fields, not parameters: they move with \code{device} +here, but \code{$to()} on the returned model moves only the residents. + +} diff --git a/man/qwen3_quantize.Rd b/man/qwen3_quantize.Rd new file mode 100644 index 0000000..2513b65 --- /dev/null +++ b/man/qwen3_quantize.Rd @@ -0,0 +1,50 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{qwen3_quantize} +\alias{qwen3_quantize} +\title{Quantize a Qwen3 text encoder to fp8 or NF4 shards} +\usage{ +qwen3_quantize( + model_path, + output_dir, + format = c("fp8", "nf4"), + shard_bytes = 1.9e+09, + force = FALSE, + verbose = TRUE +) +} +\arguments{ +\item{model_path}{Source directory with \code{config.json} and +\code{model*.safetensors}.} + +\item{output_dir}{Output directory for shards + manifest.} + +\item{format}{"fp8" or "nf4".} + +\item{shard_bytes}{Numeric. Target shard size in bytes (1.9e9 keeps +shards readable by safetensors older than 0.3.0).} + +\item{force}{Logical. Re-quantize even if a valid manifest exists.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the manifest list. +} +\description{ +Streams a transformers Qwen3 (FLUX.2 klein, Z-Image) or Qwen3-VL +(FLUX 3) text encoder tensor by tensor. The per-layer projection +weights (q/k/v/o and gate/up/down, 252 tensors in the 36-layer 4B) +are stored as float8_e4m3fn with a per-tensor absmax/448 +\code{_scale}, or as NF4 (packed uint8 + \code{_absmax}); +embeddings and norms are copied at bfloat16. The vision tower and the +tied LM head are dropped, and keys are written in the +\code{\link{qwen3_encoder}} layout (\code{model.layers.*}). The +manifest embeds the text config, so the loader needs no other file. +} +\details{ +Host memory: the 4B text model is ~8 GB in bfloat16, ~4.4 GB as fp8 +and ~2.8 GB as NF4 (the 0.8 GB embedding stays bfloat16 in both). +Load the artifact with \code{\link{load_qwen3_text_encoder}} on its +directory. + +} From 2d1ec4be2c42087b7d63f8bc7d1ce0acaac31115 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 29 Oct 2026 16:20:16 -0500 Subject: [PATCH 2/2] Bump version to 0.2.2.22 --- DESCRIPTION | 2 +- NEWS.md | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index 023e031..75bab50 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.21 +Version: 0.2.2.22 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), diff --git a/NEWS.md b/NEWS.md index 7016edb..b89cb82 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,20 @@ +# diffuseR 0.2.2.22 + +* **FLUX 3 in less host RAM: fp8 Qwen3-VL.** `qwen3_quantize()` writes + a Qwen3 or Qwen3-VL text model as an fp8 or NF4 artifact (the 252 + projection linears quantized; embeddings and norms bf16; vision tower + dropped), and `load_qwen3_text_encoder()` loads one. + `download_flux3_klein(text_precision = "fp8")` builds it by default + and `flux3_load_pipeline(text_precision = "auto")` uses it. + - Measured at 1024x1024 on the 16 GB card: host high-water 20.2 GB + with the bf16 encoder, 15.5 GB with fp8, 13.5 GB with NF4, at the + same ~24 s per image. + - fp8 text keeps the image: 27.2 dB PSNR against bf16 text, prompt + embeddings within 2.7%. NF4 text (5%) gives a different composition + for the same seed, so it is opt-in. + - The 8 GB layer-streaming plan stages the fp8 weights too: GPU peak + 4.9 GB, the same image as the resident encoder. + # diffuseR 0.2.2.21 * **LTX-2.5 in 10 GB of host RAM.** `stage_from = "disk"` on