A small diagnostic for Hugging Face Trainer data-loading bottlenecks

I kept running into a frustrating training-performance question: GPU utilization would hover around 50%, loss looked normal, and nothing was obviously broken. Was the model actually the limit, or was the GPU waiting for the next batch?

Those are opposite problems with opposite fixes. But nvidia-smi and a healthy loss curve do not distinguish them, and I do not want to open a full PyTorch Profiler trace for every ordinary run.

So I added a Hugging Face Trainer integration to TraceML, an open-source PyTorch diagnostics project I have been building.

from traceml.integrations import huggingface as traceml_hf
from traceml.integrations.huggingface import TraceMLTrainerCallback

traceml_hf.init()

trainer = Trainer(
    ...,
    callbacks=[TraceMLTrainerCallback()],
)

It keeps the rest of the Trainer setup unchanged, separates input wait from step work, and gives a practical input-bound verdict.

I made a runnable Colab with ResNet-50 and real images to test this properly. It runs twice, changing only DataLoader settings. On my included run, it was 1.83× faster and the diagnosis changed from input-bound to compute-bound.

That number will vary with the CPU/GPU ratio. The useful part is finding out which side of the line your own run is on before optimizing the wrong thing.

I would really value feedback from people running real Trainer workloads, especially cases where this diagnosis is surprising or wrong.

For now, I tried a small experiment in Colab:


A small surprising case

The lightweight “which side should I investigate first?” role still looks useful to me. I did, however, find one Hugging Face Trainer case where I would be cautious about interpreting H2D: 0.0 ms as “no host-to-device transfer happened.”

In the tested path, Accelerate moved each batch from CPU to CUDA before TraceML’s on_step_begin callback opened the trace_step() window.

With gradient accumulation set to 2, the observed order was consistently:

microbatch 1: CPU → CUDA
microbatch 2: CPU → CUDA
TraceML on_step_begin
training_step for microbatch 1
training_step for microbatch 2
optimizer step
TraceML on_step_end

The transfers were nonzero when measured independently with CUDA events, but the TraceML summary reported H2D: 0.0 ms.

I reproduced that with the main revision checked out by the final probe, commit f3a09d2, across these four single-process paths:

Accelerate path Transfer mode CPU→CUDA transfers before each TraceML step window Independently measured H2D per optimizer-step group TraceML H2D
DataLoaderShard blocking 2 about 1.74 ms 0.0 ms
DataLoaderShard non-blocking 2 about 1.51 ms 0.0 ms
DataLoaderDispatcher blocking 2 about 4.15 ms 0.0 ms
DataLoaderDispatcher non-blocking 2 about 3.78 ms 0.0 ms

I would not use the differences between those four timings as a loader-performance ranking; the useful result here is the event ordering and the existence of nonzero transfers outside the callback window.

I also ran a positive control in which a similarly sized CPU→CUDA transfer was deliberately placed inside trace_step(). In that case:

  • independent CUDA-event timing: about 1.50 ms
  • TraceML H2D timing: about 1.46 ms

So this looks less like a generally broken H2D timer and more like an integration coverage / timing-window boundary.

My practical reading would therefore be:

In this Trainer + Accelerate path, H2D: 0.0 ms did not establish that no H2D transfer occurred. It established that no H2D transfer was observed inside the current TraceML optimizer-step bracket.

The distinction that seems most useful

I think three separate properties may be needed:

Property Example states
Metric value zero / nonzero / null
Instrumentation availability available / unavailable / not applicable
Integration coverage inside the current window / outside the current window / partial / unknown

That third property matters here.

The H2D instrumentation was available and worked in the positive control, but the normal Accelerate transfer happened before the integration armed it. Treating this only as “missing versus measured zero” may therefore still lose an important distinction.

The existing issues about preserving missing Step Time signals and gating diagnosis and representing missing metrics in final summaries already seem closely related and point in a useful direction:

  • retain genuine measured zero as 0.0;
  • represent unavailable data separately;
  • abstain from diagnoses that require signals which were not observed;
  • expose per-metric availability in the final evidence.

This fixture seems to add one neighboring case:

the instrumentation exists
+ the operation exists
+ the operation occurs outside the integration's current bracket

For that case, a label such as outside_window, not_covered, or pre_step may be more informative than either 0.0 or a generic n/a.

My default design preference would be to preserve the existing meaning of one TraceML step—one optimizer step—and expose pre-step input delivery separately, rather than silently expanding the denominator. But there seem to be several reasonable routes:

  1. Add a pre-step input-delivery phase
    Measure DataLoader delivery and Accelerate device placement before on_step_begin, while keeping the current optimizer-step bracket intact.

  2. Include device placement in a broader input-delivery metric
    This may be simpler for users who primarily want to know how much time elapses before model work can begin, though it would need clear naming because DataLoader fetch and H2D are different operations.

  3. Keep the current timing window, but qualify the metric
    Report H2D as outside the current integration window rather than as a measured zero.

Even the third option would make the summary considerably safer to interpret.

Experiment setup and event ordering

Environment

The final matrix used:

Component Value
GPU Tesla T4
PyTorch 2.11.0+cu128
Transformers 5.14.1
Accelerate 1.14.0
TraceML commit f3a09d23996a6d6ddd68e8b0dcdbd8f9b5559b7e
Processes 1
Gradient accumulation 2
DataLoader workers 0
Input pinned CPU tensors
Optimizer steps per lane 6

The DataLoaderDispatcher lanes exercised that code path in one process. They were not a test of multi-rank broadcast behavior.

Observed counts

Across the four lanes:

  • 24 TraceML optimizer-step windows were opened;
  • every window had exactly two CPU→CUDA transfers before it;
  • all 48 transfers occurred while the TraceML callback bracket was closed;
  • all 48 transfers occurred while the H2D timing state was inactive;
  • all 48 calls to training_step received tensors already resident on CUDA;
  • TraceML reported H2D: 0.0 ms in every lane.

An earlier smaller probe also compared gradient accumulation 1 and 2:

GA=1:
    1 transfer
    → on_step_begin
    → 1 training_step

GA=2:
    2 transfers
    → on_step_begin
    → 2 training_step calls

That smaller fixture likewise produced nonzero independent H2D timings but 0.0 ms from TraceML.

The sequence is therefore consistent with the optimizer-step semantics described in TraceML’s Hugging Face integration code: the callback opens trace_step() in on_step_begin, and accumulated microbatches are folded into one TraceML step.

Why this ordering appears in this Trainer + Accelerate path

The event ordering also matches the pinned upstream implementations.

1. TraceML opens its step in on_step_begin

The TraceML callback describes itself as a bracket around trace_step():

  • enter in on_step_begin;
  • exit in on_step_end;
  • one TraceML step equals one optimizer step;
  • accumulated microbatches are folded into that step.

See the callback implementation and the actual context entry in on_step_begin.

2. Trainer collects the accumulation group first

In Transformers 5.14.1, the outer loop represents an optimizer update. It first calls get_batch_samples(...) to collect the requested number of microbatches.

Only after that collection does the inner loop call on_step_begin before processing the first microbatch.

The relevant section is in the pinned Trainer loop.

Conceptually:

batch_samples = collect gradient_accumulation_steps batches

for inputs in batch_samples:
    if first microbatch:
        on_step_begin()
    training_step(inputs)

3. Accelerate transfers each batch before yielding it

For DataLoaderShard, Accelerate calls send_to_device(...) before yielding current_batch to the consumer:

DataLoaderDispatcher similarly moves the batch to the device before its later processing and delivery:

Putting those three boundaries together gives:

Accelerate fetches and transfers microbatch 1
Accelerate fetches and transfers microbatch 2
Trainer finishes collecting the accumulation group
Trainer calls on_step_begin
TraceML opens trace_step
Trainer runs forward/backward for both microbatches
Trainer performs the optimizer step
TraceML closes trace_step

This explains why an H2D timer that is armed by trace_step() can work correctly for an explicitly in-window transfer while missing the ordinary Accelerate transfer.

What I think this establishes—and what it does not

What the fixture establishes

For the listed versions and single-GPU setup:

  • real CPU→CUDA transfers occurred;
  • those transfers occurred before the TraceML callback window;
  • independent CUDA-event timing measured them as nonzero;
  • TraceML’s final H2D value remained 0.0 ms;
  • moving a transfer into the active TraceML step made it measurable;
  • gradient accumulation increased the number of pre-window transfers grouped under one optimizer step;
  • the result reproduced in both DataLoaderShard and DataLoaderDispatcher code paths, with blocking and non-blocking transfer modes.

What it does not establish

It does not show that:

  • TraceML’s overall diagnosis is generally wrong;
  • the H2D timer is broken in all integrations;
  • all Hugging Face Trainer configurations have this ordering;
  • every 0.0 ms H2D value is an out-of-window value;
  • the full residual time is hidden H2D;
  • the four independent H2D timings are directly comparable performance measurements;
  • multi-process dispatcher, DDP, FSDP, DeepSpeed, XLA, or other device-placement paths behave identically;
  • asynchronous transfer is or is not successfully overlapped with later work.

In the earlier GA=1 and GA=2 fixtures, the independently measured H2D was only a fraction of the reported residual. So I would not infer:

RESIDUAL_HEAVY is wrong

from this result alone.

A narrower statement seems supported:

one component that actually occurred was attributed neither to H2D
nor to an explicitly named pre-step transfer phase

That is mainly a reporting and coverage issue, unless the missing component is large enough to change a threshold-based verdict.

A small regression fixture that may be useful

If this fits the intended semantics, the same setup could be reduced to a fairly small integration regression test.

Lane A: ordinary Trainer delivery, GA=1

Expected event relationship:

one device transfer before on_step_begin

The test should not necessarily demand that TraceML report the transfer as H2D if the integration intentionally defines H2D as in-window only. It should instead assert that the final output does not represent an unobserved transfer as an unqualified measured zero.

Lane B: ordinary Trainer delivery, GA=2

Expected event relationship:

two device transfers before one on_step_begin

This protects the optimizer-step versus microbatch boundary.

Lane C: explicit in-window positive control

Expected:

trace_step begins
CPU→CUDA transfer occurs
TraceML H2D is nonzero and close to an independent timer

This distinguishes coverage failure from timer failure.

Suggested assertions

The most robust assertions seem structural rather than exact-duration based:

  • how many transfers occurred;
  • whether the callback bracket was open;
  • whether the H2D timer was armed;
  • whether inputs were CPU or CUDA at each boundary;
  • whether the public value was measured, unavailable, or outside-window;
  • whether a diagnosis that requires H2D knows which of those states applies.

Exact millisecond thresholds would probably make the test unnecessarily hardware-sensitive.

A separate note on the 1.83× example

The input-bound → compute-bound transition in the notebook is still a useful demonstration of a combined intervention.

I would only separate three different timing claims when presenting it:

Comparison Approximate ratio
Whole TraceML process duration 140.0 / 76.4 ≈ 1.83×
Trainer train_runtime 106 / 68.66 ≈ 1.54×
TraceML total-step statistic 522.2 / 337.5 ≈ 1.55×

So 1.83× is valid as the whole-process ratio, but the training-loop improvement in that recorded run looks closer to 1.54–1.55×.

Showing both would help future readers distinguish:

process startup / download / initialization / teardown

from:

steady training-loop throughput

The notebook also changes num_workers, pinned-memory behavior, and persistent workers together. That is perfectly reasonable for a practical “before/after” recipe, but it demonstrates the effect of the combined configuration, not the individual contribution of each switch.

Because the run ends before completing a full epoch, I would also avoid attributing much of that particular improvement specifically to the epoch-to-epoch worker reuse provided by persistent workers.

None of this weakens the central lesson that the input-side change materially improved the run. It just makes the scope of the speedup claim and the intervention easier to reuse correctly.

The runnable example is in the Hugging Face data-loading bottleneck notebook.

A compact validation path for future surprising results

For future reports where a diagnosis seems surprising, this sequence may provide a lot of information without requiring a full profiling campaign.

1. Check signal status first

Before interpreting 0.0, distinguish:

measured zero
unavailable
not applicable
outside the current timing window

2. Separate denominators

Record separately:

  • whole-process duration;
  • Trainer runtime;
  • optimizer-step timing;
  • microbatch count;
  • gradient accumulation;
  • evaluation and checkpoint phases.

3. Apply one differential control

Examples:

  • real input versus synthetic in-memory tensors;
  • storage-backed input versus cached input;
  • real collator versus a lightweight collator;
  • ordinary Accelerate delivery versus an explicit in-window transfer.

If the diagnosis is useful, the relevant component should usually move in the predicted direction.

4. Use a short profiler trace only when needed

A short steady-state torch.profiler window can then answer questions such as:

  • was the GPU idle because the host had not submitted work;
  • where did the transfer occur relative to model execution;
  • was transfer overlapped;
  • did the broad input-versus-compute classification agree.

The goal would not be exact millisecond equality between tools with different timing windows. The useful comparison is whether they point to the same broad part of the execution path.

5. Change one tuning variable at a time when attribution matters

A combined “good default” configuration is useful operationally. A one-factor ablation is only necessary when the goal is to determine which setting produced the improvement.

Overall, I still think the lightweight first-pass diagnostic is a useful direction. This fixture mainly suggests that the public output should avoid making 0.0 ms carry more meaning than the integration window actually establishes.

The lowest-cost safety improvement would be to distinguish a measured zero from unavailable or out-of-window coverage. A separate pre-step input-delivery phase would provide an even more complete picture while preserving the current one-optimizer-step TraceML bracket.

If useful for the project, the GA=1 / GA=2 plus in-window-positive-control fixture looks small enough to serve as an integration regression case.

Thanks, this is an excellent and very clear reproduction. I opened #276 to track it.

Your conclusion is right: the H2D timer works when a transfer is inside the TraceML step window, but in this Trainer + Accelerate path device placement can happen before the callback opens that window. So H2D: 0.0 ms must not be read as “no transfer happened.”

We will make the metric semantics safer: CPU-only runs as N/A, and uncaptured CUDA H2D as no captured samples rather than a default zero, and add the GA regression fixture.

Notebook update: I refreshed the linked Colab to use traceml_ai, TraceMLTrainerCallback, the smaller 320px Imagenette archive, and traceml compare.

This does not resolve the H2D coverage issue discussed above. #276 tracks that distinction. The notebook’s before/after result should be read from its measured input-wait and total-step changes, not as a claim that H2D is absent.

Oh. When I ran the updated notebook, I was able to reproduce the main result pretty well:


I ran the refreshed huggingface_dataloading_bottleneck.ipynb from top to bottom on another Colab Free T4 session. It completed successfully, and the main before/after pattern reproduced clearly.

Metric Baseline Optimized Change
Trainer runtime 82.16 s 68.53 s about 16.6% lower
Total step 404.1 ms 337.9 ms about 16.4% lower
Input wait 104.8 ms 2.2 ms about 97.9% lower
Average GPU utilization 48.6% 84.7% about +36 percentage points
TraceML verdict INPUT-BOUND / CRITICAL COMPUTE-BOUND / INFO flipped as intended

On this Colab runtime, the optimized lane selected two workers and enabled pinned memory and persistent workers. The input-wait share fell from 25.8% of the typical iteration to 0.6%, while the compute share rose from 72.4% to 97.0%.

So the exact improvement was different from the output saved in the repository—which seems entirely reasonable for a hardware-dependent example—but the important direction reproduced:

large exposed input wait
→ loader configuration change
→ input wait almost disappears
→ GPU utilization rises
→ the diagnosis moves from input-bound to compute-bound

I used the Trainer runtime and the measured total-step/input-wait values for that comparison, rather than the full process durations. The full TraceML run durations were 113.6 seconds and 75.6 seconds, but the baseline process also downloaded the model while the second run benefited from the cache, so that ratio would mix setup effects with training-loop improvement.

Also, as noted in the update above, both lanes still reported H2D: 0.0 ms. I did not interpret that as evidence that no transfer occurred; I treated the result through the input-wait and total-step measurements, in line with the current H2D coverage note.

One other useful aspect of the refreshed notebook is that its conclusion now matches the scope of the experiment quite well: the three loader settings are changed together to demonstrate a practical before/after pattern, while the final guidance recommends testing them one at a time on the hardware that will actually run the job.

So, at least on this independent Colab Free T4 run, the updated notebook was reproducible in the sense that seems most important for the demo: it correctly exposed a substantial input-side stall, the proposed loader configuration nearly removed that exposed wait, and the measured end-to-end training-step cost improved accordingly.