LTX-2.5 in 10 GB of host RAM: stage weights from memory-mapped shards

CLAUDE.md

@@ -518,6 +518,21 @@ at 768x512x121: text 12.4 GB; render 20.1 GB (host) / 16.0 GB (disk)
 on 8 GB, 14.9 GB (disk) on 12 GB. Set `MALLOC_MMAP_THRESHOLD_=1048576`
 so freed tensors return to the OS.
 
+10 GB machines: `stage_from = "disk"` on `ltx23_load_pipeline()` and
+`load_gemma4_nf4()` loads no weights at all and reads every staged
+tensor (Gemma layers, connectors, transformer resident + streamed
+blocks, VAEs, vocoder) from a private memory map of the shards
+(`.mmap_reader`: bytes -> device -> `torch_view_copy` to the dtype, since
+R torch has no zero-copy `view(dtype)` and the R safetensors writer
+leaves tensor data unaligned). Under a 10 GB cgroup cap (page cache
+counted), 8 GB card: text phase 1.5 GB anon, render <= 4.6 GB anon,
+warm render 127 s (8.3 s/step) vs 82 s with host copies; latents
+bit-identical. Gotchas: offsets past 2^31 need `torch_as_strided(...,
+storage_offset = <double>)` (narrow() overflows R's int; one 2.3 GB
+LTX-2.5 shard silently pushed every read onto the 0.6 GB/s get_tensor
+fallback), and CPU reads must clone (a write to a returned view
+changes later reads of the private map).
+
 Gotchas that bit once (host RAM): freed PINNED memory stays in torch's
 CachingHostAllocator, not the OS, so a pipeline you plan to drop must
 load unpinned (RSS stayed 9 GB after freeing pinned connectors). Pinned

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.2.2.20
+Version: 0.2.2.21
 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,27 @@
+# diffuseR 0.2.2.21
+
+* **LTX-2.5 in 10 GB of host RAM.** `stage_from = "disk"` on
+  `ltx23_load_pipeline()` and `load_gemma4_nf4()` keeps no weights in
+  host memory: loading reads none, and every phase reads what it needs
+  from a private memory map of the artifact shards (Gemma 4 layers,
+  connectors, the transformer's resident part per render and its
+  streamed blocks every step, VAEs, vocoder). The OS page cache keeps
+  what fits.
+  - Measured on a simulated 8 GB card under a 10 GB memory cap that
+    counts page cache, 768x512x121: the text phase holds 1.5 GB of
+    anonymous memory (was 12.4 GB), the render phase at most 4.6 GB
+    (was 16.0 GB), and a warm render takes 127 s (82 s with host
+    copies). Output is bit-identical.
+  - `ltx23_memory_profile()` and `recommend("ltx25")` return
+    `stage_from`, `"disk"` below 18 GB of available RAM.
+    `serve(model = "ltx25")` uses it wherever the profile would split
+    the text phase off, which a persistent server cannot do.
+  - Reads go through the mapped bytes at 4.5-6.4 GB/s from the page
+    cache (safetensors' reader: ~0.7 GB/s), including shards past
+    2 GB.
+* `ltx23_load_group()` gains `lazy` and returns `source_keys`;
+  `ltx23_load_transformer_nf4()` gains `lazy`.
+
 # diffuseR 0.2.2.20
 
 * `serve(model = "ltx25")` serves LTX-2.5 on `/v1/videos`. It loads

R/checkpoint_ltx23.R

@@ -152,16 +152,20 @@ ltx23_split_keys <- function(keys) {
 #'   parameter/buffer name, or NA to skip the key deliberately.
 #' @param verbose Logical. Report progress and coverage.
 #' @param gc_every Integer. Run \code{gc()} after this many tensors.
+#' @param lazy Logical. Map keys to destinations without reading any
+#'   tensor: the destinations are left as allocated, for a caller that
+#'   stages them from disk (\code{source_keys}) before any forward.
 #'
 #' @return Invisibly, a list with \code{unmapped} (checkpoint keys that
 #'   found no destination), \code{skipped} (keys the mapper declined),
-#'   and \code{unfilled} (module parameters/buffers never written).
-#'   A perfectly loaded group has zero \code{unmapped} and zero
+#'   \code{unfilled} (module parameters/buffers never written), and
+#'   \code{source_keys} (destination name -> checkpoint key). A
+#'   perfectly loaded group has zero \code{unmapped} and zero
 #'   \code{unfilled}.
 #'
 #' @export
 ltx23_load_group <- function(ckpt, keys, module, map_key = identity,
-                             verbose = TRUE, gc_every = 50L) {
+                             verbose = TRUE, gc_every = 50L, lazy = FALSE) {
     stopifnot(inherits(ckpt, "ltx23_checkpoint"))
 
     params <- module$named_parameters()
@@ -171,6 +175,7 @@ ltx23_load_group <- function(ckpt, keys, module, map_key = identity,
     unmapped <- character(0)
     skipped <- character(0)
     filled <- character(0)
+    source_keys <- character(0)
 
     n <- length(keys)
     torch::with_no_grad({
@@ -187,6 +192,11 @@ ltx23_load_group <- function(ckpt, keys, module, map_key = identity,
                 unmapped <- c(unmapped, key)
                 next
             }
+            source_keys[dest_name] <- key
+            if (lazy) {
+                filled <- c(filled, dest_name)
+                next
+            }
 
             value <- ckpt$handle$get_tensor(key)
             if (!identical(as.integer(dest$shape), as.integer(value$shape))) {
@@ -228,7 +238,8 @@ ltx23_load_group <- function(ckpt, keys, module, map_key = identity,
         }
     }
 
-    invisible(list(unmapped = unmapped, skipped = skipped, unfilled = unfilled))
+    invisible(list(unmapped = unmapped, skipped = skipped, unfilled = unfilled,
+                   source_keys = source_keys))
 }
 
 #' Summarize checkpoint key coverage

R/fp8_ltx23.R

@@ -250,11 +250,23 @@ ltx23_open_fp8_checkpoint <- function(dir) {
         }
     }
 
+    # Disk-backed staging re-reads through a memory map (.mmap_reader),
+    # built on first use; NULL where this torch cannot map files
+    reader <- NULL
     handle <- list(
                    get_tensor = function(key) {
         h <- key_to_handle[[key]]
         if (is.null(h)) stop("Key not found in fp8 shards: ", key)
         .st_read_or_breadcrumb(function() h$get_tensor(key), key_to_path[[key]])
+    },
+                   read_to = function(key, device) {
+        if (is.null(reader)) {
+            reader <<- .mmap_reader(shard_paths) %||% FALSE
+        }
+        if (isFALSE(reader)) {
+            return(handle$get_tensor(key)$to(device = device))
+        }
+        reader$read(key, device)
     }
     )
 

R/gemma4_text_encoder.R

@@ -354,17 +354,26 @@ gemma4_text_model <- torch::nn_module(
 #'   \code{options(diffuseR.pin_staging)}.
 #' @param stream Logical. Stage layer by layer (see
 #'   \code{\link{load_gemma4_nf4}}).
+#' @param stage_from "host" or "disk" (NF4 artifacts only; see
+#'   \code{\link{load_gemma4_nf4}}).
 #' @param verbose Logical.
 #'
 #' @return A \code{gemma4_text_model}.
 #'
 #' @export
 load_gemma4_text_encoder <- function(path, device = "cpu", dtype = "bfloat16",
                                      pin = getOption("diffuseR.pin_staging", TRUE),
-                                     stream = FALSE, verbose = TRUE) {
+                                     stream = FALSE, stage_from = c("host", "disk"),
+                                     verbose = TRUE) {
+    stage_from <- match.arg(stage_from)
     if (file.exists(file.path(path, "manifest.json"))) {
         return(load_gemma4_nf4(path, device = device, dtype = dtype,
-                               pin = pin, stream = stream, verbose = verbose))
+                               pin = pin, stream = stream,
+                               stage_from = stage_from, verbose = verbose))
+    }
+    if (stage_from == "disk") {
+        stop("stage_from = \"disk\" needs a Gemma 4 NF4 artifact directory",
+             call. = FALSE)
     }
     ckpt <- ltx23_open_checkpoint(path, require_version = NULL)
     meta <- ckpt$handle$metadata[["__metadata__"]]

R/memory_ltx23.R

@@ -46,15 +46,22 @@ NULL
 #' text phase splits off; below 22 GB the resident blocks come from
 #' disk (measured on LTX-2.5 at 768x512x121 with an 8 GB card: the
 #' render phase peaks at 20.1 GB of host RAM with host copies and 16.0
-#' GB without, for ~8 s more per render).
+#' GB without, for ~8 s more per render). Below 18 GB,
+#' \code{stage_from = "disk"} (for \code{\link{ltx23_load_pipeline}} and
+#' \code{\link{load_gemma4_nf4}}) keeps no weights in host memory: every
+#' phase reads them from the memory-mapped artifacts. Measured the same
+#' way under a 10 GB memory cap that also counts page cache: the text
+#' phase holds 1.5 GB of anonymous memory, the render phase at most
+#' 4.6 GB, and a warm render takes 127 s instead of 82 s.
 #'
 #' @param vram_gb Numeric or NULL (auto-detect free VRAM).
 #' @param host_ram_gb Numeric or NULL (auto-detect available RAM; NA
 #'   when undetectable, which keeps the all-in-RAM defaults).
 #'
 #' @return Named list with device/dtype placement, \code{attn_chunk},
 #'   \code{pin_weights}, \code{resident_gb}, \code{stream_text},
-#'   \code{split_text}, \code{resident_from}, and resolution caps.
+#'   \code{split_text}, \code{resident_from}, \code{stage_from}, and
+#'   resolution caps.
 #'
 #' @export
 ltx23_memory_profile <- function(vram_gb = NULL, host_ram_gb = NULL) {
@@ -141,16 +148,25 @@ ltx23_memory_profile <- function(vram_gb = NULL, host_ram_gb = NULL) {
     } else {
         "host"
     }
+    out$stage_from <- if (gpu && ram_known &&
+                          host_ram_gb < .ltx23_ram_host_staged) {
+        "disk"
+    } else {
+        "host"
+    }
     out
 }
 
 # Host RAM (GB) for the LTX GPU tiers: the render pipeline plus the
-# Gemma 4 encoder in one process, and the render pipeline alone with
-# its resident blocks kept in host memory (LTX-2.5 768x512x121 on a
-# simulated 8 GB card, high-water marks: 20.1 GB with host copies,
-# 16.0 GB re-reading from disk, 12.4 GB for the split text phase)
+# Gemma 4 encoder in one process, the render pipeline alone with its
+# resident blocks kept in host memory, and with only the streamed
+# blocks kept there (LTX-2.5 768x512x121 on a simulated 8 GB card,
+# high-water marks: 20.1 GB with host copies, 16.0 GB re-reading the
+# resident part from disk, 12.4 GB for the split text phase). Below
+# the last, stage_from = "disk" keeps no weights in host memory.
 .ltx23_ram_one_process <- 32
 .ltx23_ram_host_resident <- 22
+.ltx23_ram_host_staged <- 18
 
 #' Tune the torch CUDA allocator for large-resident inference
 #'

R/mmap_reader.R

@@ -0,0 +1,98 @@
+# Memory-mapped safetensors reads for disk-backed staging. Each shard is
+# mapped privately (copy-on-write, so nothing can write back into the
+# artifact) as one uint8 tensor; a tensor's bytes are a narrow of that
+# map, copied to the target device and only then reinterpreted as their
+# dtype. R torch has no zero-copy Tensor$view(dtype), and the shards
+# the R safetensors writer produces do not align tensor data to the
+# element size, so a typed view of the host bytes is not available.
+# Measured on an LTX-2.5 shard: 4.5 GB/s to the GPU from the page cache
+# (safetensors' get_tensor: 0.68 GB/s), ~1.4 GB/s cold from NVMe, and
+# no anonymous host memory: the bytes stay page cache the kernel can
+# drop.
+
+.st_dtypes <- function() {
+    list(BF16 = torch::torch_bfloat16(), F16 = torch::torch_float16(),
+         F32 = torch::torch_float32(), F64 = torch::torch_float64(),
+         I64 = torch::torch_long(), I32 = torch::torch_int32(),
+         I16 = torch::torch_int16(), I8 = torch::torch_int8(),
+         U8 = torch::torch_uint8(), BOOL = torch::torch_bool(),
+         F8_E4M3 = torch::torch_float8_e4m3fn(),
+         F8_E5M2 = torch::torch_float8_e5m2())
+}
+
+# The header of a safetensors file: the byte where tensor data starts,
+# and each tensor's dtype, shape and [start, end) data offsets
+.st_header <- function(path) {
+    con <- file(path, "rb")
+    on.exit(close(con))
+    n <- readBin(con, "integer", n = 1L, size = 8L, endian = "little")
+    hdr <- jsonlite::fromJSON(rawToChar(readBin(con, "raw", n = n)),
+                              simplifyVector = FALSE)
+    hdr[["__metadata__"]] <- NULL
+    list(data_start = 8 + n, tensors = hdr)
+}
+
+# A reader over `paths`, or NULL when this torch cannot map files (the
+# caller falls back to safetensors' get_tensor). Shards are mapped on
+# first use. read(key, device) returns the tensor on `device`.
+.mmap_reader <- function(paths) {
+    from_file <- get0("torch_from_file", envir = asNamespace("torch"))
+    view_copy <- get0("torch_view_copy", envir = asNamespace("torch"))
+    if (!is.function(from_file) || !is.function(view_copy)) {
+        return(NULL)
+    }
+    dtypes <- .st_dtypes()
+    index <- new.env(parent = emptyenv(), size = 8192L)
+    for (i in seq_along(paths)) {
+        h <- .st_header(paths[[i]])
+        for (k in names(h$tensors)) {
+            info <- h$tensors[[k]]
+            off <- unlist(info$data_offsets)
+            # Byte offset kept as a double: shards can pass 2^31 bytes
+            index[[k]] <- list(shard = i,
+                               offset = h$data_start + off[[1]],
+                               len = as.integer(off[[2]] - off[[1]]),
+                               dtype = info$dtype,
+                               shape = as.integer(unlist(info$shape)))
+        }
+    }
+    maps <- vector("list", length(paths))
+    shard_map <- function(i) {
+        if (is.null(maps[[i]])) {
+            maps[[i]] <<- from_file(paths[[i]], shared = FALSE,
+                                    size = file.size(paths[[i]]),
+                                    options = list(dtype = torch::torch_uint8()))
+        }
+        maps[[i]]
+    }
+    read <- function(key, device = "cpu") {
+        e <- index[[key]]
+        if (is.null(e)) {
+            stop("Key not found in the mapped shards: ", key, call. = FALSE)
+        }
+        dtype <- dtypes[[e$dtype]]
+        if (is.null(dtype)) {
+            stop("Unsupported safetensors dtype ", e$dtype, " for ", key,
+                 call. = FALSE)
+        }
+        if (e$len == 0L) {
+            return(torch::torch_empty(e$shape, dtype = dtype, device = device))
+        }
+        # as_strided takes the offset as int64; narrow() would cap it at
+        # R's integer range
+        t <- torch::torch_as_strided(shard_map(e$shard), e$len, 1L,
+                                     storage_offset = e$offset)$to(device = device)
+        # A move to the GPU lands in fresh, aligned storage. On the CPU
+        # the bytes are still the map: copy them, so the result is
+        # aligned for the dtype view and a write to it cannot change
+        # what later reads of the map return
+        if (identical(t$device$type, "cpu")) {
+            t <- t$clone()
+        }
+        if (!identical(e$dtype, "U8")) {
+            t <- view_copy(t, dtype = dtype)
+        }
+        if (length(e$shape)) t$view(e$shape) else t$squeeze()
+    }
+    list(read = read, keys = function() ls(index))
+}

R/nf4_ltx23.R

@@ -397,13 +397,17 @@ ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = NULL,
 #'   (\code{\link{ltx23_open_fp8_checkpoint}} reads any shard artifact).
 #' @param device Character.
 #' @param verbose Logical.
+#' @param lazy Logical. Build the model and record each tensor's
+#'   artifact key without reading any weight; every weight is a
+#'   zero-stride placeholder until a disk-backed staging set reads it
+#'   onto the device. Host RAM then never holds the transformer.
 #' @param ... Passed to \code{\link{ltx23_transformer}} (tiny test configs).
 #'
 #' @return The loaded \code{ltx23_transformer}.
 #'
 #' @export
 ltx23_load_transformer_nf4 <- function(ckpt, device = "cuda", verbose = TRUE,
-                                       ...) {
+                                       lazy = FALSE, ...) {
     stopifnot(inherits(ckpt, "ltx23_checkpoint"))
     args <- utils::modifyList(ltx23_dit_config_args(ckpt$config), list(...))
     model <- do.call(.construct_skeleton, c(list(ltx23_transformer), args))
@@ -446,10 +450,15 @@ ltx23_load_transformer_nf4 <- function(ckpt, device = "cuda", verbose = TRUE,
                     # loads through the pre-swap destination map
                     nf4_mod$bias <- old$bias
                 }
-                nf4_mod$set_nf4_weight(
-                                       ckpt$handle$get_tensor(key),
-                                       ckpt$handle$get_tensor(paste0(key, "_absmax"))
-                )
+                if (lazy) {
+                    .release_host(nf4_mod$weight_nf4)
+                    .release_host(nf4_mod$weight_absmax)
+                } else {
+                    nf4_mod$set_nf4_weight(
+                                           ckpt$handle$get_tensor(key),
+                                           ckpt$handle$get_tensor(paste0(key, "_absmax"))
+                    )
+                }
                 do.call(`$<-`, list(parent, leaf, nf4_mod))
                 filled <- c(filled, mapped)
             } else {
@@ -458,7 +467,11 @@ ltx23_load_transformer_nf4 <- function(ckpt, device = "cuda", verbose = TRUE,
                     unmapped <- c(unmapped, key)
                     next
                 }
-                dest$copy_(ckpt$handle$get_tensor(key))
+                if (lazy) {
+                    .release_host(dest)
+                } else {
+                    dest$copy_(ckpt$handle$get_tensor(key))
+                }
                 filled <- c(filled, mapped)
                 source_keys[mapped] <- key
             }
@@ -493,7 +506,11 @@ ltx23_load_transformer_nf4 <- function(ckpt, device = "cuda", verbose = TRUE,
     # ~1.5GB per block at high resolution; per-block gc keeps them bounded
     options(diffuseR.block_gc = TRUE)
     if (verbose) {
-        message("Transformer ready: NF4 weights resident on ", device)
+        if (lazy) {
+            message("Transformer ready: NF4 weights left on disk until staged")
+        } else {
+            message("Transformer ready: NF4 weights resident on ", device)
+        }
     }
     model
 }