FLUX 3 in less host RAM: fp8 Qwen3-VL text encoder

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

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")),

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)

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

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)

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)
                 }
             })

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{<key>_scale}, or as NF4 (packed uint8 + \code{<key>_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
+}