FLUX 3 on 8 and 12 GB cards: stream Qwen3-VL layers and fp8 linears

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

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

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

R/dit_flux3.R

@@ -177,18 +177,25 @@ 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)
     }
     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)
         }
     }
 

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

R/quantize_flux.R

@@ -183,17 +183,15 @@ 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,
                  attention_head_dim = config$attention_head_dim,
                  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
         }

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(

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