diff --git a/CLAUDE.md b/CLAUDE.md index c672afa..080755f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -471,6 +471,24 @@ own text encoder: R bf16 52.2 dB PSNR, R fp8 36.7 dB. Edits that change little (no snow from a "covered in snow" caption) reproduce in diffusers; they come from the model without its prompt upsampler. +Small cards (planned from `vram_gb`): below 10 GB the Qwen3-VL layers +stream from pinned host memory one at a time (`stream_text`); below +12 GB only `vram_gb - 6.5` GB of fp8 linears stay resident and the +rest stream inside each forward (`resident_gb`). Simulated with a +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). + +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 +that from a 12 GB footprint, which on an 8 GB card means no gc even +when an allocation fails. Above `allocated_rate` it is a full gc, +otherwise a minor one; block outputs promoted by the per-block minor +gcs need the full one. Streaming mode uses reserved 0.6 / allocated +0.8. The streamed onload's R-side copies hold device memory until +collected, hence the minor gc after each streamed Qwen layer. + Gotchas that bit once: HF's last `hidden_states` entry carries the final norm (the Qwen3 encoder now matches for k = num_layers); building 5.6B params in fp32 before the bf16 cast needs 22 GB of host RAM, so the diff --git a/DESCRIPTION b/DESCRIPTION index f85ffa1..04adcfd 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.16 +Version: 0.2.2.17 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 1c29c74..899aa18 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,24 @@ +# diffuseR 0.2.2.17 + +* **FLUX 3 on 8 and 12 GB cards.** `flux3_load_pipeline()` plans from + the card size (`vram_gb`, detected by default): + - `stream_text` (default below 10 GB): the Qwen3-VL layers stay in + pinned host memory and visit the GPU one at a time during the + encode (~1 s), instead of needing ~9 GB at once. + - `resident_gb` (default `vram_gb - 6.5` below 12 GB): that many GB + of fp8 transformer weights stay on the GPU; the rest streams in + from pinned memory inside each forward. + - Measured on a 16 GB card with the rest held back: 8 GB runs + 1024x1024 text-to-image in ~30 s and a 1 MP edit in ~56 s; 12 GB + streams nothing and runs them in ~27 s and ~53 s. +* FLUX 3 runs R's garbage collector after every transformer block as a + minor collection (it ran a full one after the joint blocks only): + 1024x1024 text-to-image on 16 GB went from 44 s to 27 s and its peak + from 11.7 to 9.1 GB. +* The VAE encoder and decoder blocks run a minor collection after each + resnet when `diffuseR.block_gc` is on, which bounds the dead + temporaries at 1024x1024 on small cards. + # diffuseR 0.2.2.16 * **FLUX 3 Image Klein.** `txt2img_flux3()` and `img2img_flux3()` run diff --git a/R/dit_flux3.R b/R/dit_flux3.R index 0d95cb4..ef58019 100644 --- a/R/dit_flux3.R +++ b/R/dit_flux3.R @@ -177,6 +177,7 @@ flux3_transformer <- torch::nn_module( } # 5. Per-mode blocks (each mode attends within itself only) + block_gc <- isTRUE(getOption("diffuseR.block_gc")) for (i in seq_along(self$txt_mode_blocks)) { encoder_hidden_states <- self$txt_mode_blocks[[i]]( encoder_hidden_states, txt_early, txt_rope, chunk_size) @@ -184,11 +185,17 @@ flux3_transformer <- torch::nn_module( for (i in seq_along(self$mode_blocks$image)) { hidden_states <- self$mode_blocks$image[[i]]( hidden_states, image_early, image_rope, chunk_size) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } if (has_cond) { for (i in seq_along(self$mode_blocks$image_cond)) { image_cond_hidden_states <- self$mode_blocks$image_cond[[i]]( image_cond_hidden_states, cond_early, cond_rope, chunk_size) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } } @@ -211,7 +218,6 @@ flux3_transformer <- torch::nn_module( torch::torch_cat(lapply(ropes, `[[`, 2L), dim = 1L)) rm(hidden_states, image_cond_hidden_states) - block_gc <- isTRUE(getOption("diffuseR.block_gc")) for (i in seq_along(self$single_transformer_blocks)) { res <- self$single_transformer_blocks[[i]]( hidden_states = content, @@ -223,7 +229,7 @@ flux3_transformer <- torch::nn_module( encoder_hidden_states <- res[[1]] content <- res[[2]] if (block_gc) { - gc(verbose = FALSE) + gc(verbose = FALSE, full = FALSE) } } diff --git a/R/memory_flux.R b/R/memory_flux.R index c8e3c74..1f5bebd 100644 --- a/R/memory_flux.R +++ b/R/memory_flux.R @@ -23,12 +23,18 @@ NULL # 12.9 GB. ltx23_tune_gc's reserved_rate stays as the fragmentation # net; the LTX video pipelines keep their own separately measured # gates. -.flux_gc_gates <- function(footprint_gb = 12) { +.flux_gc_gates <- function(footprint_gb = 12, allocated_rate = 0.65, + reserved_rate = NULL) { if (is.null(getOption("torch.cuda_allocator_allocated_reserved_rate"))) { options(torch.cuda_allocator_allocated_reserved_rate = 1.0) } if (is.null(getOption("torch.cuda_allocator_allocated_rate"))) { - options(torch.cuda_allocator_allocated_rate = 0.65) + options(torch.cuda_allocator_allocated_rate = allocated_rate) + } + # NULL: ltx23_tune_gc sizes it from footprint_gb + if (!is.null(reserved_rate) && + is.null(getOption("torch.cuda_allocator_reserved_rate"))) { + options(torch.cuda_allocator_reserved_rate = reserved_rate) } ltx23_tune_gc(footprint_gb = footprint_gb) # start_torch() reads the three gate options ONCE at torch init and diff --git a/R/quantize_flux.R b/R/quantize_flux.R index 0491792..29bb261 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -183,8 +183,7 @@ NULL if (is.null(config)) { return(list()) } - args <- list( - in_channels = config$in_channels, + args <- list(in_channels = config$in_channels, cond_in_channels = config$cond_in_channels, num_layers = config$num_layers, num_single_layers = config$num_single_layers, @@ -192,8 +191,7 @@ NULL num_attention_heads = config$num_attention_heads, joint_attention_dim = config$joint_attention_dim, timestep_channels = config$timestep_channels, - axes_dims_rope = config$axes_dims_rope - ) + axes_dims_rope = config$axes_dims_rope) args <- Filter(function(x) !is.null(x) && length(x) > 0L, args) args <- lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x) for (field in c("mlp_ratio", "rope_theta", "eps")) { @@ -289,7 +287,13 @@ NULL } extra <- if (identical(nm, "transformer") && isTRUE(pipe$fp8_resident)) { - .flux_fp8_collect(module) + # A split transformer (FLUX 3 on small cards) stages only its + # resident linears; the streamed ones stay on the host + if (is.null(pipe$fp8_split)) { + .flux_fp8_collect(module) + } else { + .flux3_fp8_fields(pipe$fp8_split$resident) + } } else { NULL } diff --git a/R/qwen3_text_encoder.R b/R/qwen3_text_encoder.R index 19a59aa..e73ba64 100644 --- a/R/qwen3_text_encoder.R +++ b/R/qwen3_text_encoder.R @@ -29,8 +29,7 @@ NULL torch::torch_arange(start = 0, end = seq_len - 1, dtype = f32, device = device) } else { - torch::torch_tensor(as.numeric(positions), dtype = f32, - device = device) + torch::torch_tensor(as.numeric(positions), dtype = f32, device = device) } freqs <- pos$unsqueeze(2L) * inv_freq$unsqueeze(1L) # [S, D/2] # [1, 1, S, D/2] for broadcasting against [B, H, S, D/2] @@ -92,8 +91,7 @@ NULL initialize = function(hidden_size, intermediate_size) { self$gate_proj <- linear_noinit(hidden_size, intermediate_size, bias = FALSE) - self$up_proj <- linear_noinit(hidden_size, intermediate_size, - bias = FALSE) + self$up_proj <- linear_noinit(hidden_size, intermediate_size, bias = FALSE) self$down_proj <- linear_noinit(intermediate_size, hidden_size, bias = FALSE) }, @@ -197,8 +195,20 @@ qwen3_encoder <- torch::nn_module( out_layers <- sort(as.integer(out_layers)) states <- vector("list", length(out_layers)) n_layers <- length(self$model$layers) + # Layer streaming (see .flux3_text_stream): each layer's weights visit + # the device only for its own forward + stream <- self$layer_stream for (i in seq_len(max(out_layers))) { + if (!is.null(stream)) { + .staged_onload(stream[[i]], device) + } x <- self$model$layers[[i]](x, rope, mask) + if (!is.null(stream)) { + .staged_offload(stream[[i]]) + # The onload's R-side copies hold the device storage until + # collected; a minor gc frees this layer before the next + gc(verbose = FALSE, full = FALSE) + } hit <- which(out_layers == i) if (length(hit)) { # HF's last hidden_states entry carries the final norm @@ -265,8 +275,7 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", # Build at the target dtype: a 4B fp32 build before the bf16 cast # needs 16 GB of host RAM model <- do.call(.construct_skeleton, - c(list(qwen3_encoder), args, - list(dtype = .flux_dtype(dtype)))) + c(list(qwen3_encoder), args, list(dtype = .flux_dtype(dtype)))) opened <- .flux_open_sharded_dir(model_path, "model") ckpt <- structure( diff --git a/R/stream_flux3.R b/R/stream_flux3.R new file mode 100644 index 0000000..4c76f50 --- /dev/null +++ b/R/stream_flux3.R @@ -0,0 +1,91 @@ +# FLUX 3 on small cards: weights that do not fit stay in (pinned) host +# memory and visit the GPU only while they compute. The Qwen3-VL text +# model streams layer by layer through its encode; the fp8 transformer +# keeps a budget of linears resident and streams the rest inside each +# forward (ltx23_fp8_linear copies a host weight to the input's device). + +# How much to stream for a card of `vram_gb` (total). The Qwen3-VL +# encode peaks ~9.2 GB whole, so it streams below 10 GB. The fp8 +# transformer (~4.8 GB of fp8 linears) stays whole from 12 GB up; below +# that it keeps vram_gb - 6.5 GB resident, leaving room for the bf16 +# residents and a 1 MP edit's activations (8,192 content tokens). +# Measured on a simulated 8 GB card: 1.5 GB resident runs 1024x1024 +# text-to-image and the 1 MP edit; 3.5 GB runs text-to-image only. +.flux3_stream_plan <- function(vram_gb) { + if (!isTRUE(vram_gb > 0)) { + return(list(stream_text = FALSE, resident_gb = Inf)) + } + list(stream_text = vram_gb < 10, + resident_gb = if (vram_gb >= 12) Inf else max(0, vram_gb - 6.5)) +} + +# Staging pairs for a list of tensors (see staging.R): pinned host copies +# when `pin`, otherwise the tensors' own host storage +.stage_pairs <- function(tensors, pin) { + pin <- isTRUE(pin) && torch::cuda_is_available() + lapply(tensors, function(p) { + if (pin) { + pinned <- .pin_host(p) + p$set_data(pinned) + } else { + pinned <- p$detach() + } + list(live = p, pinned = pinned) + }) +} + +# Stream a Qwen3 encoder's layers: one staging set per layer goes on +# the module (its forward swaps each layer in and out); the returned +# pairs (embeddings and final norm) onload for the encode phase +.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(encoder$model$embed_tokens$parameters, + encoder$model$norm$parameters), pin) +} + +# Split a transformer's fp8 linears into a resident set (in module +# order, up to `resident_gb` of fp8 bytes) and a streamed set. Streamed +# weights are pinned in place so their per-forward copies run at DMA +# rate. +.flux3_fp8_split <- function(transformer, resident_gb, pin) { + mods <- list() + walk <- function(m) { + for (child in m$children) { + if (!is.null(child$weight_fp8)) { + mods[[length(mods) + 1L]] <<- child + } else { + walk(child) + } + } + } + walk(transformer) + bytes <- vapply(mods, function(m) m$weight_fp8$numel(), 0) + keep <- cumsum(bytes) <= resident_gb * 1024 ^ 3 + streamed <- mods[!keep] + if (isTRUE(pin) && torch::cuda_is_available()) { + for (m in streamed) { + m$weight_fp8$set_data(.pin_host(m$weight_fp8)) + } + } + list(resident = mods[keep], streamed = streamed, + resident_gb = sum(bytes[keep]) / 1024 ^ 3, + streamed_gb = sum(bytes[!keep]) / 1024 ^ 3) +} + +# The resident fp8 fields (weights and scales), for staging +.flux3_fp8_fields <- function(mods) { + unlist(lapply(mods, function(m) list(m$weight_fp8, m$weight_scale)), + recursive = FALSE) +} + +# Move the resident fp8 linears' weights (unstaged pipelines) +.flux3_fp8_to_device <- function(mods, device) { + for (m in mods) { + m$weight_fp8 <- m$weight_fp8$to(device = device) + m$weight_scale <- m$weight_scale$to(device = device) + } + invisible(NULL) +} diff --git a/R/txt2img_flux3.R b/R/txt2img_flux3.R index e631ae0..ac9c58e 100644 --- a/R/txt2img_flux3.R +++ b/R/txt2img_flux3.R @@ -73,6 +73,16 @@ NULL #' @param dtype Compute dtype, "bfloat16" or "float32". NULL (default): #' bfloat16 on CUDA, float32 on the CPU. "bfloat16" on the CPU halves #' host memory (the raw transformer is ~11 GB instead of ~22 GB). +#' @param vram_gb Numeric or NULL. Card size the streaming defaults are +#' planned for (default: the detected total). +#' @param stream_text Logical or NULL. Stream the Qwen3-VL layers through +#' the GPU one at a time from host memory, for cards that cannot hold +#' the whole encoder (default: below 10 GB). +#' @param resident_gb Numeric or NULL. GB of fp8 transformer weights +#' kept on the GPU; the rest stays in (pinned) host memory and streams +#' 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 verbose Logical. #' #' @return A \code{flux3_pipeline} list. @@ -83,7 +93,9 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, pin = NULL, encoder = TRUE, revision = .flux3_revision, - dtype = NULL, verbose = TRUE) { + dtype = NULL, vram_gb = NULL, + stream_text = NULL, resident_gb = NULL, + verbose = TRUE) { .flux2_rev(revision) precision <- match.arg(precision) if (precision != "bf16") { @@ -91,6 +103,11 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", .flux3_artifact_prefix()) } text_device <- text_device %||% device + gpu <- device != "cpu" + plan <- .flux3_stream_plan(if (gpu) vram_gb %||% .detect_vram() else 0) + stream_text <- isTRUE(stream_text %||% plan$stream_text) && gpu && + phase_offload && !identical(text_device, "cpu") + resident_gb <- resident_gb %||% plan$resident_gb if (is.null(model_dir)) { model_dir <- if (precision == "bf16") { dirname(.flux3_cached("transformer/config.json", revision)) @@ -111,8 +128,17 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", } if (device == "cuda") { # Gate footprint sized to the largest phase (the ~8 GB Qwen3-VL - # encode), as for klein - .flux_gc_gates(footprint_gb = 12) + # encode), as for klein. On a streaming (small) card torch must + # collect well before the limit: R gc runs only on an allocator + # miss with reserved memory above reserved_rate (12 GB / 8 GB + # would mean never), as a full gc above allocated_rate. Block + # outputs promoted by the per-block minor gcs need that full gc. + if (is.finite(resident_gb)) { + .flux_gc_gates(footprint_gb = 12, allocated_rate = 0.8, + reserved_rate = 0.6) + } else { + .flux_gc_gates(footprint_gb = 12) + } } component_device <- if (phase_offload) "cpu" else device compute_dtype <- dtype %||% if (device == "cpu") "float32" else "bfloat16" @@ -136,6 +162,14 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", pin = FALSE, fp8_resident = FALSE, verbose = verbose ) pipe$fp8_resident <- identical(pipe$format, "fp8") && device == "cuda" + if (pipe$fp8_resident && is.finite(resident_gb)) { + pipe$fp8_split <- .flux3_fp8_split(pipe$transformer, resident_gb, pin) + if (verbose) { + message(sprintf("fp8 transformer: %.1f GB resident, %.1f GB streamed", + pipe$fp8_split$resident_gb, + pipe$fp8_split$streamed_gb)) + } + } if (verbose) { message("Loading FLUX 3 VAE...") @@ -179,11 +213,18 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", .flux3_cached("text_encoder/tokenizer.json", revision)) components <- intersect(c("transformer", "decoder", "encoder"), names(pipe)) - if (!identical(text_device, "cpu")) { + if (!identical(text_device, "cpu") && !stream_text) { components <- c(components, "text_encoder") } pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, components, verbose = verbose) + if (stream_text) { + if (verbose) { + message("Qwen3-VL streams layer by layer") + } + pipe$staging <- c(pipe$staging, + list(text_encoder = .flux3_text_stream(pipe$text_encoder, pin))) + } structure(pipe, class = "flux3_pipeline") } @@ -398,7 +439,11 @@ txt2img_flux3 <- function(prompt, pipeline = NULL, width = 1024L, fp8_manual <- isTRUE(pipeline$fp8_resident) && is.null(staging[["transformer"]]) if (fp8_manual) { - .flux_fp8_to_device(transformer, device) + if (is.null(pipeline$fp8_split)) { + .flux_fp8_to_device(transformer, device) + } else { + .flux3_fp8_to_device(pipeline$fp8_split$resident, device) + } } if (verbose) { message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width, diff --git a/R/vae_decoder.R b/R/vae_decoder.R index 2cea400..b12cc04 100644 --- a/R/vae_decoder.R +++ b/R/vae_decoder.R @@ -140,8 +140,14 @@ VAEUpBlock <- torch::nn_module( }, forward = function(x) { + # At 1024x1024 each resnet leaves ~1-3 GB of dead temporaries until + # R collects them; a minor gc per resnet bounds that on small cards + block_gc <- isTRUE(getOption("diffuseR.block_gc")) for (i in seq_along(self$resnets)) { x <- self$resnets[[i]](x) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } if (!is.null(self$upsamplers)) { x <- self$upsamplers[[1]](x) diff --git a/R/vae_flux2_encoder.R b/R/vae_flux2_encoder.R index a4eb2e4..5345b1c 100644 --- a/R/vae_flux2_encoder.R +++ b/R/vae_flux2_encoder.R @@ -37,8 +37,13 @@ VAEDownBlock <- torch::nn_module( }, forward = function(x) { + # Minor gc per resnet, as in VAEUpBlock + block_gc <- isTRUE(getOption("diffuseR.block_gc")) for (i in seq_along(self$resnets)) { x <- self$resnets[[i]](x) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } if (!is.null(self$downsamplers)) { x <- self$downsamplers[[1]](x) diff --git a/inst/tinytest/test_flux3_stream.R b/inst/tinytest/test_flux3_stream.R new file mode 100644 index 0000000..81f5e85 --- /dev/null +++ b/inst/tinytest/test_flux3_stream.R @@ -0,0 +1,102 @@ +# FLUX 3 small-card streaming: the VRAM plan, Qwen3-VL layer streaming +# and the fp8 resident/streamed split. On the CPU the staged moves are +# host-to-host, so outputs must match the unstreamed model exactly. + +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) + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# --- Plan --------------------------------------------------------------------- + +plan <- diffuseR:::.flux3_stream_plan +expect_equal(plan(16), list(stream_text = FALSE, resident_gb = Inf)) +expect_equal(plan(12), list(stream_text = FALSE, resident_gb = Inf)) +expect_equal(plan(8), list(stream_text = TRUE, resident_gb = 1.5)) +expect_equal(plan(10)$stream_text, FALSE) +expect_equal(plan(10)$resident_gb, 3.5) +expect_equal(plan(6)$resident_gb, 0) +# No GPU detected: nothing streams +expect_equal(plan(0), list(stream_text = FALSE, resident_gb = Inf)) + +# --- Qwen3-VL layer streaming ----------------------------------------------------- + +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)) { + enc <- load_qwen3_text_encoder(ckpt_dir, device = "cpu", dtype = "float32", + verbose = FALSE) + ids <- torch::torch_tensor(matrix(c(5L, 9L, 13L, 2L, 7L), nrow = 1), + dtype = torch::torch_long()) + ref <- torch::with_no_grad(enc(ids, out_layers = c(1L, 2L, 4L))) + + base <- diffuseR:::.flux3_text_stream(enc, pin = FALSE) + # One staging set per layer; embeddings + final norm in the base set + expect_equal(length(enc$layer_stream), length(enc$model$layers)) + expect_equal(length(base), 2L) + n_layer_tensors <- sum(vapply(enc$layer_stream, length, 0L)) + expect_equal(n_layer_tensors, + length(enc$parameters) - 2L) + diffuseR:::.staged_onload(base, "cpu") + out <- torch::with_no_grad(enc(ids, out_layers = c(1L, 2L, 4L))) + for (k in seq_along(ref)) { + expect_true(max_abs_diff(out[[k]], ref[[k]]) == 0) + } + enc$layer_stream <- NULL +} + +# --- fp8 split ----------------------------------------------------------------------- + +tiny_dir <- system.file("tinytest", "fixtures", "flux3_tiny_ckpt", + package = "diffuseR") +if (tiny_dir == "") tiny_dir <- "fixtures/flux3_tiny_ckpt" +fixture_path <- system.file("tinytest", "fixtures", "flux3_model.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/flux3_model.safetensors" +f8_ok <- diffuseR:::.st_can_write("float8_e4m3fn") +if (f8_ok && dir.exists(tiny_dir) && file.exists(fixture_path)) { + fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + fp8_dir <- file.path(tempdir(), "flux3-tiny-fp8-stream") + unlink(fp8_dir, recursive = TRUE) + manifest <- flux_quantize(tiny_dir, output_dir = fp8_dir, format = "fp8", + verbose = FALSE) + model <- flux_load_transformer(flux_open_quantized(fp8_dir), device = "cpu", + dtype = "float32", pin = FALSE, verbose = FALSE) + run <- function() { + torch::with_no_grad(model(hidden_states = fx$hidden, + encoder_hidden_states = fx$encoder, timestep = fx$timestep, + img_ids = fx$img_ids, txt_ids = fx$txt_ids, + image_cond_hidden_states = fx$cond, image_cond_ids = fx$cond_ids, + target_timestep = fx$target_timestep)) + } + ref <- run() + + all <- diffuseR:::.flux3_fp8_split(model, Inf, pin = FALSE) + expect_equal(length(all$resident), manifest$cast) + expect_equal(length(all$streamed), 0L) + none <- diffuseR:::.flux3_fp8_split(model, 0, pin = FALSE) + expect_equal(length(none$streamed), manifest$cast) + + half <- all$resident_gb / 2 + split <- diffuseR:::.flux3_fp8_split(model, half, pin = FALSE) + expect_true(split$resident_gb <= half) + expect_true(length(split$resident) > 0L && length(split$streamed) > 0L) + expect_equal(length(split$resident) + length(split$streamed), manifest$cast) + expect_equal(length(diffuseR:::.flux3_fp8_fields(split$resident)), + 2L * length(split$resident)) + expect_true(max_abs_diff(run(), ref) == 0) + ltx23_release_dequant_buffers() + unlink(fp8_dir, recursive = TRUE) +} +options(diffuseR.block_gc = NULL) diff --git a/man/flux3_load_pipeline.Rd b/man/flux3_load_pipeline.Rd index 57aa0ce..56a44a0 100644 --- a/man/flux3_load_pipeline.Rd +++ b/man/flux3_load_pipeline.Rd @@ -14,6 +14,9 @@ flux3_load_pipeline( encoder = TRUE, revision = .flux3_revision, dtype = NULL, + vram_gb = NULL, + stream_text = NULL, + resident_gb = NULL, verbose = TRUE ) } @@ -49,6 +52,19 @@ the validated snapshot).} bfloat16 on CUDA, float32 on the CPU. "bfloat16" on the CPU halves host memory (the raw transformer is ~11 GB instead of ~22 GB).} +\item{vram_gb}{Numeric or NULL. Card size the streaming defaults are +planned for (default: the detected total).} + +\item{stream_text}{Logical or NULL. Stream the Qwen3-VL layers through +the GPU one at a time from host memory, for cards that cannot hold +the whole encoder (default: below 10 GB).} + +\item{resident_gb}{Numeric or NULL. GB of fp8 transformer weights +kept on the GPU; the rest stays in (pinned) host memory and streams +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{verbose}{Logical.} } \value{