I tried the ImageNet example with ResNet152 on 8GPUs but it is much slower than fb.resnet.torch (1.5s vs 0.8s per iter).
The replicate 7 in DataParallel could be the bottleneck that costs half of the forward time. The Broadcast 2 function is called for every parameter and buffer, while in fb.resnet.torch, the parameters are flattened first and the bcast is only called once.
It’s elegant to implement the Broadcast as an Op/Function. I wonder if it is possible to overlap the communication with computation during forward/backward? Or it is necessary to flatten the parameters in order to improve the efficiency?