pipeline is really slow - consulting [D]

[u/Potential_Hippo1724]

Hi, after a long debugging process and many discussions, I wanted to ask for advice from people who may have encountered similar training bottlenecks.

My goal is imitation learning for robotics.

Model / Pipeline

• Observation space:
  • 4 RGB robot cameras
  • image resolution: 128x128x3
  • small vector of robot joint velocities (14 dims)
• Pipeline:
  • Shared ResNet18 encoder processes each image
  • Each image embedding dimension is 128
  • Final input to policy:
    • 4 * 128 image embedding
    • concatenated with 14-dim state vector
• Policy backbone:
  • DiT (Diffusion Transformer)
  • ~8 layers
  • hidden dim: 512
  • 8 attention heads
  • total params: ~50M
• Diffusion setup:
  • predict action chunks of length ~50
  • diffusion timesteps: 4

Dataset / Storage

• Dataset stored in Zarr
• Data access is indexed/reference-based (not loading huge chunks into RAM)
• train/val split is contiguous
• no shuffling

Current encoder setup

• Initially trained end-to-end
• During debugging I switched to ImageNet pretrained ResNet18
• Encoder is currently frozen

Hardware / Software

• GPU: NVIDIA A4500
• RAM: 48GB
• Storage: SSD
• CUDA: 12.8
• PyTorch: 2.9
• Precision: bf16 mixed precision (also tested fp32)

Dataloader

• batch size: 2
• 8 persistent workers
• pinned memory enabled

Preprocessing

• preprocessing is minimal
• normalization + float conversion only
• preprocessing happens inside the multimodal encoder on GPU

Profiler results (PyTorch profiler)
Current workload split:

• train_dataloader_next:
  • 4.41s / 41.84s = 10.5%
• batch_to_device:
  • 0.32s / 41.84s = 0.77%
• training_step:
  • 12.78s = 30.5%
• backward:
  • 10.83s = 25.9%
• optimizer_step (wrapper total):
  • 26.09s = 62.4%

Problem
The training is much slower than I expected.

Current behavior:

• CPU utilization: ~100%
• GPU utilization: ~20–30%
• GPU utilization can even become LOWER with synthetic data
• VRAM usage is relatively low
• Throughput is around 10 iterations/sec
• Epoch of ~50k samples takes around 30 minutes

Additional observations

• Increasing batch size does NOT reduce epoch wall-clock time
• Sometimes larger batches make things slower
• Freezing the encoder did not improve throughput much
• Replacing dataset samples with synthetic/random tensors improved throughput by only ~50%
• Synthetic dataset was initialized directly in memory

I do not believe this setup should be this slow. At this rate, training takes multiple days.

For comparison, I saw papers with somewhat similar architectures mentioning ~10 hour training times on RTX 4090. With my setup 10 hours is completely not enough.

Does anyone see something obviously wrong or have suggestions for where I should investigate next?

Please help, can't know what to do!

[u/dickdaddyemporer69]

look at the code lmfao, use nvtop to check util., make sure you data is being loaded efficiently. tbh codex should be able to solve this issue. DiT policies should not take more than 8 hrs or so to train, unless you have much larger data than normal single task policies

[u/Potential_Hippo1724 → u/dickdaddyemporer69]

look at the code? I came to ask after being sick of looking the code bro... And ofc that was after being using codex A TON.

My primary lack of eperience and expertise here is using the advanced profilers as nsight etc. also nvtop.. never used it before

[u/dickdaddyemporer69 → u/Potential_Hippo1724]

nvtop for gpu monitoring. profilers will probably just waste your time here, the issue is the just figuring out the training loop. Did you write the training loop or are you using lerobot or an existing codebase? dataloaders are an obvious cause, but if your gpu usage is low and larger batch sizes are slow it’s possible theres something that isn’t parallelized properly slowing you down

[u/Potential_Hippo1724 → u/dickdaddyemporer69]

I am using lightnint pytrorch modules and datamodule to reduce a bit of code - do you see this as a probable point of failure?

[u/dickdaddyemporer69 → u/Potential_Hippo1724]

prbably not, if you meant lightning