ML Benchmarks: SLAF vs State-of-the-Art Dataloaders¶

SLAF provides state-of-the-art (SOTA) performance in data loading throughput for machine learning workflows, reaching 2.6x speedups relative to current SOTA, particularly for training transformer-based single-cell foundation models. What follows are comprehensive benchmarks comparing SLAF against state-of-the-art dataloaders including scDataset, BioNeMo SCDL, AnnDataLoader, AnnLoader, and TileDB DataLoader.

Motivation¶

The goal of these benchmarks is to demonstrate that SLAF can stream tokens to modern GPUs at a rate sufficient to prevent idle time between training loops. For a 1B parameter model like scGPT, fast enough means delivering training batches within 50 ms to keep the GPU utilization high. This benchmark establishes SLAF's ability to meet the throughput requirements for efficient foundation model training on massive single-cell datasets.

Dataset and Hardware¶

Dataset: Tahoe-100M¶

We downloaded one of the 7 h5ad files comprising the Tahoe-100M dataset made accessible by ARC Institute. This slice of the dataset contains 5,481,420 cells and 62,710 genes, with approximately 8B non-zero expression values. All benchmarks reported below used this dataset except for the TileDB dataloader since we couldn't successfully convert a 5M-cell dataset to the Tile DB SOMA Experiment format with 32G RAM. For the TileDB DataLoader alone, we report numbers on a smaller 50k-cell synthetic dataset.

Conversion and Optimization¶

We used the SLAF converter (see Migrating to SLAF) to convert the h5ad file to SLAF format. The Lance table fragments (Lance's term for partitions) were optimized for compression/query tradeoffs, with 5-10M non-zeros (rows) per fragment in the expression table. While inherently parallelizable, conversion is currently single process, and took about 10 minutes for this dataset.

Hardware Configuration¶

• Machine: Apple MacBook Pro with M1 Max
• Memory: 32 GB RAM
• Storage: 1 TB NVMe SSD
• OS: macOS 13.6.1

Internal Benchmarks¶

Methodology¶

We used a batch size of 32 with an enhanced warmup and measurement procedure to ensure accurate and consistent results, especially for the Mixture of Scanners (MoS) strategy:

• Initial Warmup: 15 batches to initialize the dataloader
• Extended Warmup: 10 seconds to allow MoS to fully stabilize all fragment generators
• Measurement Period: 40 seconds of pure performance measurement (excluding warmup time)
• Total Runtime: 50 seconds per benchmark (10s warmup + 40s measurement)

This methodology ensures that all dataloader strategies reach steady-state performance before measurement begins, eliminating variance from incomplete initialization.

Tokenization Strategy Comparison¶

We benchmarked different tokenization strategies to understand the performance impact of various preprocessing options:

Raw Mode Performance Scaling¶

Raw mode bypasses tokenization and returns Polars DataFrames that have the exact schema as sparse CSR tensors, demonstrating SLAF's base data loading performance.

Fragment vs Batch Loading Comparison¶

SLAF supports two loading strategies: fragment-based and batch-based loading. Fragment-based loading processes entire Lance fragments at once, while batch-based loading processes multiple Lance batches sequentially.

Tokenized Mode: Tokens/sec Scaling¶

Tokenized mode provides pre-tokenized sequences ready for GPU training, demonstrating SLAF's end-to-end pipeline performance.

Entropy Measurement: Training Batch Randomness¶

To ensure models don't converge to local minima due to biased and highly correlated training batches, we want to make training batches as random as possible. However, random reads are more expensive than sequential reads, so we need to balance randomness with performance.

To address this challenge, we developed a novel dataloader strategy called the Mixture of Scanners (MoS) approach, which randomly tasks a small randomized group of scanners to populate a queue of training batches by reading from different starting points of the dataset. A deeper dive into our approach to optimize dataloaders is available here and a more detailed write up of the MoS dataloader is here.

To measure entropy without using metadata, we simulate random cell IDs and measure L1 distance between pairs of cell IDs both within and across adjacent training batches for our different dataloaders to show how each dataloader strategy performs relative to a purely sequential (lowerbound) vs a truly random approach (upperbound).

We ran a test on 10,000 batches with a batch_size of 32 from a 5.4M cell dataset and found these results:

Entropy Measurement Results:

Normalized Entropy Scores [0=Sequential, 1=Random]:

Throughput Performance Results:

External Benchmarks¶

Alternate Dataloaders¶

We compared SLAF against five state-of-the-art dataloaders:

1. AnnLoader - Experimental PyTorch DataLoader for AnnData objects from anndata.experimental
2. AnnDataLoader - From scvi-tools, designed for training variational autoencoder (VAE)-style models
3. scDataset - Recently released high-performance dataloader with multiprocessing support
4. TileDB DataLoader - An internal custom PyTorch DataLoader for TileDB SOMA experiments
5. BioNeMo SCDL - NVIDIA's single-cell data loading framework for scalable training of foundation models

Methodology¶

To match the benchmarks from the scDataset paper as closely as possible, we used a batch_size=64 across all comparisons. For scDataset itself, we used the optimal parameters in our hardware (block_size=8, fetch_factor=64, which were different from the ones found to be optimal in the paper). However, we couldn't use num_workers=12 out of the box because h5ad datasets aren't pickle-able and PyTorch DataLoaders expect this since they use multiprocessing.

Enhanced Measurement Procedure: All external benchmarks now use the same enhanced measurement procedure as internal benchmarks for fair comparison:

• Initial Warmup: 15 batches to initialize each dataloader
• Extended Warmup: 10 seconds to allow all systems to reach steady state
• Measurement Period: 30 seconds of pure performance measurement (excluding warmup time)
• Total Runtime: 40 seconds per benchmark (10s warmup + 30s measurement)

This ensures fair and consistent performance comparisons across all dataloader systems.

Tier 1: Raw Data Loading Comparison¶

Raw data loading performance measures the base throughput of each system without any tokenization overhead.

Tier 2: GPU-Ready Output Comparison¶

Raw data loading benchmarks are great, provided that we intend to train on gene expression counts directly. However, for modern foundation models like Geneformer, scGPT, Transcriptformer, or STATE, cell sentences are constructed using tokens that represent gene identity and expression bins. A lot of these workflows require dataframe-friendly operations like sorting, windowing, ranking, and filtering. Our view is that it much better to situate these computations within the (typically) CPU-bound dataloader, rather than expect the GPU in the training loop to do the heavy lifting. Accordingly, SLAF dataloaders take a tokenizer and transform raw data into training-ready token sequences.

This GPU-ready throughput measures end-to-end performance including tokenization (that includes windowing, ranking, vocabulary mapping and padding), which is critical for training workflows involving models that turn cells into sentences.

Even though SLAF's tokenizing dataloaders do more work (tokenization), we find that their throughput remains competitive with scDataset's raw-data dataloader, achieving comparable performance despite the additional processing overhead.

Some Discrepancies¶

Conclusion¶

Cloud-Native Architecture¶

While these benchmarks use local SSD, the Lance format is native to cloud storage. Early tests suggest that latency between S3 and EC2 in the same region is not appreciably different from local storage. This opens up a cloud-native, store-once, query-multiple-times zero-copy architecture that eliminates data duplication.

GPU Throughput Requirements¶

What's a good enough cells/sec rate to keep an 8×H100 node at $2/hr busy? Assuming 50 ms per training loop for a model like scGPT:

• 8 GPUs × 32 cells/batch × 20 batches/sec = 5,120 cells/sec would maximize GPU utilization
• Tahoe-100M training: 100M cells ÷ 5,120 cells/sec = ~5.4 hours per epoch ~ $86 / epoch
• Anything faster than 5,120 cells/sec opens up multi-node training possibilities, trading off cost and wall clock time.

This raises the question: can we build towards a $100 scGPT model through efficient multi-node training enabled by high-throughput data loading? More on this soon!

High Concurrency and Multi-User Training¶

The Lance format's high concurrency, optimized for production multimodal data lakes with high QPS, enables not only multi-node training but multiple users training multiple models simultaneously without their own copies of the dataset. This contrasts with h5ad, which requires:

1. Local storage: The dataset must be local to the CPU instance loading it for attached GPUs
2. Non-concurrent access: One copy of the dataset per user

SLAF, with Lance under the hood, enables a truly scalable architecture for foundation model training on massive single-cell datasets.

These benchmarks demonstrate SLAF's position as the leading solution for high-performance single-cell data loading, enabling efficient training of foundation models on massive datasets with minimal resource requirements and maximum scalability.

SLAF is a young project with a bus factor of 1. You can help improve that by using it and contributing to it. Read about the SLAF vision in this blog post and contribute at github.com/slaf-project/slaf.