From 04b1813fdc4959b356258c430610d4f743988c8a Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Tue, 8 Sep 2026 17:56:40 -0500 Subject: [PATCH] Document TorchScript and allocator economics in the performance vignette Move the LTX-2.3 tuning notes out of the personal skill hub and into "Performance Levers": compiling the hot step, when tracing loses, allocator cost models and the once-at-load callback gates, pool regrowth between phases, compact per-token conditioning, and verifying tuning inputs. ltx23_tune_gc() now cites the vignettes instead of a path on Troy's machine. --- R/memory_ltx23.R | 4 +- man/ltx23_tune_gc.Rd | 4 +- vignettes/performance-levers.md | 132 ++++++++++++++++++++++++++++++++ 3 files changed, 136 insertions(+), 4 deletions(-) diff --git a/R/memory_ltx23.R b/R/memory_ltx23.R index 476ef65..ce74ee6 100644 --- a/R/memory_ltx23.R +++ b/R/memory_ltx23.R @@ -98,8 +98,8 @@ ltx23_memory_profile <- function(vram_gb = NULL) { #' Tune the torch CUDA allocator for large-resident inference #' -#' Stops the allocator GC storm (cf. ~/skills/torch -#' torch-jit-gc-performance.md): lantern proactively calls R's gc() +#' Stops the allocator GC storm (see the "Performance Levers" vignette +#' and chatterbox's "performance" vignette): lantern proactively calls R's gc() #' whenever reserved memory exceeds \code{torch.cuda_allocator_reserved_rate} #' (default 0.20) of the card. With ~75\% of VRAM occupied by resident #' weights that fires on nearly every allocation. Raising the rate to the diff --git a/man/ltx23_tune_gc.Rd b/man/ltx23_tune_gc.Rd index 9aa575b..e9c0309 100644 --- a/man/ltx23_tune_gc.Rd +++ b/man/ltx23_tune_gc.Rd @@ -15,8 +15,8 @@ ltx23_tune_gc(footprint_gb = 12, total_gb = NULL) Invisibly, the applied reserved rate (NULL if skipped). } \description{ -Stops the allocator GC storm (cf. ~/skills/torch -torch-jit-gc-performance.md): lantern proactively calls R's gc() +Stops the allocator GC storm (see the "Performance Levers" vignette +and chatterbox's "performance" vignette): lantern proactively calls R's gc() whenever reserved memory exceeds \code{torch.cuda_allocator_reserved_rate} (default 0.20) of the card. With ~75\% of VRAM occupied by resident weights that fires on nearly every allocation. Raising the rate to the diff --git a/vignettes/performance-levers.md b/vignettes/performance-levers.md index a9eac35..b727680 100644 --- a/vignettes/performance-levers.md +++ b/vignettes/performance-levers.md @@ -129,3 +129,135 @@ Treat the result as the machine's advice: pass its fields to the loaders and generators. The FLUX-family loaders' `pin` (and FLUX.1's `text_device`), `flux_memory_profile()`, and `serve()`'s LTX artifact selection consume it automatically. + +## Where the time goes: TorchScript and the allocator + +The levers above decide what fits. A second set decides how fast the +fitted configuration runs. They came out of making the 22B NF4 LTX +DiT fast on a 16 GB card, where a same-seed render went from 1331 s +to 221 s without changing a weight. + +### Compile the hot step + +R torch pays twice per eager op: the dispatch itself and the R handle +each intermediate becomes, which the garbage collector then has to +sweep. `jit_compile()` removes both. The LTX denoiser step compiles +into one TorchScript call: a `while` loop over the 48 blocks inside +the script, weights passed as a flat `List[Tensor]` with a fixed +per-block slot layout, and the NF4 dequant math in-script. Measured on +the same step: eager 99.9 s at 88% of wall time in R's gc, compiled +8.8 s at 1%. + +Builtins lantern's TorchScript accepts beyond the documented list: +`torch.scaled_dot_product_attention` (fused, so materialized +`[B,H,S,S]` score matrices and their chunked-attention workarounds go +away), `torch.gelu(x, approximate="tanh")`, `torch.sigmoid`, +`torch.silu`, `torch.rsqrt`, `torch.linear`, +`torch.bitwise_right_shift`, `torch.bitwise_and`, `torch.index_select`, +and the `.long()` method cast. That set is enough for NF4 dequant +fully in-script; chunk the loop to bound the int64 index temporaries. + +Marshalling traps, all silent or opaque: + +- A named R list marshals as `Dict[str, Tensor]`, not `List[Tensor]`. + `lapply()` over an `nn_module_list` yields named children, so + `unname()` the packed weight list. Parity-test the packer against + the script's slot indices. +- R integers marshal as `List[int]`; wrap scalars in `jit_scalar()`. +- Dtype constants such as `torch.long` do not resolve; use method + casts and `type_as()`. +- `Optional[Tensor]` parameters accept R `NULL`. Cross-def calls + within one compilation unit work. Tuple returns come back as R + lists. +- Compile once per session and cache the compilation unit. Passing + 5000+ tensor references per call costs milliseconds. + +### When tracing loses + +`jit_trace()` converts static feed-forward modules (VAE decoders, +vocoders) wholesale with exact parity, but it bakes runtime shapes +into the graph, so cache one trace per input shape. Trace a closure +rather than a bare `nn_module`, and set `requires_grad_(FALSE)` on the +parameters first. Two hazards decide whether it pays: + +- An R gc during recording corrupts captured argument values. Under + memory pressure the allocator callback can fire mid-trace and the + graph records garbage narrow offsets. Run `gc()` and + `cuda_empty_cache()` immediately before tracing, wrap the trace in + `tryCatch()`, and validate every fresh trace against the eager + output once, falling back to eager for that shape on mismatch. +- A trace captures the module's weights and pins them on the device + across phase offloads, so it must be released when the component + leaves the GPU. A once-per-render decode then re-pays trace plus + validation every run: traced was net slower than eager (39.8 s vs + 32.6 s). `jit_compile()` from source is shape-generic and captures + nothing, so prefer it wherever you can afford to write the script. + +### Allocator cost models + +`PYTORCH_CUDA_ALLOC_CONF` picks between two cost models, and the right +one depends on whether intermediates still exist as R handles: + +- `expandable_segments:True`: cheap growth, expensive frees (page + unmaps). Suits eager paths with fragmentation churn. +- `backend:native`: expensive growth (about 15 ms per `cudaMalloc`), + cheap frees. Suits a JIT path that keeps intermediates out of R. + +The same decode measured 32.7 s at 86% gc on expandable segments and +21.0 s at 50% on native. The setting is process-wide and read before +the first CUDA allocation, so set it per model at load time. + +When a phase is gc-bound, ask where the gc time goes before touching +options. If the trigger frequency is the problem, the callback gates +help: `torch.cuda_allocator_reserved_rate` plus the two most people +forget, `torch.cuda_allocator_allocated_rate` and +`torch.cuda_allocator_allocated_reserved_rate` (both default 0.8). +Raising those two to 0.95 took the expandable-segments decode from +32.7 s to 21.5 s. If the freeing work itself is the problem, no +option helps; the fix is fewer or cheaper frees, through the backend +choice or through JIT so the handles never exist. + +torch reads all three gates once, in its `.onLoad`. Any package that +imports torch has it loaded before its own tuner runs, and the set +silently no-ops while `getOption()` still returns the value you set. +`ltx23_tune_gc()` therefore pushes the gates into the live allocator +as well as setting the options. + +### Pool regrowth between phases + +`cuda_empty_cache()` between pipeline phases returns every block to +the driver, so the next phase regrows the pool one `cudaMalloc` per +tensor. Measured: 5530 tensors, 83.5 s for one 11 GB component +onload, and back to 83 s after a single empty-cache. Two fixes: + +- Pre-warm once by allocating one footprint-sized tensor and freeing + it into the cache (0.2 s, one `cudaMalloc`); later per-tensor + allocations carve from it (83.5 s to 4.6 s). +- Between phases, `gc()` without `cuda_empty_cache()`. The caching + allocator reuses freed blocks across phases on its own. + +The corollary is that "phase offload traffic" may not be transfer +time at all. Pinned staging is worth having and roughly doubles +transfer rate, but the transfers here were about 1 s all along; the +110 s was regrowth. Time the transfer in isolation before optimizing +it. + +### Compact per-token conditioning + +Prefix conditioning gives conditioned tokens timestep 0 and free +tokens timestep t, a per-token timestep vector with only a few +distinct values. Materializing per-token modulation as +`[B,S,num_params,D]` cost about 2 GB at 13.5k tokens and ran out of +memory. Passing the distinct timesteps plus an integer index and doing +`index_select` per modulation vector inside the script materializes +one `[B,S,D]` slice at a time (about 110 MB). Whenever a per-token +tensor is an index into few variants, ship the variants and the index +and select late. + +### Verify the inputs to tuning formulas + +A VRAM-detection helper once returned a hardcoded 8 GB guess when its +optional dependency was missing, and every allocator rate computed +from it was wrong for weeks. If a formula tunes by footprint over +total, log both numbers at tuning time and fail loudly on a fallback +guess. `nvidia-smi` is always there to ask.