CPU-side ColorJitter is a dataloader bottleneck — could the GPU path be restored?

[xiaojunjie]

DROID action SFT currently runs ColorJitter inside the dataloader workers, which is
computationally expensive.

1. In principle: ColorJitter is a random augmentation to begin with, so moving it to
   the GPU should only introduce CPU/CUDA kernel-level numerical differences, without
   changing the training semantics — is that right?

2. In practice: ColorJitter and the tail after it can be moved out of the dataloader
   worker and run in the main process on the GPU. The pipeline stage and the ordering stay exactly
   the same; only the execution device changes. The worker also passes along its RNG
   state, which keeps the final difference limited to CPU/CUDA kernel rounding.
   The cost is that worker-to-main-process traffic doubles, since the tensor returned
   by the worker goes from [T, C, 3H/2, W] to [3T, C, H, W].

Does this approach look viable?

[lfengad]

Yes to (1), with one caveat, and (2) works but costs more than you estimated.

(1) Correct — it's a random augmentation, so the device swap is a re-parameterization, not a semantic change. But don't ship the worker's RNG state: ColorJitter draws from the global torch RNG, so restoring it in the main process clobbers the stream the model uses (flow-matching noise, cfg_dropout_rate, dropout). Ship the sampled parameters instead — 4 factors plus the op-order permutation. Also keep them per sample; jittering the collated batch once would cut color diversity by batch_size.

(2) Viable, but there are two costs:

• The worker doesn't return the composite — _convert_video casts to uint8 and VideoResize resizes/pads after it. Your pre-jitter tensor is float32, so it's ~7×, not 2×, unless the uint8 cast also moves ahead of the jitter. And since the jitter precedes composition, compose + convert + resize/pad all have to move with it.
• More importantly, this trades a CPU bottleneck for GPU critical path + IPC bandwidth. The GPU is the scarce resource in this recipe, and moving the jitter there adds an RGB→HSV→RGB pass to every step, while IPC traffic goes up 2–7×. That may just relocate the stall from the worker into the step rather than remove it.

Cheaper first step, same reasoning: reorder instead of relocating. Run the jitter after composition — L/R are already downscaled 2×, so that's exactly half the pixels — and after the uint8 cast, for another 4× off the memory traffic. Few-line diff, no IPC change. Worth profiling hue=0.08 first though; it's the one op with no integer fast path and likely most of the cost.