Hello @mrshenli. Thank you for your reply.
Model parallelism will probably also not work because even though the torch calls are async, I have to call backward() at the end of the for loop, creating a blocking bottleneck. The computation graph would otherwise get too large for a single GPU.
Your reply made me realize I had silently forgotten about DataParallel. Since my batches contain whole graphs, I somehow forgot I could further divide the graph into batches of size 2. I had stopped thinking about DataParallel when I found out about the advantages of DistributedDataParallel.