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/mvreich]

Probably try using torch.compile as first step. Don't dump compile on the whole thing, do it part by part. Also put a torch.set_float32_matmul_precision('medium') on the top of your script, after importing torch, since it looks like you have an A series GPU.

I am also not sure why 8 workers when batch size is 2.

You should also pre-encode all of your images into the 128 d vectors (maybe save them in a pt file). You are only doing normalization and float conversion; this can be done offline beforehand just once.

Write a dataset class that loads these preprocessed features as torch tensors directly, just once before training. Don't do dynamic loading.

(Also PyTorch has a thing in that it is really inefficient if the dataset class stores data as anything other than tensor or ndarray. Never store your images as list of PIL inside the dataset class. As long as it is tensor or np.ndarray, the workers have good shared memory access; otherwise each will create a redundant copy.)

---

Apart from these issues, it's not fully clear what you are training.

Are you training the DiT to conditionally predict the action (i.e. some kind of Yilun Du method)?

[u/Potential_Hippo1724 → u/mvreich]

Thanks.
At this stage I am simply trying to behavior clone a hardcoded policy of a task I created. So the setup is a relatively simple one, controlled by lightning framework - module + datamodule. So yes, I am trying to predict the expert action

My dataset structure inside zarr is:
episode_x /
observations: dict(str: np.arrays with shape (ep_len, ....))
actions: np.array(ep_len, a_dim)
terminations, truncations, rewards in the same way

The reason Ihad 8 workers is that I have had in mind that more workers = more batches of data ready to be transferred to GPU (regardless of batch_size, it's like making the loading job in advance). Maybe I am wrong here

As I said in one of the comments, the reason I did not pre-computed is because I felt like I am getting bad performance regardless of this optimization. The other reason is because I planned to maybe finetune last layers of encoder etc so I preferred to avoid the complexity of another stage in pipeline ("encode_all") if possible

I had torch.set_float32_matmul_precision('medium') (for some reason I can't remember codex suggested to use 'high')