Thank you for your reply.
If what you say it’s true I would expect the computation time to be linear also given the batch size, however there’s no difference in execution time between 32 or 1024 (as well in the hidden_dim of each of the 10 layers I’m using: from 256, to 1024). I imagine that there’s a lot of overhead in calling stuff, compared to executing it. In the end what I’m trying to do is computing a single linear layer which for abstraction is split in 10 that should run in parallel. I’m sure that if I build a model with 10 small layers or one with 1 big layer the performance would be very different. I’ve tried to look into things such as vmap (but doesn’t work because it doesn’t support indexing different layers) or torchscript.fork (which helps a little, 15-25%, but introduces a lot of overhead for the forking since I have very simple and small layers) but without any luck.
This is how my forward pass and backward pass for each layer look like:
https://drive.google.com/file/d/1QKLGwHfke_AdxQ-RQMScqmwZ-vCXMqBu/view?usp=sharing 1
(the second half in the forward pass is the loss (energy) computation, which happens in every layer)
I have backward pass after I collected the sum of all the individual losses, but in theory it could be called independently on each single one (I tried but it is still executed sequentially, not even in parallel with the forward pass of the next layer).
I don’t know if it can help in giving me tips to improve performance.
Thank you so much for your time.