ludwig-ai--ludwig/.plan-lazy-preprocessing.md

Ludwig Lazy / Streaming Preprocessing Plan

Problem Statement

Ludwig's preprocessing pipeline is not lazy: before the training loop starts, it decodes and transforms every sample in the dataset into in-memory NumPy arrays. This causes:

1. OOM on media datasets — 1 000 images at original resolution × their tensor size saturates RAM before a single batch reaches the GPU.
2. Fixed preprocessing bottleneck — transforms run single-threaded on the CPU ahead of training, blocking the GPU.
3. No online statistics — normalization stats (mean/std) require a full pass that materialises the entire dataset.
4. Impossible streaming — datasets that exceed RAM (or that arrive over a network) cannot be used at all today.

---

How Other Frameworks Solve This

PyTorch / torchvision

Dataset.__getitem__ is the canonical lazy primitive. Each DataLoader worker calls it independently, so only batch_size × num_workers × prefetch_factor samples are alive at once. Transforms (resize, normalize, augment) are composed in the Dataset constructor and run inside worker processes, overlapping with the GPU forward pass.

Statistics (ImageNet mean/std) are pre-computed offline and shipped as constants. For custom datasets a single Welford accumulation scan (O(1) memory) is standard.

TensorFlow tf.data

Declarative lazy pipeline: .map(fn, num_parallel_calls=AUTOTUNE) → .batch() → .prefetch(AUTOTUNE). Nothing is executed until the training iterator pulls a batch. .cache() optionally persists the decoded dataset to disk after the first epoch.

HuggingFace datasets

Arrow IPC memory-map: only the pages actually accessed are loaded from disk. Audio/image columns store raw bytes or file paths — decoded waveforms are produced per-sample at access time via cast_column("audio", Audio()). Passing decode=False gives raw bytes with zero decoding cost.

MONAI (3D medical imaging)

Lazy resampling accumulates sequential spatial transforms (rotate → crop → resize) as pending homogeneous-matrix operations and fuses them into a single interpolation, eliminating intermediate materialisation and reducing quality loss.

torchaudio

Transforms (MelSpectrogram, Resample, MFCC) are nn.Module subclasses — composable, JIT-scriptable, and GPU-movable. Applying them post-batch on GPU is 10-20× faster than per-sample on CPU.

---

Design Principles for Ludwig's New Pipeline

1. Store paths, not arrays. The Arrow dataset (or in-memory DataFrame) holds file paths / metadata, never decoded tensors. Decode happens inside the DataLoader worker.
2. Transforms as nn.Module. All feature-specific transforms become PyTorch modules: composable, testable, batchable, GPU-ready.
3. Online statistics with Welford's algorithm. Replace full-dataset decode-then-accumulate with a single streaming pass that needs O(1) memory per feature.
4. Prefetching and parallelism. DataLoader num_workers > 0 + prefetch_factor provides CPU/GPU overlap for free once decode is inside __getitem__.
5. Backward compatibility. Existing configs that work today must continue to work. New behaviour is opt-in at the feature level via lazy: true or automatic for audio/image features.

---

Implementation Phases

---

Phase 0 — Foundations (no user-facing change)

Goal: establish shared infrastructure that later phases build on.

0-A. Welford online stats accumulator

• Add ludwig/data/statistics.py with WelfordAccumulator (supports merge across shards).
• API: acc.update(x: np.ndarray), acc.result() → (mean, std, min, max, count), acc.merge(other).
• Unit-test with known distributions.

0-B. Feature transform protocol

• Define FeatureTransform(nn.Module) base class in ludwig/features/transforms.py.
• Existing preprocessing logic stays unchanged for now; this just establishes the interface.
• __call__(self, x: Tensor) → Tensor; serialisable via torch.save.

0-C. DataLoader worker-safe file handle pool

• Research whether Ludwig's existing multiprocessing setup needs worker_init_fn to reset open file handles or RNG state (same issue as HuggingFace's IterableDataset shard split).
• Document the policy; no code change yet.

Exit criteria: WelfordAccumulator tests pass; FeatureTransform ABC exists.

---

Phase 1 — Lazy Audio (highest OOM risk)

Goal: audio columns decode inside the DataLoader worker, not during preprocessing.

1-A. AudioFeature path-only mode

• After the preprocessing step, the Arrow/DataFrame column holds file paths (strings), not decoded tensors.
• Remove _process_in_memory eager decode; replace with a lightweight path-validity check.
• The existing read_audio_files_to_dict path (used today) becomes a fallback for non-lazy mode.

1-B. AudioDataset.__getitem__ decode

• Create ludwig/data/datasets/audio_dataset.py (or extend LudwigDataset).
• In __getitem__, call torchaudio.load(path) → resample → pad/trim to target length.
• Apply FeatureTransform chain (mel spectrogram, normalization).

1-C. Online stats for audio normalization

• Replace current two-pass (decode all → compute stats) with a streaming Welford scan: torchaudio.load one file at a time, update accumulator, discard tensor.
• Triggered once at prepare_data time; stats stored in the feature config cache.

1-D. Integration test

• Run pytest tests/integration_tests/test_audio_feature.py — no OOM on the medium-size audio fixtures.
• Add a test that processes 10 000 short clips and checks RSS stays under 2 GB.

Exit criteria: audio smoke tests pass on datasets with >10 h of audio; RSS bounded by batch_size × clip_length × sr × 4 bytes.

---

Phase 2 — Lazy Image

Goal: image columns store paths; decode + augment happens per-sample in DataLoader workers.

2-A. ImageFeature path-only mode

• Preprocessing stores file paths (already partially the case for CSV/HDF5 workflows).
• Remove upfront _process_in_memory full-batch decode; keep only metadata inference (height, width, channels) via a single sampled read.

2-B. ImageDataset.__getitem__ decode

• PIL.Image.open(path).convert("RGB") → torchvision.transforms pipeline.
• Apply FeatureTransform (resize, center crop, normalize).
• Multi-channel and TIFF support via existing read_image_from_path helper.

2-C. Online stats for image normalization

• Welford accumulator over pixel values per channel; one forward pass, batch-level updates (read image → reshape to (H×W, C) → acc.update).
• Offer use_imagenet_stats: true shortcut (ships precomputed constants) for the common case.

2-D. Augmentation support

• ImageFeature config gains an optional augmentation: block (equivalent to torchvision.transforms.v2).
• Augmentation runs inside __getitem__ → free via DataLoader parallelism.

Exit criteria: image smoke tests pass; openfake and synthia complete without OOM at original resolution.

---

Phase 3 — Lazy Text / Tabular (lower priority)

Goal: tokenisation and numerical transforms move into the DataLoader worker.

3-A. Text tokenisation in __getitem__

• Currently: tokenise all text upfront into integer id arrays → store in HDF5.
• New: store raw strings; tokenise on-the-fly in __getitem__.
• Tradeoff: tokenisation is cheap but repeated; cache tokenised output to Arrow after first epoch if cache_encoder_embeddings: true.

3-B. Numerical normalisation as nn.Module

• NormalizationTransform(mean, std) replaces the current in-place column mutation.
• Stats computed once via Welford in prepare_data; transform applied in __getitem__.

3-C. Category embedding stays eager

• Category integer encoding is O(1) per sample and zero-memory; no change needed.

Exit criteria: existing tabular test suite passes unchanged.

---

Phase 4 — Streaming Dataset Source

Goal: support HuggingFace streaming datasets (streaming=True) as a first-class source, enabling datasets that exceed local disk space.

4-A. HFStreamingDataset(IterableDataset)

• Wraps a datasets.IterableDataset in a PyTorch IterableDataset.
• Shard splitting: worker_init_fn slices the stream across DataLoader workers (HF provides dataset.shard(num_shards=N, index=i)).
• Shuffle buffer: configurable per-dataset (small for media, large for text).

4-B. Online stats for streaming sources

• Welford accumulator updated during the stats-scan pre-training pass (a full iteration over the stream without gradient accumulation).
• Stats cached to disk keyed by dataset id + revision + split.

4-C. No-shuffle streaming mode

• For datasets too large even for a shuffle buffer, expose shuffle: false (warn user).

Exit criteria: AudioSet, VoxPopuli, LibriSpeech stream through without touching disk; stats computed in one pass.

---

Phase 5 — MONAI-style Transform Fusion (stretch goal)

Goal: fuse adjacent spatial transforms (resize → crop → normalize) into a single tensor operation to eliminate intermediate materialisation.

5-A. Transform graph analysis

• At setup time, walk the FeatureTransform chain; identify composable sequences.
• A sequence is composable when all transforms implement as_matrix() → Tensor[4×4].

5-B. Single-pass spatial transform

• Fuse the homogeneous matrices; apply via F.grid_sample once.
• Benchmark vs. sequential PIL transforms on ImageNet validation set.

5-C. Augmentation randomness

• Random transforms (flip, rotation jitter) break deterministic fusion; keep them as post-fusion nn.Module steps.

Exit criteria: ≥ 15% wall-clock speedup on image-only training runs vs. Phase 2 baseline.

---

Data-Flow Diagram (target state)

Disk / HF Hub
    │
    ▼
┌──────────────────────────────────────────────────────┐
│  prepare_data  (single process, once per dataset)    │
│                                                      │
│  • Build Arrow metadata table (paths, labels, ids)   │
│  • Welford scan → normalization stats                │
│  • Write stats to feature config cache              │
└───────────────────────────┬──────────────────────────┘
                            │ Arrow file / paths
                            ▼
┌──────────────────────────────────────────────────────┐
│  LudwigDataset.__getitem__  (per DataLoader worker)  │
│                                                      │
│  • Read path from Arrow row                          │
│  • torchaudio.load / PIL.Image.open                  │
│  • FeatureTransform chain (resize, mel spec, norm)   │
│  • Return sample dict of tensors                     │
└───────────────────────────┬──────────────────────────┘
                            │ batch of tensors (pinned RAM)
                            ▼
┌──────────────────────────────────────────────────────┐
│  GPU training loop                                   │
│  • H2D transfer (pin_memory → non-blocking)          │
│  • Forward pass                                      │
│  • Gradient accumulation                             │
└──────────────────────────────────────────────────────┘

---

Key Files to Create / Modify

File	Action
ludwig/data/statistics.py	NEW — WelfordAccumulator
ludwig/features/transforms.py	NEW — FeatureTransform base + concrete transforms
ludwig/features/audio_feature.py	MODIFY — remove eager decode, add lazy path
ludwig/features/image_feature.py	MODIFY — remove eager decode, add lazy path
ludwig/data/datasets/base_dataset.py	MODIFY — add __getitem__ decode hooks
ludwig/data/datasets/audio_dataset.py	NEW (or extend base)
ludwig/data/datasets/image_dataset.py	NEW (or extend base)
ludwig/data/datasets/hf_streaming_dataset.py	NEW — Phase 4
tests/unit/data/test_statistics.py	NEW — Welford unit tests
tests/integration_tests/test_lazy_audio.py	NEW
tests/integration_tests/test_lazy_image.py	NEW

---

Risks and Mitigations

Risk	Mitigation
DataLoader worker forking copies open file handles / CUDA contexts	worker_init_fn resets handles; delay CUDA init until after fork
Welford online stats diverge from current batch stats	Unit-test Welford against numpy batch stats on synthetic data
Per-sample disk I/O is slower than batched HDF5 read	Benchmark; add optional post-first-epoch disk cache for decoded tensors
Existing HDF5 preprocessing cache invalidated	Introduce cache version key; fall back to eager mode on cache miss
IterableDataset with shuffle buffer OOMs on large-media streams	Default shuffle buffer scales with declared feature size; audio → 500, image → 200, text → 10 000

---

Open Questions

1. HDF5 cache vs. Arrow: Should the decoded-tensor cache (post-Phase 2) write Arrow or HDF5? Arrow supports mmap natively; HDF5 has chunked access. Arrow preferred for new code.
2. Normalization stat freshness: If the dataset changes between runs, cached Welford stats are stale. Hash the dataset id + revision + split into the cache key.
3. Multi-output features: When one column feeds multiple output features (e.g. image → classification + detection), the decode must happen once. Ensure LudwigDataset returns a shared tensor, not independent copies.
4. Tokenisation caching: Tokenising every sample every epoch wastes CPU. Cache to Arrow after the first epoch under ${cache_dir}/<dataset_hash>/tokenized/.