Parallel layers/modules execution

[liukidar]

Hello!

I'm developing a library to build and execute predictive coding networks. One of the main features of PCNs is that they use only local computations (i.e., each layer's forward and backward pass is independent of the others). As a consequence, it is theoretically possible to execute them in parallel (layers could be of any kind, convolutional, linear, etc. etc.). However, even when jitting, it seems that they are still executed sequentially. It would be great to overcome this issue, as the network would train L times faster (where L is the number of layers).

In pseudo-code, what I'm trying to achieve is more or less the same:

def forward():
    e1= self.layer1.forward()
    e2= self.layer2.forward()
    ....
    eL= self.layerL.forward()
    loss = e1 + e2 + ... +eL
   
    return loss

forward_and_backward = grad(forward())

updates = forward_and_backward()
...

I've been trying with very small batch sizes and hidden dimensions (on a RTX TITAN), and the total time of a 'forward_and_backward' step (averaged over an epoch of training) scales linearly with the number of layers, while the GPU utilization can be as low as 10%.

Any help would be much appreciated.