Izzet Yildirim
Notes

Capturing and analyzing fine-grained deep learning I/O

·14 min read

A training job’s I/O is spread over places most tools don’t look: worker processes that come and go every epoch, Python code that decodes samples, and a GPU that runs asynchronously from all of it. The clearest way I know to show how to capture and read it is to do it once, end to end, on a real run, so that’s what this note is. The question running through it is the one that matters most for training: not how long loading takes, but how much of it the GPU actually waits for.

The workload is a small 3D UNet trained on CT volumes shaped like KiTS19,1 the dataset behind MLPerf’s UNet3D benchmark.2 I traced it with DFTracer and analyzed it with DFAnalyzer at four data-loader settings: 0, 1, 2, and 4 workers. Everything ran on my laptop, with an RTX 5060 Laptop GPU and an 8-core CPU, under WSL2 with Ubuntu 24.04. I ran each setting three times and dropped the page cache before every epoch.

The workload

  • Data: 160 volumes of 160³ voxels, each a float32 image plus a label mask, stored as compressed .npz files: 14.6 MiB each, 2.3 GiB in total. They’re synthetic, so anyone can regenerate them.
  • Per sample: np.load reads and decompresses the whole volume, then the sample is cropped to a random 96³ patch and flipped at random.
  • Training: batch size 2, two epochs (160 steps), bf16 autocast, AdamW. The loop reads the loss every step to log it, as most training loops do. That detail matters later.

Install

pip install torch --index-url https://download.pytorch.org/whl/cu128
pip install dftracer
pip install git+https://github.com/llnl-asr/dfanalyzer.git

Annotate the loop

DFTracer intercepts POSIX calls on its own. The annotations add the layers above them, and the ai module gives them standard names, so DFAnalyzer’s ai preset can find them without any per-job configuration. Here is the training script with the model and arguments left out:

from dftracer.python import ai, dftracer

@ai.data.preprocess.derive(name="augment")
def augment(img, lbl, patch, rng):
    ...  # random 96³ crop and flips

class Volumes(Dataset):
    @ai.data.item
    def __getitem__(self, i):
        with np.load(self.files[i]) as z:  # read and decompress the whole volume
            img, lbl = z["image"], z["label"]
        img, lbl = augment(img, lbl, self.patch, np.random.default_rng(i))
        return torch.from_numpy(img), torch.from_numpy(lbl.astype(np.int64))

@ai.device.transfer
def to_gpu(x, y):
    return x.cuda(non_blocking=True), y.cuda(non_blocking=True)

@ai.compute.forward
def forward(model, x, y, loss_fn):
    with torch.autocast("cuda", dtype=torch.bfloat16):
        return loss_fn(model(x), y.squeeze(1))

@ai.compute.backward
def backward(loss, opt):
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return loss.item()  # logging the loss waits for the GPU

@ai.checkpoint.capture
def save(model, path):
    torch.save(model.state_dict(), path)

The loop itself:

# paths come from DFTRACER_* variables
log = dftracer.initialize_log(None, None, -1)

@ai.pipeline.train
def train():
    for epoch in range(epochs):
        ai.pipeline.epoch.start()
        for step, (x, y) in ai.dataloader.fetch.iter(enumerate(loader)):
            ai.update(step=step, epoch=epoch)
            x, y = to_gpu(x, y)
            backward(forward(model, x, y, loss_fn), opt)
        save(model, f"{ckpt}/epoch{epoch}.pt")
        ai.pipeline.epoch.stop()

train()
log.finalize()

ai.dataloader.fetch.iter wraps the loader and records two spans per step. fetch.iter is the time spent waiting for the next batch. fetch.block is the time spent in the loop body, with the batch in hand.

Run it

DFTRACER_ENABLE=1 \
DFTRACER_LOG_FILE=traces/w4/trace \
DFTRACER_DATA_DIR=data:ckpt \
  python train.py --workers 4

DFTRACER_DATA_DIR lists the paths whose POSIX calls get traced, here the dataset and the checkpoints. You can set it to all, but then you also trace every Python import and every CUDA library load. DFTRACER_LOG_FILE is a prefix, and each process writes its own compressed file next to it: trace-<hash>-app.pfw.gz for the main process and trace-<hash>-fork.pfw.gz for each worker.

Make cold mean cold. The dataset is 2.3 GiB and the machine has 8 GB of memory, so without intervention the second epoch reads from the page cache instead of the disk. Before every epoch, I evict the files with posix_fadvise(DONTNEED), then check with mincore how much of them is still in memory, and stop the run if more than 5% is. Every check read 0.0%. The eviction runs in a subprocess with the DFTracer variables removed, so it leaves no events in the trace. It still takes 0.26 s inside each epoch span.

What the trace looks like

These lines come from the four-worker run, two from the main process and two from a worker:

{"id":17,"name":"fetch.iter","cat":"dataloader","type":6,"pid":839,"tid":839,"ts":1792392200632907,"dur":236763,"ph":1,"args":{"hhash":"67bc2312b5a18356","epoch":0,"step":0,"count":1}}
{"id":16,"name":"backward","cat":"compute","type":6,"pid":839,"tid":839,"ts":1792392200930139,"dur":183599,"ph":1,"args":{"hhash":"67bc2312b5a18356","epoch":0,"step":0}}
{"id":221,"name":"item","cat":"data","type":6,"pid":853,"tid":853,"ts":1792392200632323,"dur":95402,"ph":1,"args":{"hhash":"67bc2312b5a18356"}}
{"id":4,"name":"read","cat":"POSIX","type":3,"pid":853,"tid":853,"ts":1792392200632516,"dur":1738,"ph":1,"args":{"hhash":"67bc2312b5a18356","ret":4096,"count":4096,"fhash":"0e7e86450ddea747"}}

Each line is one event:

  • ts and dur are in microseconds, and cat is the layer.
  • args carries the step and epoch from ai.update in the main process, plus, for POSIX calls, the bytes transferred (ret) and a file hash. Metadata records elsewhere in the trace map each hash to a path.
  • The first fetch.iter took 237 ms against an average of 3 ms, because no batch was ready yet when training started.

The volume is modest at this scale. With four workers, the main process wrote 13 KB, and each worker wrote 115 KB: 8,845 events, including 5,440 reads and 3,200 seeks across its 40 samples.

The layer breakdown

To turn those events into layers, point DFAnalyzer at a run’s trace directory:

dfanalyzer analyzer/preset=ai trace_path=traces/w0 cluster=none

The ai preset maps the annotation names to layers. With no workers, the first table it prints reads like this (trimmed to the main columns):

Layer                Time (s)       Ops   Size (MiB)
Training               37.599         1
Epoch                  37.082         2
Fetch Data             24.739       160
Data Loader            23.215       320
Reader                  0.671       320
POSIX Reads             3.658    69,761       4,680
Device                  0.025       160
Compute                12.227       320
Checkpoint              0.018         2

Training takes 37.6 s, and Fetch Data (24.7 s) plus Compute (12.2 s) account for nearly all of it.

The data layers nest. Fetch Data is the loop waiting for a batch. Inside it, Data Loader is the time spent loading samples (item). Inside that, Reader covers the augmentation and POSIX Reads the reads from the files. Read top-down, the table says the loop waited 24.7 s, and 23.2 s of that went to loading samples. Reads took 3.7 s of the loading and augmentation 0.7 s. That leaves 18.9 s, about 59 ms per sample, of work inside np.load that happens between the reads.

The trace can’t see inside np.load, but a profiler can. Per sample, about 50 ms goes to zlib inflating the arrays, 5 ms to checking their CRC32 checksums, and 4 ms to building the NumPy arrays. The image barely compresses (15.6 MiB to 14.4 MiB), yet inflating it takes more than four times as long as reading it. Loading the same files from a warm page cache, with no storage involved at all, still takes 63–68 ms per sample.

So a system-call trace alone would report 3.7 s of I/O, and the fetch layer alone 24.7 s. Both numbers are correct, but only the layers between them say what the 24.7 s is made of.

Unoverlapped time

What matters for I/O in training isn’t how long loading takes. It’s how much of that time the GPU waits for. A batch loaded while the GPU is still computing the previous step costs nothing, and only the unoverlapped part slows training down. Data-loader workers exist to create that overlap: they load the next batches in separate processes while the main process trains.

In the trace, the unoverlapped part is Fetch Data. Each fetch.iter span is the training loop blocked on the loader, so together they add up to the time the loop waited for data. That covers the loading that wasn’t hidden, plus handing each batch over: collating it, moving it out of the worker, and pinning it for the copy to the GPU. The loading work itself is the sum of the item spans across all processes. Averaged over three runs:

Workers Loading work Waited for data Compute Job time Compute share
0 23.0 s 24.5 s 12.2 s 37.5 s 33%
1 22.8 s 14.4 s 12.3 s 27.7 s 44%
2 23.9 s 1.4 s 12.7 s 15.2 s 84%
4 24.8 s 0.5 s 12.6 s 14.3 s 88%

The loading work barely changes: 23–25 s in every setting, slightly more with more workers, which compete for the CPU. What changes is how much of it the loop waits for. With no workers, the loop waits for all of it, and a little more, because the main process also collates and pins each batch. With four workers, it waits 0.5 s. Compute stays near 12.5 s throughout, so the job time follows the wait, and compute share, the fraction of the job spent computing, climbs from a third to 88%.

Runs of the same setting differed by at most 0.6 s. In every run, Fetch Data agreed within 0.15 s with a timer I kept in the script independently of the tracer.

Here is half a second of training at each setting. The waits and compute come from the measured per-step means, and loading is drawn at about 150 ms per batch per worker:

Half a second of training at each setting (ms)Wait for dataComputeWait for dataComputeWaitComputeWaitComputeWaitComputeLoadLoadLoadWaitComputeWaitComputeWaitComputeWaitComputeWaitComputeLoadLoadLoadLoadLoadLoadWaitComputeComputeComputeComputeComputeComputeLoadLoadLoadLoadLoadLoadLoadLoad0 workers1 worker2 workers4 workers0 ms50 ms100 ms150 ms200 ms250 ms300 ms350 ms400 ms450 ms500 ms

The arithmetic explains why most of the gain arrives by two workers. A sample takes 73–76 ms to load, so one worker produces a batch of two in about 150 ms, while the GPU finishes a step in 77–80 ms:

  • One worker can’t keep up, and each step waits 90 ms.
  • Two workers just about can, and each step waits 7–14 ms. Being right at the edge, two workers are also the least stable setting: one of the three runs waited twice as long as the other two.
  • Four workers leave headroom. Each step waits 4 ms, and each worker sits idle about half the time.

Where each job’s time went:

0 workers (37.5 s)33%65%2%Compute [12.21]Waiting for data [24.45]Other [0.87]
1 worker (27.7 s)44%52%4%Compute [12.29]Waiting for data [14.37]Other [1.06]
2 workers (15.2 s)84%9%7%Compute [12.69]Waiting for data [1.39]Other [1.07]
4 workers (14.3 s)88%8%4%Compute [12.61]Waiting for data [0.52]Other [1.14]

Waiting for data falls from 65% of the job to 4%, and compute rises from 33% to 88%. “Other” is about a second in every setting: checkpoints, my cache eviction, and, with workers, starting a new set of them at each epoch.

Drill down over time

DFAnalyzer computes its tables for one or more views, which are ways of grouping the same events. proc_name groups them by process, file_name by file, and time_range by fixed windows of time_granularity seconds. Overlap is a question about when things happen, so the time view is the one to read:

from dftracer.analyzer import init_with_hydra

dfa = init_with_hydra(hydra_overrides=[
    "analyzer=dftracer", "analyzer/preset=ai", "analyzer.time_granularity=1",
    "cluster=none", "trace_path=traces/w4",
])
result = dfa.analyze_trace(view_types=["time_range"])
windows = result.flat_views[("time_range",)]

Each row is one window, and the columns are named {layer}_{metric}_{stat}. For time, two statistics matter. _time_sum adds up the time of every process in the window, and _time_proc_max takes the busiest process. They differ when several processes work in parallel, and they answer different questions. Here are three seconds from the four-worker run, with the Data Loader time both ways:

Second Compute Fetch Data Data Loader, all workers Data Loader, busiest worker
1 1.03 s 0.004 s 1.88 s 0.51 s
2 0.95 s 0.004 s 1.76 s 0.45 s
3 1.03 s 0.004 s 1.90 s 0.52 s

In each of those seconds, the four workers together spent almost 2 s loading samples, more than the second itself. Yet none of them was busy for more than about half of it, and the loop waited 4 ms. The sum measures how much work was done. The busiest process measures what a synchronous step could have waited on, so for overlap, read the _time_proc_max columns.

The console’s layer table works the same way: a layer’s time there is the busiest process’s time, added up over the windows. That is why, with four workers, it shows 0.9 s of POSIX Reads, while the eight workers spent 3.3 s in reads altogether.

With one worker, every second looks alike: about 0.45 s of compute, 0.52 s of waiting, and the single worker loading for 0.85 s. The worker is busy most of the time and still can’t keep up. With four, the windows also show the epoch boundary. In second 6, compute drops to 0.37 s and the loop waits 0.21 s while a fresh set of workers produces its first batch.

Down the layers

Back to the layer breakdown, now at every setting. Each number is the busiest process’s time per window, added up, and averaged over three runs:

Workers Fetch Data Data Loader Reader POSIX Reads
0 24.5 s 23.0 s 0.7 s 3.6 s
1 14.4 s 22.8 s 0.6 s 3.4 s
2 1.4 s 12.1 s 0.4 s 1.7 s
4 0.5 s 7.5 s 0.2 s 0.9 s

Without workers, every layer sits on the loop’s own path, so each one fits inside the one above it. With workers, Data Loader and everything under it run in other processes, and the nesting breaks at the top. With one worker, the worker loads for 22.8 s, but the loop waits only 14.4 s, so at least 8.4 s of the loading ran while the GPU computed. From two workers on, each worker carries a share of the loading: 12.1 s for the busiest with two workers, 7.5 s with four. The loop’s wait is a small fraction of that.

The gap between Fetch Data and Data Loader is the overlap, read straight off the table. When the two are close, as without workers, the loading is exposed. The layers below it, here mostly decompression, are then what the loop pays for. When Fetch Data falls far below Data Loader, the loading is hidden, and the layers under it no longer set the pace of training.

Gotchas

Count the processes first. The trace should have one file per process you expect. PyTorch starts new workers every epoch by default, so here that is 1, 3, 5, and 9. The workers do nearly all the reading: with four of them, the main process read none of the dataset, so a trace that misses the workers misses almost all the I/O.

A CPU span isn’t GPU time. CUDA kernels run asynchronously, so a span ends when the work has been queued, not when it has run:

  • Forward vs. backward: forward averaged 2.9 ms per step and backward 73.5 ms. Backward isn’t 25 times more work. The loss.item() inside it waits for every kernel queued so far, so forward’s GPU time lands in backward’s span.
  • The transfer: the copy to the GPU averaged 0.15 ms because it was only queued.
  • Without the loss read: I removed it from the run without workers. Compute dropped to 1.23 s and the job from 37.5 s to 26.7 s, because the GPU now worked through the next fetch, inside spans labeled Fetch Data.

The model, the data, and the trace format were the same in both runs, but the two traces tell different stories about overlap. Know where your loop synchronizes.

What I take from it

  1. Measure what the loop waits for, not how long loading takes. Only the unoverlapped part costs training time.
  2. Trace every process, and check the process count before reading anything else.
  3. Annotate the step, the fetch, the sample, and the transform, and let DFTracer capture the POSIX calls underneath. The layers between them are where the explanations are.
  4. Find where your loop synchronizes with the GPU before you read compute and fetch times.
  5. Read top-down: the layer breakdown first, then the time windows, then each layer below the wait.

References

  1. N. Heller et al., “The KiTS19 Challenge Data: 300 Kidney Tumor Cases with Clinical Context, CT Semantic Segmentations, and Surgical Outcomes,” 2019. arXiv:1904.00445 ↩

  2. P. Mattson et al., “MLPerf Training Benchmark,” MLSys 2020. arXiv:1910.01500 ↩