I need to parallelize a model along the samples of a batch on CPU. The reason for that is that my model is a recursive network so I don’t have tensors which I can stack into a batch but tree structures I have to unfold during training.
Basically, the code I want to achieve looks like this.
for b in batches:
  losses = 0  # needs to be synchronized among processes
  for sample in batch:  # sample is a recursive structure, not a tensor
     # do this in parallel on batch_size-many CPU cores
     loss += model(sample)
  # sum/average all losses from the subprocesses
  loss /= num_processes
  loss.backward()
  optimizer.step()

Basically, it should work like a multiprocessing.Pool. However, I cannot use this as gradients cannot be shared along processes. I know that DistributedDataParallel exists, but I don’t see how I can use this to achieve my goal.
Thanks!