François CholletVerified account

Overall I think our parallelism model is sound (per-device forward and backwards pass on sub-batches + CPU concat, keep params on CPU)

But it would need auditing, in particular to check whether we correctly place the grads computation on each device (now we leave it to TF)

His NVIDIA DevBox is PCIe 16x, Azure is PCIe 8x, custom builds often even worse. Scaling also impacted by small batch sizes (64 -> 32).

Overally we focus on small batch sizes, as I see huge sample sizes in many real world applications, including internally at @RossumAi.

We have significant overhead (shuffling the params from CPU to GPU at every step; split+concat). Need time(process(sub_batch)) >> overhead

Not doable to reduce the overhead, so to have a speedup you need to keep your per-sub-batch processing time high (large models or batches)

This is just async prefetch to 1 GPU. Then it needs StagingArea placed at each GPU. Also we need async loading to TF (Dataset API/queue).