topic ML Training low File I/O and Throughout in Machine Learning

ML Training low File I/O and Throughout

aswinkks — Mon, 17 Aug 2026 18:55:08 GMT

Hi everyone,

I have an image-based deep learning workload running on Azure Databricks, while the training dataset must remain in AWS S3 due to some constraints. We cannot move or replicate the dataset to Azure.
Our current architecture is roughly:
AWS S3 (images) → Unity Catalog Volume → Azure Databricks GPU compute → PyTorch training

The main issue we're seeing is very high file I/O latency and relatively low training throughput. Since the dataset contains a large number of individual image files, reading the images through the UC Volume appears to involve significant network overhead and many individual file reads.
I initially expected Mosaic Streaming / StreamingDataset to improve this because the dataset can be converted into MDS shards and the shards can be downloaded progressively to local storage while training continues.

However, interestingly, in our testing:
Direct image loading from the UC Volume is currently faster than Mosaic Streaming.
I'm trying to understand whether we are missing an important configuration or whether the cross-cloud architecture itself is the primary bottleneck.

A few questions:
1. What is the recommended architecture for training on S3 data from Azure Databricks when the data cannot be moved to Azure?

2. For Mosaic Streaming, what are the recommended values/strategies for:
shard size
num_workers
predownload
cache_limit
shuffle configuration

3. Is there an optimal MDS shard size for image datasets to minimize S3/network overhead?

4. Would increasing DataLoader workers and prefetching significantly improve throughput in this cross-cloud scenario?

5. Would it be better to use local caching of UC Volume files instead of Mosaic Streaming for a multi-epoch image training workload?

6. Are there any Databricks-recommended approaches for measuring whether the bottleneck is S3 → Azure network bandwidth, file-level latency, CPU image decoding, or GPU starvation?

Our main objective is to maximize GPU utilization and training throughput without moving the source dataset out of S3.
Any recommendations, benchmarks, or reference architectures for this type of cross-cloud training setup would be greatly appreciated

Re: ML Training low File I/O and Throughout

GabFernandes — Mon, 17 Aug 2026 19:28:14 GMT

Hi @aswinkks ,

The core issue is clear: cross-cloud file-level I/O (AWS S3 → Azure compute) with many small files is the worst-case scenario for training throughput. Here's a structured breakdown addressing each of your questions.

1. Recommended Architecture for S3 Data from Azure Databricks

Your current path S3 → UC Volume (external, read-only) → FUSE mount → Training adds multiple latency layers: cross-cloud network hop (~50-150ms per request), FUSE overhead, and per-file request amplification.

Recommended approaches (best → acceptable):

Strategy	Latency	Complexity
Pre-convert to large MDS/WebDataset shards, stage to Azure ADLS, read locally	Lowest	Medium (one-time ETL
Pre-convert to MDS shards on S3, use Mosaic Streaming with aggressive local caching	Medium	Medium
UCVolumeDataset with local caching (Databricks AI Runtime)	Medium-High (1st epoch), Low (subsequent)	Lowest
Direct UC Volume reads (no caching)	Highest	Lowest

If you truly cannot replicate data to Azure, option B (MDS shards on S3 + Mosaic Streaming + local cache) or option C (UCVolumeDataset) are your best bets. The key insight is: you must amortize the cross-cloud transfer across epochs via local caching.

2 & 3. Mosaic Streaming Configuration for Cross-Cloud Image Datasets

The reason Mosaic Streaming may be slower in your tests is likely shard size too small and/or predownload too low, meaning the pipeline stalls waiting for cross-cloud downloads.

Recommended settings for cross-cloud image workloads:

import streaming # --- Writing MDS shards --- # Target 128-256 MB per shard for images (larger = fewer cross-cloud round trips) # Default is 67 MB — too small for high-latency links writer = streaming.MDSWriter( out="s3://your-bucket/mds-dataset", columns={"image": "jpeg", "label": "int"}, size_limit=256 * 1024 * 1024, # 256 MB per shard ) # --- Reading / Training --- dataset = streaming.StreamingDataset( remote="s3://your-bucket/mds-dataset", # remote source local="/local_disk0/mds-cache", # LOCAL SSD cache (critical!) shuffle=True, shuffle_block_size=262144, # 256K samples per shuffle block predownload=16, # download 16 batches ahead (increase for high latency!) cache_limit="100gb", # keep shards locally across epochs num_canonical_nodes=None, # auto batch_size=64, ) dataloader = streaming.StreamingDataLoader( dataset, batch_size=64, num_workers=8, # match to CPU cores available for I/O prefetch_factor=4, # each worker prefetches 4 batches pin_memory=True, persistent_workers=True, # avoid re-fork overhead between epochs )

4. DataLoader Workers and Prefetching Impact

Yes — significantly, but only if paired with local caching. Without caching, more workers just amplify the cross-cloud request volume, potentially hitting S3 rate limits or saturating bandwidth without reducing latency.

With local caching:

num_workers=6-8 + prefetch_factor=4 is the Databricks-recommended starting point (from the AI Runtime DataLoader defaults)
This overlaps GPU compute with data fetch/decode, hiding I/O latency
pin_memory=True + persistent_workers=True eliminate per-batch allocation and per-epoch fork overhead

If using Databricks Serverless GPU (AI Runtime 5+):

from serverless_gpu.data import UCVolumeDataset, DataLoader # Automatic local caching + optimized prefetch (num_workers=6, prefetch_factor=4 by default) path_dataset = UCVolumeDataset("/Volumes/catalog/schema/volume/images") loader = DataLoader(path_dataset, batch_size=64)

This is Databricks' purpose-built solution for exactly your scenario — it caches each file to local NVMe on first access and serves from cache for all subsequent reads.

5. Local Caching vs. Mosaic Streaming for Multi-Epoch Image Training

For multi-epoch training, local caching wins after epoch 1. Here's the tradeoff:

Approach	Epoch 1	Epoch 2+	Memory-efficient shuffle
Mosaic Streaming (MDS)	Moderate (sequential shard download)	Fast (if cache_limit retains shards)	Yes (built-in)
UCVolumeDataset + local cache	Slower (per-file cross-cloud fetch)	Very fast (local NVMe)	No (must implement)
Manual shutil.copytree upfront	Slowest start (copies everything)	Fastest (fully local)	No

Recommendation: If your dataset fits on local disk — use UCVolumeDataset (simplest, no format conversion). If it doesn't fit — use Mosaic Streaming with a generous cache_limit and LRU eviction.

The reason your tests showed UC Volume faster than Mosaic Streaming is likely because:

UC Volume FUSE has internal read-ahead/prefetch that partially hides latency for sequential access
Mosaic Streaming's shard download + decompression adds overhead that only pays off when shards are large and predownload is tuned high
If MDS shards were small (default 67 MB), you're making just as many cross-cloud requests but with added decompression cost

6. Diagnosing the Bottleneck

Run these in parallel during training to identify the chokepoint:

# === GPU Utilization (should be >90% if data pipeline is healthy) === # In a separate cell or terminal: # !nvidia-smi dmon -s u -d 2 # GPU util every 2 sec # === DataLoader timing (insert around your training loop) === import time for epoch in range(num_epochs): data_time_total = 0 compute_time_total = 0 t0 = time.perf_counter() for batch in dataloader: data_time = time.perf_counter() - t0 data_time_total += data_time # --- GPU forward/backward --- t1 = time.perf_counter() loss = model(batch) loss.backward() optimizer.step() compute_time = time.perf_counter() - t1 compute_time_total += compute_time t0 = time.perf_counter() print(f"Epoch {epoch}: data_load={data_time_total:.1f}s, " f"compute={compute_time_total:.1f}s, " f"ratio={data_time_total/(data_time_total+compute_time_total)*100:.1f}% Waiting on data")

Summary: Recommended Action Plan

Quick win: Switch to UCVolumeDataset + serverless_gpu.data.DataLoader (if on Serverless GPU) or manually pre-copy a subset to /local_disk0/ before training
Medium-term: Convert images to MDS with size_limit=256MB, tune predownload=16, set cache_limit to fill local SSD
Long-term: If budget allows, stage the MDS shards to an Azure ADLS location (even if the raw images must stay in S3, derived MDS shards might be allowed?) — this eliminates the cross-cloud hop entirely

The fundamental law here: with cross-cloud access, you're paying ~100ms per request vs ~1ms for local storage. The only real fix is to minimize the number of remote requests (large shards) and cache aggressively (avoid re-fetching across epochs).

If my answer was helpful, please consider marking it as accepted solution!

Re: ML Training low File I/O and Throughout

aswinkks — Tue, 18 Aug 2026 06:53:45 GMT

The best and the most detailed explanation on Mosaic Streaming. Thank you

Re: ML Training low File I/O and Throughout

snehamore811 — Tue, 18 Aug 2026 07:31:16 GMT

Yes according to my understanding the cross cloud architecture is the creating the issue,but there is one more noticeable thing Mosaic Streaming is not automatically faster than UC Volume access.

In your particular workload, I would actually test

UC Volume → local NVMe cache → PyTorch as the primary architecture before investing further in Mosaic Streaming.

Databricks' current Azure guidance is quite aligned with that. For unstructured data such as images, Databricks recommends UCVolumeDataset, which copies files from the UC Volume to local storage on first access and serves subsequent epochs from the local cache. Databricks explicitly notes that /Volumes access is network-bandwidth limited and recommends local caching for multi-epoch training.

Re: ML Training low File I/O and Throughout

ThiamLee — Wed, 19 Aug 2026 08:37:40 GMT

I’d first check whether the bottleneck is S3/network latency or image decoding. For multi-epoch training, local caching + larger MDS shards might help, but the cross-cloud setup could still be the main issue.