GPU memory and runtime

What to change when a job runs out of memory, and how to find out what a run will actually cost before you commit to it.

Measure rather than guess

Requirements vary by model, by dataset size, and by GPU, and no table can substitute for one measured run. Nextflow will tell you:

scfoundry embed --method scgpt --data demo/colon_1000.h5ad \
  -- -with-trace trace.txt -with-report report.html

trace.txt, written in the run directory, records peak resident memory and wall time per task. Run one representative dataset, read the trace, then size the rest. This is worth five minutes before a sweep that would otherwise fail overnight.

For GPU memory specifically, watch the device while a task runs:

nvidia-smi --query-gpu=memory.used,memory.total --format=csv -l 5

When you run out of GPU memory

Lower the batch size. It is the first and usually the only knob you need, and which option to use depends on the task:

Task

Option

embed, transfer (the embedding step)

--batch-size

finetune, training

--batch-size

finetune, prediction

--predict_batch_size

embed --method scgpt_integrated

--integration_batch_size

Halving is a reasonable first step:

scfoundry embed --method uce --data demo/colon_1000.h5ad --batch-size 8

If a method’s default is already small — Cell2Sentence runs at 8, UCE at 16 — you are near the floor, and the answer is a larger GPU rather than a smaller batch.

Note

Some methods declare no batch_size at all: cellama, scfoundation, genept, novae and pca. Passing --batch-size to them is accepted but does nothing. Fine-tuning has its own gaps — CellPLM ignores both fine-tuning batch options, CELLama ignores --predict_batch_size.

Two methods have a second lever worth knowing:

  • scPRINT — --max_len (default 4000) caps genes per cell. Lowering it cuts memory roughly linearly, at some cost in fidelity.

  • UCE — --nlayers (default 33) selects how many transformer layers to use. Reducing it changes what the model computes, so treat it as a different configuration rather than a memory fix.

What drives cost

Three properties predict how expensive a method will be, all of them checkable rather than guessed.

Model size. Some are named in the checkpoint: Geneformer-V2-316M is 316M parameters, Cell2Sentence wraps a 410M-parameter Pythia model, CellFM is 80M, CellPLM 85M. On-disk checkpoint size is a decent proxy for the rest — Downloading model weights tabulates all of them.

Sequence length. Transformer memory scales with the number of gene tokens per cell. scPRINT’s --max_len 4000 and UCE’s long padded sequences are the reason those two carry the smallest default batch sizes.

Default batch size, read backwards. The defaults encode the authors’ and our own experience of what fits. A method defaulting to 2048 (SCimilarity) is cheap per cell; one defaulting to 8 (Cell2Sentence) is not. The full table is in the method reference.

Cheapest to most expensive, as a rough ordering: pca and genept (no neural inference over cells at all, CPU-friendly) → scvi, scimilarity → scbert, cellfm, cellplm, sccello → scgpt, geneformer, scfoundation, langcell → scprint, uce, c2s. Among the integration methods, Harmony and Seurat are CPU jobs of minutes; scvi_denovo trains a small model; scgpt_integrated fine-tunes a transformer and costs GPU hours.

Scaling with cell count

Most of the pipeline is linear in cells: inference batches through the dataset, so doubling cells roughly doubles time at constant memory.

The exceptions are in benchmark and geometry, both CPU tasks:

  • ASW computes a full pairwise distance matrix, quadratic in cells; benchmark densifies X before scoring. An embedding of 500,000 cells × 512 dimensions is about 2 GB dense before any metric runs. --batch-max-cells bounds the batch metrics, which are the most expensive part.

  • geometry subsamples to --max-cells (20,000 by default) precisely so that its pairwise and neighbourhood statistics stay tractable.

Scoring large datasets therefore wants a large-memory CPU node rather than a GPU node.

CPU and threading

A few processes accept explicit parallelism:

  • --geneformer_nproc (default 1) — tokenisation processes.

  • --sccello_tokenize_num_proc (default 1) — the same idea for scCello.

  • --scprint_num_workers (default 0) — dataloader workers.

Everything else is single-process. Under Docker the --shm-size=16g option exists precisely so that multi-worker dataloaders do not crash on shared memory.

Sharing a GPU

Nothing in the shipped configuration stops several runs launched at once from landing on the same device.

Tip

On a multi-GPU machine, pin a run to one device with --gpu 0. Without it, Docker gets --gpus=all and Apptainer gets --nv with no CUDA_VISIBLE_DEVICES, so the job may land on a card someone else is using. To make it permanent, set gpu_id in the workspace nextflow.config.

Runs are independent, so the simplest discipline is to launch them one after another from a loop; when you do run several at once, give each its own --gpu. The one task that parallelises internally is transfer, which embeds reference and query as two concurrent tasks — -- -process.maxForks 1 serialises them.

On a cluster, let the scheduler handle it instead — see Running on an HPC cluster.

Disk

Four consumers, in descending order of appetite:

Model weights — about 71 GB for the full set, dominated by SCimilarity (41 GB) and UCE (17 GB). Download only what you need, and share one copy between workspaces with --weights-dir.

Container images — roughly 20 GB across the images you actually use, in cache/.shared/nxf_singularity/. Written once.

Work directories — transient, but large while a run is in flight, since each task stages its inputs. cleanup = true removes them on success, which is why the default exists. A failed run keeps its work/; delete the run directory once you are done with it.

Results — small. Embeddings are a few hundred megabytes at most; fine-tuned models are the size of a checkpoint; metric tables are kilobytes.

See also