From 8365ecc53888c13ac45180b8c4369a5d5cee4ff5 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 29 Oct 2026 14:23:05 -0500 Subject: [PATCH 1/2] LTX-2.5 in 10 GB of host RAM: stage every weight from memory-mapped shards stage_from = "disk" on ltx23_load_pipeline() and load_gemma4_nf4() loads no weights and reads each staged tensor (Gemma 4 layers, connectors, transformer resident and streamed blocks, VAEs, vocoder) from a private memory map of the artifact shards when its phase needs it. .mmap_reader copies a tensor's bytes to the device and reinterprets the dtype there (torch_view_copy), so unaligned shard data works and no anonymous host memory holds weights. Measured on a simulated 8 GB card under a 10 GB memory cap that counts page cache, 768x512x121: text phase 1.5 GB anon, render at most 4.6 GB anon, warm render 127 s (82 s with host copies), latents bit-identical. ltx23_memory_profile() and recommend("ltx25") return stage_from ("disk" below 18 GB available), and serve(model = "ltx25") stages from disk where a split text phase would be needed. --- CLAUDE.md | 15 +++++ R/checkpoint_ltx23.R | 19 ++++-- R/fp8_ltx23.R | 12 ++++ R/gemma4_text_encoder.R | 13 +++- R/memory_ltx23.R | 28 +++++++-- R/mmap_reader.R | 98 +++++++++++++++++++++++++++++++ R/nf4_ltx23.R | 31 +++++++--- R/quantize_gemma4.R | 46 +++++++++++++-- R/recommend.R | 25 ++++---- R/serve.R | 24 +++++--- R/staging.R | 46 ++++++++++++--- R/stream_ltx23.R | 46 ++++++++------- R/txt2vid_ltx23.R | 69 +++++++++++++++++----- inst/tinytest/test_dit_ltx23.R | 27 +++++++++ inst/tinytest/test_memory_ltx23.R | 12 ++++ inst/tinytest/test_mmap_reader.R | 83 ++++++++++++++++++++++++++ inst/tinytest/test_recommend.R | 3 + man/load_gemma4_nf4.Rd | 6 ++ man/load_gemma4_text_encoder.Rd | 4 ++ man/ltx23_load_group.Rd | 12 +++- man/ltx23_load_pipeline.Rd | 14 +++++ man/ltx23_load_transformer_nf4.Rd | 13 +++- man/ltx23_memory_profile.Rd | 11 +++- man/recommend.Rd | 16 ++--- man/serve.Rd | 6 +- 25 files changed, 575 insertions(+), 104 deletions(-) create mode 100644 R/mmap_reader.R create mode 100644 inst/tinytest/test_mmap_reader.R diff --git a/CLAUDE.md b/CLAUDE.md index ac869d9..758c5ef 100644 --- a/CLAUDE.md +++ b/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 = )` (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 diff --git a/R/checkpoint_ltx23.R b/R/checkpoint_ltx23.R index 3c14762..f536076 100644 --- a/R/checkpoint_ltx23.R +++ b/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 diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index aeb737e..fd9e3aa 100644 --- a/R/fp8_ltx23.R +++ b/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) } ) diff --git a/R/gemma4_text_encoder.R b/R/gemma4_text_encoder.R index c852cbe..0a52180 100644 --- a/R/gemma4_text_encoder.R +++ b/R/gemma4_text_encoder.R @@ -354,6 +354,8 @@ 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}. @@ -361,10 +363,17 @@ gemma4_text_model <- torch::nn_module( #' @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__"]] diff --git a/R/memory_ltx23.R b/R/memory_ltx23.R index f8178eb..58702bd 100644 --- a/R/memory_ltx23.R +++ b/R/memory_ltx23.R @@ -46,7 +46,13 @@ 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 @@ -54,7 +60,8 @@ NULL #' #' @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 #' diff --git a/R/mmap_reader.R b/R/mmap_reader.R new file mode 100644 index 0000000..60e3f5e --- /dev/null +++ b/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)) +} diff --git a/R/nf4_ltx23.R b/R/nf4_ltx23.R index e717e81..84debd8 100644 --- a/R/nf4_ltx23.R +++ b/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 } diff --git a/R/quantize_gemma4.R b/R/quantize_gemma4.R index bffdbdd..de3136a 100644 --- a/R/quantize_gemma4.R +++ b/R/quantize_gemma4.R @@ -168,6 +168,10 @@ gemma4_quantize_nf4 <- function(path, output_dir = NULL, shard_bytes = 1.9e9, #' layer by layer instead of whole: an encode on the GPU then holds #' only the embeddings and one layer at a time, for cards that cannot #' hold the ~8 GB NF4 model. +#' @param stage_from "host" keeps the staged weights in host memory; +#' "disk" keeps none and reads each layer from the artifact (memory +#' mapped) when the encode needs it, so the encoder costs almost no +#' host RAM. "disk" needs \code{device = "cpu"} and a CUDA encode. #' @param verbose Logical. #' #' @return A \code{gemma4_text_model}. @@ -175,7 +179,14 @@ gemma4_quantize_nf4 <- function(path, output_dir = NULL, shard_bytes = 1.9e9, #' @export load_gemma4_nf4 <- function(artifact_dir, device = "cuda", 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) + lazy <- stage_from == "disk" + if (lazy && device != "cpu") { + stop("stage_from = \"disk\" stages through the GPU from a CPU ", + "model; use device = \"cpu\"", call. = FALSE) + } ckpt <- ltx23_open_fp8_checkpoint(artifact_dir) manifest <- jsonlite::fromJSON(file.path(artifact_dir, "manifest.json"), simplifyVector = FALSE) @@ -194,6 +205,8 @@ load_gemma4_nf4 <- function(artifact_dir, device = "cuda", dtype = "bfloat16", model <- .construct_skeleton(gemma4_text_model, config, dtype = torch_dtype) cast_keys <- sub("_absmax$", "", grep("_absmax$", ckpt$keys, value = TRUE)) + # Model tensor name -> artifact key, for disk-backed staging + source_keys <- character(0) torch::with_no_grad({ for (key in cast_keys) { # layers....weight @@ -202,8 +215,16 @@ load_gemma4_nf4 <- function(artifact_dir, device = "cuda", dtype = "bfloat16", old <- parent[[segs[4]]] nf4 <- ltx23_nf4_linear(old$weight$shape[1], old$weight$shape[2], bias = FALSE) - nf4$set_nf4_weight(ckpt$handle$get_tensor(key), - ckpt$handle$get_tensor(paste0(key, "_absmax"))) + stem <- sub("\\.weight$", "", key) + source_keys[paste0(stem, c(".weight_nf4", ".weight_absmax"))] <- + c(key, paste0(key, "_absmax")) + if (lazy) { + .release_host(nf4$weight_nf4) + .release_host(nf4$weight_absmax) + } else { + nf4$set_nf4_weight(ckpt$handle$get_tensor(key), + ckpt$handle$get_tensor(paste0(key, "_absmax"))) + } do.call(`$<-`, list(parent, segs[4], nf4)) } @@ -216,7 +237,12 @@ load_gemma4_nf4 <- function(artifact_dir, device = "cuda", dtype = "bfloat16", if (is.null(dest)) { next # e.g. text_embedding_projection.*, for the connectors } - dest$copy_(ckpt$handle$get_tensor(key)) + source_keys[key] <- key + if (lazy) { + .release_host(dest) + } else { + dest$copy_(ckpt$handle$get_tensor(key)) + } filled <- c(filled, key) } unfilled <- setdiff(names(dests), filled) @@ -229,7 +255,17 @@ load_gemma4_nf4 <- function(artifact_dir, device = "cuda", dtype = "bfloat16", model$to(device = device) model$eval() - if (stream && device == "cpu") { + if (lazy) { + src <- list(keys = source_keys, handle = ckpt$handle) + attr(model, "staging") <- if (stream) { + .stage_layer_stream(model, model$layers, "layers.", pin, src) + } else { + .stage_set(c(model$parameters, model$buffers), pin, src) + } + if (verbose) { + message("Gemma 4 weights staged from disk") + } + } else if (stream && device == "cpu") { attr(model, "staging") <- .stage_layer_stream(model, model$layers, "layers.", pin) if (verbose) { diff --git a/R/recommend.R b/R/recommend.R index e708829..94d6ed9 100644 --- a/R/recommend.R +++ b/R/recommend.R @@ -31,11 +31,12 @@ #' transformer resident and streams the rest, and the result carries #' the \code{\link{ltx23_memory_profile}} fields that configure that #' (\code{resident_gb}, \code{stream_text}, \code{split_text}, -#' \code{resident_from}). Its staging is always pinned on the GPU -#' tiers: block streaming from pageable memory is slow, and the -#' host-RAM fields already bound the footprint (measured at 768x512x121 -#' on an 8 GB card: 20.1 GB high-water with resident host copies, 16.0 -#' GB with \code{resident_from = "disk"}). +#' \code{resident_from}, \code{stage_from}). Its staging is always +#' pinned on the GPU tiers: block streaming from pageable memory is +#' slow, and the host-RAM fields already bound the footprint (measured +#' at 768x512x121 on an 8 GB card: 20.1 GB high-water with resident +#' host copies, 16.0 GB with \code{resident_from = "disk"}; with +#' \code{stage_from = "disk"} no weights are held in host memory). #' #' The pinning decision: phase-swapped weights are page-locked host #' copies (see \code{\link{staging}}) that transfer at DMA rate - @@ -69,8 +70,9 @@ #' \code{host_ram_gb}, \code{fork_suggested} (logical: the installed #' safetensors cannot read the tier the card could otherwise run), and #' \code{note} (the suggestion string, or NULL). For \code{"ltx25"} it -#' also has \code{resident_gb}, \code{stream_text}, \code{split_text} -#' and \code{resident_from} (see \code{\link{ltx23_memory_profile}}). +#' also has \code{resident_gb}, \code{stream_text}, \code{split_text}, +#' \code{resident_from} and \code{stage_from} (see +#' \code{\link{ltx23_memory_profile}}). #' The \code{fork_suggested} field name is #' historical -- the remedy used to be the cornball-ai fork and is now #' \code{install.packages("safetensors")}, since the fixes reached CRAN @@ -175,13 +177,14 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", if (model == "ltx25") { if (isTRUE(chosen$cpu)) { prof <- list(resident_gb = Inf, stream_text = FALSE, - split_text = FALSE, resident_from = "host") + split_text = FALSE, resident_from = "host", + stage_from = "host") } else { prof <- ltx23_memory_profile(vram_gb, host_ram_gb = host_ram_gb) } - out[c("resident_gb", "stream_text", "split_text", "resident_from")] <- - prof[c("resident_gb", "stream_text", "split_text", - "resident_from")] + fields <- c("resident_gb", "stream_text", "split_text", + "resident_from", "stage_from") + out[fields] <- prof[fields] } out } diff --git a/R/serve.R b/R/serve.R index 6e51b84..049f265 100644 --- a/R/serve.R +++ b/R/serve.R @@ -66,8 +66,10 @@ #' \code{gemma4-nf4} artifacts and places them by #' \code{\link{recommend}("ltx25")}: on 8-12 GB cards the transformer #' streams blocks and, below 12 GB, Gemma 4 streams layers. The -#' encoder and render pipeline share one process, which needs about -#' 31 GB of host RAM. +#' encoder and render pipeline share one process: with 32 GB of host +#' RAM available they keep their weights there, with less they read +#' them from the memory-mapped artifacts as each phase needs them +#' (\code{stage_from = "disk"} in \code{\link{ltx23_load_pipeline}}). #' @param device Character. "cuda" or "cpu". #' @param token Character or NULL. Shared secret; when set, requests #' must send \code{Authorization: Bearer }. @@ -322,7 +324,8 @@ serve <- function(port = 7812L, model = c("flux2", "zimage", "flux1", "ltx", # the GPU and streams the rest; below 12 GB Gemma 4 streams its layers. # The server holds the encoder and the render pipeline in one process # (~31 GB of host RAM with pinned staging); the split text phase of -# ltx23_text_embeds() does not apply to a persistent server. +# ltx23_text_embeds() does not apply to a persistent server, so with +# less RAM than that it stages every weight from disk instead. .dserve_load_ltx25 <- function(device, max_prompts) { data_dir <- tools::R_user_dir("diffuseR", "data") ckpt <- file.path(data_dir, "ltx2.5-nf4-full") @@ -338,18 +341,21 @@ serve <- function(port = 7812L, model = c("flux2", "zimage", "flux1", "ltx", } gpu <- identical(device, "cuda") && torch::cuda_is_available() rec <- recommend("ltx25", vram_gb = if (gpu) NULL else 0) - if (isTRUE(rec$split_text)) { + # A persistent server cannot split the text phase off, so where the + # profile asks for that it keeps no weights in host memory at all + stage_from <- if (isTRUE(rec$split_text)) "disk" else rec$stage_from + if (stage_from == "disk") { message(sprintf(paste0( - "%.0f GB of host RAM available; serving LTX-2.5 holds Gemma 4 ", - "and the render pipeline together (~31 GB measured)."), - rec$host_ram_gb)) + "%.0f GB of host RAM available; LTX-2.5 weights are read from ", + "disk as each phase needs them."), rec$host_ram_gb)) } pipe <- ltx23_load_pipeline(ckpt, device = device, pin = rec$pin, resident_gb = rec$resident_gb, resident_from = rec$resident_from, - verbose = FALSE) + stage_from = stage_from, verbose = FALSE) te <- load_gemma4_nf4(gem, device = "cpu", pin = rec$pin, - stream = rec$stream_text, verbose = FALSE) + stream = rec$stream_text, stage_from = stage_from, + verbose = FALSE) tok <- gemma4_tokenizer(gem) enc_dev <- if (gpu) "cuda" else "cpu" embeds <- .dserve_embed_cache(function(prompt) { diff --git a/R/staging.R b/R/staging.R index 9a489b1..4ed0382 100644 --- a/R/staging.R +++ b/R/staging.R @@ -172,8 +172,8 @@ NULL # Disk-backed pair (.stage_from_disk): read it again. The host # copies are garbage once on the device; collect them as we # go so the read never holds more than a slice in RAM - pair$live$set_data(pair$read()$to(device = device, - dtype = pair$live$dtype)) + pair$live$set_data(pair$read(device)$to(device = device, + dtype = pair$live$dtype)) pair$state$loaded <- TRUE n_read <- n_read + 1L if (n_read %% 64L == 0L) { @@ -192,8 +192,10 @@ NULL # Staging pairs that keep NO host copy: `keys` names each tensor's # artifact key (a model's weight_source$keys) and `handle` reads it. # Offload drops the device copy and onload re-reads the tensor from -# disk (the OS page cache, when RAM allows). Tensors without a key -# (computed buffers) get ordinary pairs. Releases the host copies now. +# disk (the OS page cache, when RAM allows): through the handle's +# memory map when it has one (read_to, see .mmap_reader), else +# get_tensor. Tensors without a key (computed buffers) get ordinary +# pairs. Releases the host copies now. .stage_from_disk <- function(tensors, keys, handle, pin) { on_disk <- names(tensors) %in% names(keys) pairs <- vector("list", length(tensors)) @@ -206,7 +208,11 @@ NULL pair <- list(live = tensors[[nm]], state = state, read = local({ k <- key - function() handle$get_tensor(k) + if (is.function(handle$read_to)) { + function(device) handle$read_to(k, device) + } else { + function(device) handle$get_tensor(k) + } })) .release_host(pair$live) pairs[[nm]] <- pair @@ -215,6 +221,26 @@ NULL pairs } +# Staging pairs for `tensors` (named by their full model names): from +# disk when `src` is a weight source (list(keys, handle)), else host +# pairs, pinned when `pin` +.stage_set <- function(tensors, pin, src = NULL) { + if (is.null(src)) { + .stage_pairs(tensors, pin) + } else { + .stage_from_disk(tensors, src$keys, src$handle, pin) + } +} + +# The tensors of `module` (parameters and buffers) whose full names +# start with any of `prefixes` +.tensors_under <- function(module, prefixes) { + tensors <- c(module$parameters, module$buffers) + keep <- Reduce(`|`, lapply(prefixes, function(p) startsWith(names(tensors), p)), + FALSE) + tensors[keep] +} + # Point a tensor at a zero-stride placeholder of the same shape and # dtype: its storage is a single element, so the old data is freed .release_host <- function(t) { @@ -228,13 +254,15 @@ NULL # (whose parameter names start with `prefix` in the model) gets its own # staging set on model$layer_stream, which the encoder's forward swaps # in and out around that layer; the returned pairs cover the rest of the -# model (embeddings, final norm) for the encode-phase onload. -.stage_layer_stream <- function(model, layers, prefix, pin) { +# model (embeddings, final norm) for the encode-phase onload. With a +# weight source `src` (list(keys, handle)) every set re-reads from disk. +.stage_layer_stream <- function(model, layers, prefix, pin, src = NULL) { model$layer_stream <- lapply(seq_along(layers), function(i) { - .stage_pairs(c(layers[[i]]$parameters, layers[[i]]$buffers), pin) + .stage_set(.tensors_under(model, paste0(prefix, i - 1L, ".")), pin, + src) }) tensors <- c(model$parameters, model$buffers) - .stage_pairs(tensors[!startsWith(names(tensors), prefix)], pin) + .stage_set(tensors[!startsWith(names(tensors), prefix)], pin, src) } #' Return a pinned component to the CPU diff --git a/R/stream_ltx23.R b/R/stream_ltx23.R index 9375e62..58e8e9b 100644 --- a/R/stream_ltx23.R +++ b/R/stream_ltx23.R @@ -20,18 +20,20 @@ # Stream the text connectors (5.9 GB in bf16): the two input projections # (1.5 and 0.8 GB) as units, each connector's blocks one at a time. -# Returns the staging pairs for the rest (registers, norms, rope). -.ltx23_connector_stream <- function(connectors, pin) { +# Returns the staging pairs for the rest (registers, norms, rope). With +# a weight source `src` (list(keys, handle)) every set re-reads from +# disk instead of holding a host copy. +.ltx23_connector_stream <- function(connectors, pin, src = NULL) { units <- c("video_text_proj_in", "audio_text_proj_in") connectors$unit_stream <- stats::setNames(lapply(units, function(u) { - m <- connectors[[u]] - .stage_pairs(c(m$parameters, m$buffers), pin) + .stage_set(.tensors_under(connectors, paste0(u, ".")), pin, src) }), units) for (cn in c("video_connector", "audio_connector")) { con <- connectors[[cn]] con$layer_stream <- lapply(seq_along(con$transformer_blocks), function(i) { - b <- con$transformer_blocks[[i]] - .stage_pairs(c(b$parameters, b$buffers), pin) + .stage_set(.tensors_under(connectors, + paste0(cn, ".transformer_blocks.", i - 1L, ".")), + pin, src) }) } tensors <- c(connectors$parameters, connectors$buffers) @@ -39,7 +41,7 @@ startsWith(names(tensors), "audio_text_proj_in.") | startsWith(names(tensors), "video_connector.transformer_blocks.") | startsWith(names(tensors), "audio_connector.transformer_blocks.") - .stage_pairs(tensors[!streamed], pin) + .stage_set(tensors[!streamed], pin, src) } # Bytes held by a module's parameters and buffers @@ -54,9 +56,10 @@ # consumed by the forward) and returns the staging pairs for everything # else, which replace the transformer's whole-module staging set. With # `from_disk`, that resident part keeps no host copy: it is re-read from -# the artifact at each onload (.stage_from_disk). +# the artifact at each onload (.stage_from_disk); `stream_from_disk` +# does the same for the streamed blocks, every step. .ltx23_block_stream <- function(transformer, resident_gb, pin, - from_disk = FALSE) { + from_disk = FALSE, stream_from_disk = FALSE) { blocks <- transformer$transformer_blocks n <- length(blocks) block_bytes <- vapply(seq_len(n), function(i) .module_bytes(blocks[[i]]), 0) @@ -64,9 +67,20 @@ budget <- resident_gb * 1024 ^ 3 - other_bytes n_resident <- sum(cumsum(block_bytes) <= budget) + src <- NULL + if (from_disk || stream_from_disk) { + src <- transformer$weight_source + if (is.null(src)) { + stop("Staging from disk needs an NF4 transformer ", + "(ltx23_load_transformer_nf4 records its artifact keys)", + call. = FALSE) + } + } streamed <- if (n_resident < n) seq(n_resident + 1L, n) else integer(0) stream_staging <- lapply(streamed, function(i) { - .stage_pairs(c(blocks[[i]]$parameters, blocks[[i]]$buffers), pin) + .stage_set(.tensors_under(transformer, + paste0("transformer_blocks.", i - 1L, ".")), + pin, if (stream_from_disk) src) }) # Everything not in a streamed block, by identity of the tensor names names_all <- c(names(transformer$parameters), names(transformer$buffers)) @@ -80,17 +94,7 @@ } else { NULL } - staging <- if (from_disk) { - src <- transformer$weight_source - if (is.null(src)) { - stop("resident_from = \"disk\" needs an NF4 transformer ", - "(ltx23_load_transformer_nf4 records its artifact keys)", - call. = FALSE) - } - .stage_from_disk(tensors[!in_stream], src$keys, src$handle, pin) - } else { - .stage_pairs(tensors[!in_stream], pin) - } + staging <- .stage_set(tensors[!in_stream], pin, if (from_disk) src) list(staging = staging, n_resident = n_resident, n_streamed = length(streamed), resident_gb = (other_bytes + sum(block_bytes[seq_len(n_resident)])) / 1024 ^ 3, diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 55a2671..31daafe 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -279,6 +279,18 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { #' (the OS page cache serves it when RAM allows). \code{"disk"} saves #' host RAM equal to the resident part (up to ~12 GB) for a re-read #' per render. Needs \code{phase_offload} on CUDA. +#' @param stage_from Where every phase-offloaded weight lives between +#' uses: \code{"host"} (default) in host memory; \code{"disk"} nowhere +#' but the artifact. Loading then reads no weights, and each onload +#' (the transformer's resident part per render, its streamed blocks +#' every step, the connectors, VAEs and vocoder per phase) reads from +#' the memory-mapped shards, so the pipeline holds almost no host +#' RAM and the OS page cache keeps what fits. Implies +#' \code{resident_from = "disk"}; the connectors always stream. For +#' machines with less host RAM than the model; each step re-reads the +#' streamed blocks (~1.4 GB/s cold from NVMe, ~4.5 GB/s from the page +#' cache). Needs an NF4 artifact directory and \code{phase_offload} +#' on CUDA. #' @param verbose Logical. #' #' @return A list with the loaded modules and the checkpoint config, @@ -292,14 +304,17 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", pin = TRUE, attn_chunk = NULL, phase_offload = TRUE, resident_gb = Inf, resident_from = c("host", "disk"), + stage_from = c("host", "disk"), verbose = TRUE) { resident_from <- match.arg(resident_from) - from_disk <- resident_from == "disk" + stage_from <- match.arg(stage_from) + stage_disk <- stage_from == "disk" + from_disk <- resident_from == "disk" || stage_disk if ((is.finite(resident_gb) || from_disk) && !(phase_offload && device == "cuda")) { - stop("resident_gb and resident_from = \"disk\" stage the transformer ", - "through the GPU; they need device = \"cuda\" and ", - "phase_offload = TRUE", call. = FALSE) + stop("resident_gb, resident_from = \"disk\" and stage_from = ", + "\"disk\" stage the transformer through the GPU; they need ", + "device = \"cuda\" and phase_offload = TRUE", call. = FALSE) } if (phase_offload) { component_device <- "cpu" @@ -380,20 +395,27 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", } # Modules arrive as skeletons already at torch_dtype res <- ltx23_load_group(ckpt, groups[[name]], module, - map_key = map_key, verbose = verbose) + map_key = map_key, verbose = verbose, + lazy = stage_disk) if (length(res$unmapped) || length(res$unfilled)) { stop(name, ": incomplete load (", length(res$unmapped), " unmapped, ", length(res$unfilled), " unfilled)") } module$to(device = dev) module$eval() + module$weight_source <- list(keys = res$source_keys, handle = ckpt$handle) module } + if (stage_disk && !(fp8 && identical(ckpt$format, "nf4"))) { + stop("stage_from = \"disk\" needs an NF4 artifact directory", + call. = FALSE) + } if ("dit" %in% components) { if (fp8 && identical(ckpt$format, "nf4")) { pipe$transformer <- ltx23_load_transformer_nf4( - ckpt, device = component_device, verbose = verbose + ckpt, device = component_device, verbose = verbose, + lazy = stage_disk ) } else if (fp8) { pipe$transformer <- ltx23_load_transformer_fp8( @@ -433,12 +455,16 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", voc <- ltx23_vocoder_with_bwe() voc$to(dtype = torch_dtype) res <- ltx23_load_group(ckpt, groups$vocoder, voc, - map_key = ltx23_map_vocoder_key, verbose = verbose) + map_key = ltx23_map_vocoder_key, verbose = verbose, + lazy = stage_disk) if (length(res$unmapped) || length(res$unfilled)) { stop("vocoder: incomplete load") } voc$to(device = component_device, dtype = torch::torch_float32()) voc$eval() + # Disk reads arrive in the artifact's dtype; the onload casts + # them to the live float32 + voc$weight_source <- list(keys = res$source_keys, handle = ckpt$handle) pipe$vocoder <- voc } @@ -452,24 +478,37 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", # blocks carry their own staging sets (staged even unpinned, so # the phase onload never moves them) split <- .ltx23_block_stream(pipe$transformer, resident_gb, - pin = pin_staging, from_disk = from_disk) + pin = pin_staging, from_disk = from_disk, + stream_from_disk = stage_disk) staging$transformer <- split$staging if (verbose) { message(sprintf(paste0("Transformer: %d blocks resident (%.1f GB%s), ", - "%d streamed (%.1f GB/step)"), + "%d streamed (%.1f GB/step%s)"), split$n_resident, split$resident_gb, if (from_disk) ", re-read from disk" else "", - split$n_streamed, split$streamed_gb)) + split$n_streamed, split$streamed_gb, + if (stage_disk) ", from disk" else "")) } } # The 5.9 GB connector phase does not fit beside an 8 GB card's # activations either: stream it for budgets under 7 GB (also when - # the connectors load alone, for ltx23_text_embeds) - if (resident_gb < 7 && !is.null(pipe$connectors)) { - staging$connectors <- .ltx23_connector_stream(pipe$connectors, - pin = pin_staging) + # the connectors load alone, for ltx23_text_embeds), and whenever + # it comes from disk + if ((resident_gb < 7 || stage_disk) && !is.null(pipe$connectors)) { + staging$connectors <- .ltx23_connector_stream( + pipe$connectors, pin = pin_staging, + src = if (stage_disk) pipe$connectors$weight_source + ) if (verbose) { - message("Text connectors stream block by block") + message("Text connectors stream block by block", + if (stage_disk) " from disk" else "") + } + } + if (stage_disk) { + for (nm in intersect(c("vae", "audio_vae", "vocoder"), names(pipe))) { + m <- pipe[[nm]] + staging[[nm]] <- .stage_set(c(m$parameters, m$buffers), pin_staging, + m$weight_source) } } if (pin_staging) { diff --git a/inst/tinytest/test_dit_ltx23.R b/inst/tinytest/test_dit_ltx23.R index 52ed496..61a91e6 100644 --- a/inst/tinytest/test_dit_ltx23.R +++ b/inst/tinytest/test_dit_ltx23.R @@ -245,6 +245,33 @@ if (requireNamespace("safetensors", quietly = TRUE)) { # Back again for later tests diffuseR:::.staged_onload(split_disk$staging, "cpu") model$block_stream <- NULL + + # stage_from = "disk": the streamed block re-reads too, through the + # memory-mapped reader (read_to), every forward + reader <- diffuseR:::.mmap_reader(st_file) + if (!is.null(reader)) { + model$weight_source <- list(keys = stats::setNames(names(tensors), names(tensors)), + handle = list(get_tensor = function(k) stop("mapped reads only"), + read_to = reader$read)) + split_all_disk <- diffuseR:::.ltx23_block_stream(model, other_gb + 1.5 * blk_gb, + pin = FALSE, from_disk = TRUE, stream_from_disk = TRUE) + w2 <- model$transformer_blocks[[2]]$parameters[[1]] + expect_equal(w2$stride(), rep(0L, w2$dim())) + diffuseR:::.staged_onload(split_all_disk$staging, "cpu") + res_mm <- run_model( + encoder_attention_mask = fx$m_enc_mask, + audio_encoder_attention_mask = fx$m_enc_mask + ) + expect_true(max_abs_diff(res_mm$sample, res$sample) == 0) + expect_true(max_abs_diff(res_mm$audio_sample, res$audio_sample) == 0) + # The streamed block went back to its placeholder after the forward + expect_equal(w2$stride(), rep(0L, w2$dim())) + model$block_stream <- NULL + # Host copies for later tests + for (nm in names(tensors)) { + tensors[[nm]]$set_data(reader$read(nm, "cpu")) + } + } model$weight_source <- NULL unlink(st_file) } diff --git a/inst/tinytest/test_memory_ltx23.R b/inst/tinytest/test_memory_ltx23.R index 4629c38..a2457e7 100644 --- a/inst/tinytest/test_memory_ltx23.R +++ b/inst/tinytest/test_memory_ltx23.R @@ -78,19 +78,31 @@ expect_equal(mid$resident_from, "host") small <- ltx23_memory_profile(12, host_ram_gb = 16) expect_true(small$split_text) expect_equal(small$resident_from, "disk") +# Below the render pipeline's 16 GB with the resident part on disk, no +# weights stay in host memory at all +expect_equal(small$stage_from, "disk") +expect_equal(mid$stage_from, "host") +expect_equal(big$stage_from, "host") +expect_equal(ltx23_memory_profile(8, host_ram_gb = 10)$stage_from, "disk") # Unknown RAM keeps the defaults; the CPU tier never stages unknown <- ltx23_memory_profile(8, host_ram_gb = NA) expect_false(unknown$split_text) expect_equal(unknown$resident_from, "host") +expect_equal(unknown$stage_from, "host") cpu <- ltx23_memory_profile(4, host_ram_gb = 8) expect_false(cpu$split_text) expect_equal(cpu$resident_from, "host") +expect_equal(cpu$stage_from, "host") # The loader refuses block streaming it cannot do expect_error(ltx23_load_pipeline("unused", device = "cpu", resident_gb = 4), "phase_offload") expect_error(ltx23_load_pipeline("unused", device = "cpu", resident_from = "disk"), "phase_offload") +expect_error(ltx23_load_pipeline("unused", device = "cpu", + stage_from = "disk"), "phase_offload") +expect_error(load_gemma4_nf4("unused", device = "cuda", stage_from = "disk"), + "device = \"cpu\"") # The split text phase needs connectors, and an encoder or raw states no_conn <- structure(list(), class = "ltx23_pipeline") diff --git a/inst/tinytest/test_mmap_reader.R b/inst/tinytest/test_mmap_reader.R new file mode 100644 index 0000000..075384f --- /dev/null +++ b/inst/tinytest/test_mmap_reader.R @@ -0,0 +1,83 @@ +# .mmap_reader: tensors read through a private memory map of a +# safetensors file match safetensors' own reader exactly, for every +# dtype the artifacts use, including tensors whose data is not aligned +# to their element size (the R writer does not pad the header). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed() || + !requireNamespace("safetensors", quietly = TRUE)) { + exit_file("torch or safetensors not available") +} +library(torch) + +torch_manual_seed(3) +src <- list( + a_u8_odd = torch_randint(0, 255, 7L, dtype = torch_uint8()), + b_bf16 = torch_randn(3, 5)$to(dtype = torch_bfloat16()), + c_f32 = torch_randn(4, 2, 3), + d_u8 = torch_randint(0, 255, c(2L, 3L), dtype = torch_uint8()), + e_i64 = torch_tensor(c(-3, 0, 9), dtype = torch_long()), + f_scalar = torch_tensor(2.5), + g_f16 = torch_randn(5)$to(dtype = torch_float16()) +) +f1 <- tempfile(fileext = ".safetensors") +f2 <- tempfile(fileext = ".safetensors") +# Lengthen one key until the header length leaves typed data misaligned +# (the case the reader exists for) +misaligned <- function(f) { + h <- diffuseR:::.st_header(f) + el <- c(U8 = 1, BF16 = 2, F32 = 4, I64 = 8, F16 = 2) + any(vapply(h$tensors, function(i) { + (h$data_start + i$data_offsets[[1]]) %% el[[i$dtype]] != 0 + }, logical(1))) +} +for (pad in 0:7) { + names(src)[1] <- paste0("a_u8_odd", strrep("x", pad)) + safetensors::safe_save_file(src[1:4], f1) + if (misaligned(f1)) break +} +expect_true(misaligned(f1)) +safetensors::safe_save_file(src[5:7], f2) + +reader <- diffuseR:::.mmap_reader(c(f1, f2)) +if (is.null(reader)) { + exit_file("this torch cannot map files") +} +expect_equal(sort(reader$keys()), sort(names(src))) + +for (k in names(src)) { + got <- reader$read(k, "cpu") + expect_equal(got$dtype, src[[k]]$dtype, info = k) + expect_equal(as.integer(got$shape), as.integer(src[[k]]$shape), info = k) + expect_true(torch_equal(got, src[[k]]), info = k) +} +expect_error(reader$read("nope", "cpu"), "not found") + +# A write to a read tensor reaches neither the file nor later reads +x <- reader$read("d_u8", "cpu") +x$add_(1L) +fresh <- safetensors::safetensors$new(f1, framework = "torch")$get_tensor("d_u8") +expect_true(torch_equal(fresh, src$d_u8)) +expect_true(torch_equal(reader$read("d_u8", "cpu"), src$d_u8)) + +if (torch::cuda_is_available()) { + for (k in names(src)) { + got <- reader$read(k, "cuda") + expect_equal(got$device$type, "cuda", info = k) + expect_true(torch_equal(got$cpu(), src[[k]]), info = k) + } +} + +# Disk-backed staging through the reader: placeholders after offload, +# exact values after onload +live <- lapply(src[c("b_bf16", "c_f32")], function(t) t$clone()) +pairs <- diffuseR:::.stage_from_disk(live, c(b_bf16 = "b_bf16", c_f32 = "c_f32"), + list(get_tensor = function(k) stop("unused"), + read_to = reader$read), pin = FALSE) +expect_equal(live$c_f32$stride(), c(0L, 0L, 0L)) +diffuseR:::.staged_onload(pairs, "cpu") +expect_true(torch_equal(live$c_f32, src$c_f32)) +expect_true(torch_equal(live$b_bf16, src$b_bf16)) +diffuseR:::.staged_offload(pairs) +expect_equal(live$b_bf16$stride(), c(0L, 0L)) + +unlink(c(f1, f2)) diff --git a/inst/tinytest/test_recommend.R b/inst/tinytest/test_recommend.R index 9f060bf..a034092 100644 --- a/inst/tinytest/test_recommend.R +++ b/inst/tinytest/test_recommend.R @@ -306,12 +306,15 @@ expect_equal(r8$resident_from, "host") low <- recommend("ltx25", 7.6, cran, host_ram_gb = 20) expect_true(low$split_text) expect_equal(low$resident_from, "disk") +expect_equal(low$stage_from, "host") expect_true(low$pin) +expect_equal(recommend("ltx25", 7.6, cran, host_ram_gb = 10)$stage_from, "disk") cpu <- recommend("ltx25", 0, cran, host_ram_gb = 64) expect_equal(cpu$devices$transformer, "cpu") expect_false(cpu$pin) expect_false(cpu$split_text) expect_equal(cpu$resident_from, "host") +expect_equal(cpu$stage_from, "host") # Other models do not grow the LTX fields expect_false("resident_gb" %in% names(recommend("ltx", 16, cran))) expect_true("precision" %in% names(formals(download_ltx2))) diff --git a/man/load_gemma4_nf4.Rd b/man/load_gemma4_nf4.Rd index bfdc023..0670ceb 100644 --- a/man/load_gemma4_nf4.Rd +++ b/man/load_gemma4_nf4.Rd @@ -9,6 +9,7 @@ load_gemma4_nf4( dtype = "bfloat16", pin = getOption("diffuseR.pin_staging", TRUE), stream = FALSE, + stage_from = c("host", "disk"), verbose = TRUE ) } @@ -29,6 +30,11 @@ layer by layer instead of whole: an encode on the GPU then holds only the embeddings and one layer at a time, for cards that cannot hold the ~8 GB NF4 model.} +\item{stage_from}{"host" keeps the staged weights in host memory; +"disk" keeps none and reads each layer from the artifact (memory +mapped) when the encode needs it, so the encoder costs almost no +host RAM. "disk" needs \code{device = "cpu"} and a CUDA encode.} + \item{verbose}{Logical.} } \value{ diff --git a/man/load_gemma4_text_encoder.Rd b/man/load_gemma4_text_encoder.Rd index 3b320e4..745dfc2 100644 --- a/man/load_gemma4_text_encoder.Rd +++ b/man/load_gemma4_text_encoder.Rd @@ -9,6 +9,7 @@ load_gemma4_text_encoder( dtype = "bfloat16", pin = getOption("diffuseR.pin_staging", TRUE), stream = FALSE, + stage_from = c("host", "disk"), verbose = TRUE ) } @@ -31,6 +32,9 @@ encode at DMA speed. Default follows \item{stream}{Logical. Stage layer by layer (see \code{\link{load_gemma4_nf4}}).} +\item{stage_from}{"host" or "disk" (NF4 artifacts only; see +\code{\link{load_gemma4_nf4}}).} + \item{verbose}{Logical.} } \value{ diff --git a/man/ltx23_load_group.Rd b/man/ltx23_load_group.Rd index b03bb9a..5bfbdfa 100644 --- a/man/ltx23_load_group.Rd +++ b/man/ltx23_load_group.Rd @@ -9,7 +9,8 @@ ltx23_load_group( module, map_key = identity, verbose = TRUE, - gc_every = 50L + gc_every = 50L, + lazy = FALSE ) } \arguments{ @@ -26,12 +27,17 @@ parameter/buffer name, or NA to skip the key deliberately.} \item{verbose}{Logical. Report progress and coverage.} \item{gc_every}{Integer. Run \code{gc()} after this many tensors.} + +\item{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.} } \value{ 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}. } \description{ diff --git a/man/ltx23_load_pipeline.Rd b/man/ltx23_load_pipeline.Rd index e839265..83dcbea 100644 --- a/man/ltx23_load_pipeline.Rd +++ b/man/ltx23_load_pipeline.Rd @@ -14,6 +14,7 @@ ltx23_load_pipeline( phase_offload = TRUE, resident_gb = Inf, resident_from = c("host", "disk"), + stage_from = c("host", "disk"), verbose = TRUE ) } @@ -64,6 +65,19 @@ keeps none and re-reads it from the artifact at the next onload host RAM equal to the resident part (up to ~12 GB) for a re-read per render. Needs \code{phase_offload} on CUDA.} +\item{stage_from}{Where every phase-offloaded weight lives between +uses: \code{"host"} (default) in host memory; \code{"disk"} nowhere +but the artifact. Loading then reads no weights, and each onload +(the transformer's resident part per render, its streamed blocks +every step, the connectors, VAEs and vocoder per phase) reads from +the memory-mapped shards, so the pipeline holds almost no host +RAM and the OS page cache keeps what fits. Implies +\code{resident_from = "disk"}; the connectors always stream. For +machines with less host RAM than the model; each step re-reads the +streamed blocks (~1.4 GB/s cold from NVMe, ~4.5 GB/s from the page +cache). Needs an NF4 artifact directory and \code{phase_offload} +on CUDA.} + \item{verbose}{Logical.} } \value{ diff --git a/man/ltx23_load_transformer_nf4.Rd b/man/ltx23_load_transformer_nf4.Rd index cdab665..1ba9a57 100644 --- a/man/ltx23_load_transformer_nf4.Rd +++ b/man/ltx23_load_transformer_nf4.Rd @@ -3,7 +3,13 @@ \alias{ltx23_load_transformer_nf4} \title{Load the LTX-2.3 transformer with resident NF4 weights} \usage{ -ltx23_load_transformer_nf4(ckpt, device = "cuda", verbose = TRUE, ...) +ltx23_load_transformer_nf4( + ckpt, + device = "cuda", + verbose = TRUE, + lazy = FALSE, + ... +) } \arguments{ \item{ckpt}{An NF4 \code{ltx23_checkpoint} @@ -13,6 +19,11 @@ ltx23_load_transformer_nf4(ckpt, device = "cuda", verbose = TRUE, ...) \item{verbose}{Logical.} +\item{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.} + \item{...}{Passed to \code{\link{ltx23_transformer}} (tiny test configs).} } \value{ diff --git a/man/ltx23_memory_profile.Rd b/man/ltx23_memory_profile.Rd index 8aa1a51..cd6e672 100644 --- a/man/ltx23_memory_profile.Rd +++ b/man/ltx23_memory_profile.Rd @@ -14,7 +14,8 @@ when undetectable, which keeps the all-in-RAM defaults).} \value{ 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. } \description{ Selects transformer precision, component placement, and attention @@ -53,6 +54,12 @@ of the GPU-resident transformer blocks. Below 32 GB available the 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. } diff --git a/man/recommend.Rd b/man/recommend.Rd index 1923785..1d582c9 100644 --- a/man/recommend.Rd +++ b/man/recommend.Rd @@ -34,8 +34,9 @@ A list with \code{model}, \code{precision}, \code{devices} \code{host_ram_gb}, \code{fork_suggested} (logical: the installed safetensors cannot read the tier the card could otherwise run), and \code{note} (the suggestion string, or NULL). For \code{"ltx25"} it - also has \code{resident_gb}, \code{stream_text}, \code{split_text} - and \code{resident_from} (see \code{\link{ltx23_memory_profile}}). + also has \code{resident_gb}, \code{stream_text}, \code{split_text}, + \code{resident_from} and \code{stage_from} (see + \code{\link{ltx23_memory_profile}}). The \code{fork_suggested} field name is historical -- the remedy used to be the cornball-ai fork and is now \code{install.packages("safetensors")}, since the fixes reached CRAN @@ -74,11 +75,12 @@ has only the nf4 artifact: below 14 GB it keeps part of the nf4 transformer resident and streams the rest, and the result carries the \code{\link{ltx23_memory_profile}} fields that configure that (\code{resident_gb}, \code{stream_text}, \code{split_text}, -\code{resident_from}). Its staging is always pinned on the GPU -tiers: block streaming from pageable memory is slow, and the -host-RAM fields already bound the footprint (measured at 768x512x121 -on an 8 GB card: 20.1 GB high-water with resident host copies, 16.0 -GB with \code{resident_from = "disk"}). +\code{resident_from}, \code{stage_from}). Its staging is always +pinned on the GPU tiers: block streaming from pageable memory is +slow, and the host-RAM fields already bound the footprint (measured +at 768x512x121 on an 8 GB card: 20.1 GB high-water with resident +host copies, 16.0 GB with \code{resident_from = "disk"}; with +\code{stage_from = "disk"} no weights are held in host memory). The pinning decision: phase-swapped weights are page-locked host copies (see \code{\link{staging}}) that transfer at DMA rate - diff --git a/man/serve.Rd b/man/serve.Rd index 1679f57..e14b7ab 100644 --- a/man/serve.Rd +++ b/man/serve.Rd @@ -28,8 +28,10 @@ served yet. "ltx25" loads the \code{ltx2.5-nf4-full} and \code{gemma4-nf4} artifacts and places them by \code{\link{recommend}("ltx25")}: on 8-12 GB cards the transformer streams blocks and, below 12 GB, Gemma 4 streams layers. The -encoder and render pipeline share one process, which needs about -31 GB of host RAM.} +encoder and render pipeline share one process: with 32 GB of host +RAM available they keep their weights there, with less they read +them from the memory-mapped artifacts as each phase needs them +(\code{stage_from = "disk"} in \code{\link{ltx23_load_pipeline}}).} \item{device}{Character. "cuda" or "cpu".} From 9b9a2272a21dfc076de4050d59b6c41aa4026301 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 29 Oct 2026 14:23:14 -0500 Subject: [PATCH 2/2] Bump version to 0.2.2.21 --- DESCRIPTION | 2 +- NEWS.md | 24 ++++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index 54f6a84..023e031 100644 --- a/DESCRIPTION +++ b/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")), diff --git a/NEWS.md b/NEWS.md index d8f0daf..7016edb 100644 --- a/NEWS.md +++ b/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